from pathlib import Path
import io

main_prompt = """I'm giving you an image of several written answers to an exam.

Each answer is separated by a black horizontal line, and underneath,
to the left, is indicated the ID of the answer, from `01` to `50`.

I want you to score each answer, from 0 to 4, you may score half
points, such as 2.5. Even if a result is wrong, if the reasoning is
correct and could lead to a right answer, you should give at least
half the points.

You also need to give feedback to the student, in french :
 - which part of his answer is wrong,
 - why is it wrong
 - possibly, what he should have done instead.
Your feedback may contain LaTeX fragments written like `$a^2 + b^2 = c^2$`.

If your score is not 4, you should always provide some feedback
explaining what's missing.

For each piece of feedback, if it is related to a specific part of the
answer that is wrong, you may provide a `box_2d`, to locate this
specific part of the answer. This `box_2d` should be in the form
[ymin, xmin, ymax, xmax] normalized to 0-1000. If you do not provide
one, set `box_2d` to `null`.

If the answer is correct, there is no need to provide feedback. You do
not have to give positive feedback, but if you do, do not provide a
`box_2d` for it.

For example, if the student says a function is continuous when it
isn't, provide the coordinates where the word «continuous» is. If a
calculation went wrong, gives the coordinates of the step where it
goes wrong, and as feedback, what went wrong.

Avoid giving feedback about confusing letters `n` with `m`, `x` with
`n` or `h` with `k`. If it looks wrong, assume you read it wrong,
unless the distinction is very important.

You should also give me a measure of confidence, from 0 to 1 that you
were able to correctly understand the answer. A score below 0.5 means
that you think it is likely that you couldn't understand an important
part.

In some case, you may find that either
 - The student didn't answer the right question. Set the score to 0.
   Since it could be a labeling error, indicate is by setting `error`
   to \"wrong-label\".
 - You can find an answer to another question of the exercice (taking
   more than a couple of lines). Score the question you are supposed
   to score, but set `error` to \"additional-answer\".
 - The answer to the question is empty, or the student has only
   rewritten the statement of the question. In this case, set `error`
   to \"empty-answer\" and do not provide any kind of feedback.
If there's no error, set `error` to `\"\"`.

You will answer using json describing a list of dictionary with a key
\"id\", and a key \"result\" that contains the \"score\", the \"confidence\", a
list \"feedback\", and possibly an \"error\". Like this example :

[{ \"id\": \"01\",
   \"result\": {\"score\" : 2.5,
              \"confidence\" : 0.8,
              \"feedback\": [{text: \"Un retour générique. Il faut apprendre le cours.\", box_2d: null},
              {text: \"Non, la fonction n'est pas forcément continue\", pos: [145, 280, 340, 500]}],
               \"error\": \"\"}
 },
 { \"id\": \"04\",
   \"result\": {\"score\" : 4.,
              \"confidence\" : 0.9,
              \"feedback\" : []
              \"error\": \"\" }
 }
]

Here is the text of the exercice (or the relevant part of the problem)
of the exam :

```
<<text>>
```

Here is a possible correct answer :

```
<<corr>>
```
<<persp>>

You are asked to score the question or exercice labeled `<<label>>`,
do not score or give feedback to any other question."""

def make_prompt(input_dir,full_label):
    def read_longest_prefix_file(subdir):
        dir_path = input_dir / subdir
        matches = [f for f in dir_path.iterdir()
                   if f.is_file()
                   and full_label.startswith(f.name)
                   and f.suffix not in [".pdf", ".tex"]]
        if not matches:
            return ""
        return max(matches, key=lambda f: len(f.name)).read_text(encoding="utf-8", errors="replace")

    text = read_longest_prefix_file("Text")
    corr = read_longest_prefix_file("Sol")
    persp = read_longest_prefix_file("Persp")

    if persp != "":
        persp = "\n\nHere are additional scoring instructions : \n\n```\n" + persp +"\n```\n"
    return main_prompt.replace("<<text>>", text).replace("<<corr>>", corr).replace("<<persp>>", persp).replace("<<label>>", full_label)


from pydantic import BaseModel, Field, TypeAdapter
from typing import List, Optional, Tuple

class FeedbackItem(BaseModel):
    text: str = Field(description="Feedback content")
    box_2d: Optional[List[int]] = Field(None, description="box coordinates or null")

class ResultData(BaseModel):
    score: float = Field(description="The numeric score")
    confidence: float = Field(description="Confidence level")
    feedback: List[FeedbackItem] = Field(description="List of feedback items")
    error: str = Field(description="Indicates if an error occurred")

class EvaluationEntry(BaseModel):
    id: str = Field(description="Entry identifier")
    result: ResultData = Field(description="Result details")

# These nested definitions  do not work with the batch api, unroll them
UNROLLED_SCHEMA = {
                    "type": "ARRAY",
                    "items": {
                        "type": "OBJECT",
                        "properties": {
                            "id": {"type": "STRING", "description": "Entry identifier"},
                            "result": {
                                "type": "OBJECT",
                                "properties": {
                                    "score": {"type": "NUMBER", "description": "The numeric score"},
                                    "confidence": {"type": "NUMBER", "description": "Confidence level"},
                                    "error": {"type": "STRING", "description": "Indicates if an error occurred"},
                                    "feedback": {
                                        "type": "ARRAY",
                                        "description": "List of feedback items",
                                        "items": {
                                            "type": "OBJECT",
                                            "properties": {
                                                "text": {"type": "STRING", "description": "Feedback content"},
                                                "box_2d": {
                                                    "type": "ARRAY",
                                                    "items": {"type": "INTEGER"},
                                                    "nullable": True,
                                                    "description": "box coordinates or null"
                                                }
                                            },
                                            "required": ["text"]
                                        }
                                    }
                                },
                                "required": ["score", "confidence", "feedback", "error"]
                            }
                        },
                        "required": ["id", "result"]
                    }
                }

from google.genai import types

# The root model for parsing is be: List[EvaluationEntry]
def generate_request(input_dir, file, full_label):
    """Generates request for Gemini."""
    prompt = make_prompt(input_dir, full_label)
    image_path = Path(file)

    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_bytes(
                    data=image_path.read_bytes(),
                    mime_type="image/jpeg"
                ),
                types.Part.from_text(text=prompt),
            ],
        )
    ]

    generate_content_config = types.GenerateContentConfig(
        temperature=1.0,
        top_p=0.95,
        seed=0,
        max_output_tokens=65535,
        response_mime_type= "application/json",
        response_json_schema= TypeAdapter(List[EvaluationEntry]).json_schema()
    )
    return (contents, generate_content_config)


from pdf2image import convert_from_path
from PIL import Image
import json


def get_single_image_bytes(pdf_path):
    """Converts a multi-page PDF into a single stitched JPEG in memory."""
    imgs = convert_from_path(pdf_path, dpi=200) # Same DPI as grouping.py
    if not imgs:
        raise ValueError(f"No pages in {pdf_path}")

    if len(imgs) == 1:
        combined = imgs[0]
    else:
        max_width = max(img.width for img in imgs)
        total_height = sum(img.height for img in imgs)
        combined = Image.new('RGB', (max_width, total_height), 'white')
        y_offset = 0
        for img in imgs:
            combined.paste(img, (0, y_offset))
            y_offset += img.height

    img_byte_arr = io.BytesIO()
    combined.save(img_byte_arr, format='JPEG', quality=85)
    return img_byte_arr.getvalue()


def request_for_box_correction(pdf_path, original_feedbacks):
    img_bytes = get_single_image_bytes(pdf_path)

    localized_feedbacks = [f for f in original_feedbacks if f["box_2d"]]

    prompt = f"""
    Here is a single student's submission to a question in a written exam. The following JSON contains feedback items with bounding boxes (box_2d) that are incorrect. Each piece of feedback is supposed to be related to a piece of the answer that is wrong.

    For example, if the student says a function is continuous when it
isn't, the coordinates should be where the word «continuous» is. If a
calculation went wrong, the coordinates should be where the step where
it goes wrong, and the feedback is what went wrong.

    Please analyze the image and return the same feedback json content, but with ONLY the box_2d coordinates corrected for this specific image.
    Coordinates must be [ymin, xmin, ymax, xmax] scaled to 1000. If a box is invalid/not found, return null for it.
    Original feedback:

    {json.dumps(localized_feedbacks, indent=2)}
    """



    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_bytes(data=img_bytes, mime_type="image/jpeg"),
                types.Part.from_text(text=prompt),
            ],
        )
    ]

    config = types.GenerateContentConfig(
        temperature=1.0,
        response_mime_type="application/json",
        response_json_schema=TypeAdapter(List[FeedbackItem]).json_schema()
    )
    return contents,config

def request_for_wrong_label(pdf_path, label, enonce, labels_txt):

    prompt = f"""This image is a part of the answer of a student to a written exam.

It was initially labeled '{label}' but I suspect this label is wrong. Perhaps the student himself wrote the wrong label.

You need to analyse this image, and find the label of the question it answers. Do not trust the label written by the student but instead check the content of its answer and the notation he uses to identify the correct label of the question the student answered.

Return ONLY the exact label string.

Here is the full content of the exam :

{enonce}

Here is a list of all possible labels. You need to answer with one of these :

{labels_txt}
"""

    contents = [types.Content(role="user", parts=[
            types.Part.from_bytes(data=get_single_image_bytes(pdf_path), mime_type="image/jpeg"),
            types.Part.from_text(text=prompt)])]
    config = types.GenerateContentConfig(temperature=1.0)
    return contents, config

def request_for_additional_answer(pdf_path, label, enonce, labels_txt):
    prompt = f"""This image is a part of the answer of a student to a written exam.

It was initially labeled '{label}' but I suspect this image also contains answers to another, or several other questions.

You need to analyse this image, and find the list of the labels of the questions it answers. Return ONLY the list of the exact label strings.

If the end of the image only contains the first line of an answer to another question, ignore it.

Here is the full content of the exam :

{enonce}

Here is a list of all possible labels. You need to answer with a list one of these :

{labels_txt}
"""
    contents = [types.Content(role="user", parts=[
        types.Part.from_bytes(data=get_single_image_bytes(pdf_path), mime_type="image/jpeg"),
        types.Part.from_text(text=prompt)
    ])]
    config = types.GenerateContentConfig(temperature=1.0, response_mime_type="application/json")
    return contents, config