from pathlib import Path
import io
from . import utils

main_prompt = """Je te fournis une image contenant plusieurs réponses manuscrites à un examen.

Chaque réponse est séparée de la précédente par une ligne horizontale noire.
En dessous de cette ligne, à gauche, figure l'identifiant de la réponse,
compris entre `01` et `50`.

Attribue à chaque réponse une note de 0 à 4. Les demi-points sont autorisés,
par exemple 2.5. Même si le résultat est faux, accorde au moins la moitié
des points si le raisonnement est correct et pourrait conduire au bon résultat.

Rédige tous les commentaires destinés à l'élève en français. Indique :
 - quelle partie de sa réponse est fausse ;
 - pourquoi elle est fausse ;
 - éventuellement, ce qu'il aurait fallu faire à la place.
Les commentaires peuvent contenir des fragments LaTeX, par exemple
`$a^2 + b^2 = c^2$`.

Si la note n'est pas 4, fournis toujours un commentaire expliquant ce qui
manque, sauf dans le cas `empty-answer` décrit ci-dessous.

Lorsqu'un commentaire concerne une erreur située dans une partie précise
de la réponse, tu peux fournir un champ `box_2d` pour la localiser.
Ses coordonnées doivent être de la forme [ymin, xmin, ymax, xmax],
normalisées entre 0 et 1000. Sinon, attribue la valeur `null` à `box_2d`.

Si la réponse est correcte, aucun commentaire n'est nécessaire. Tu n'es pas
obligé de faire des commentaires positifs ; si tu en fais, ne leur associe
pas de `box_2d`.

Par exemple, si l'élève affirme à tort qu'une fonction est continue,
localise le mot « continue ». Si un calcul est faux, localise l'étape où
l'erreur apparaît et explique cette erreur dans le commentaire.

Évite les commentaires portant sur une confusion entre les lettres `n`
et `m`, `x` et `n`, ou `h` et `k`. En cas de doute, suppose que tu as mal
lu, sauf si la distinction est très importante.

Certains cas nécessitent une valeur particulière du champ `error` :
 - L'élève n'a pas répondu à la bonne question : attribue la note 0 et
   indique `wrong-label`, car il peut s'agir d'une erreur de label.
 - La réponse contient aussi une réponse à une autre question de
   l'exercice, sur plus de quelques lignes : note la question demandée,
   mais indique `additional-answer`.
 - La réponse est vide, ou l'élève a seulement recopié l'énoncé : indique
   `empty-answer` et ne fournis aucun commentaire.
S'il n'y a aucune de ces erreurs, attribue la chaîne vide `""` à `error`.

Réponds uniquement en JSON, sous la forme d'une liste d'objets contenant
les clés `id` et `result`. L'objet `result` contient `score`, la liste
`feedback` et `error`. Chaque commentaire contient `text` et `box_2d`.
Conserve exactement ces clés, les identifiants et les valeurs techniques
de `error` : ne les traduis pas. Le contenu de chaque champ `text` doit
être en français, même si certains documents fournis sont dans une autre langue.

Exemple :
```json
[
  {
    "id": "01",
    "result": {
      "score": 2.5,
      "feedback": [
        {"text": "Il manque la vérification des hypothèses du théorème.", "box_2d": null},
        {"text": "Non, la fonction n'est pas forcément continue.", "box_2d": [145, 280, 340, 500]}
      ],
      "error": ""
    }
  },
  {
    "id": "04",
    "result": {"score": 4.0, "feedback": [], "error": ""}
  }
]
```

Voici l'énoncé de l'exercice ou la partie pertinente du problème :

```
<<text>>
```

Voici un corrigé possible :

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

Tu dois noter uniquement la question ou l'exercice portant le label
`<<label>>`. Ne note aucune autre question et ne formule aucun commentaire
sur les autres questions."""

from .utils import get_label_text_content, get_label_sol_content, get_label_persp_content

def make_prompt(input_dir,full_label):
    text = get_label_text_content(input_dir, full_label) or ""
    corr = get_label_sol_content(input_dir, full_label) or ""
    persp = get_label_persp_content(input_dir, full_label) or ""
    # print("Debug : l/t/c/p", full_label, text, corr, persp)

    if persp:
        persp = "\n\nVoici des consignes de notation complémentaires : \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="Commentaire destiné à l’élève, rédigé en français.")
    box_2d: Optional[List[int]] = Field(None, description="Coordonnées [ymin, xmin, ymax, xmax] normalisées entre 0 et 1000, ou null.")

class ResultData(BaseModel):
    score: float = Field(description="Note numérique de la réponse, sur 4 points.")
    feedback: List[FeedbackItem] = Field(description="Liste des commentaires destinés à l’élève, rédigés en français.")
    error: str = Field(description="Type d’erreur : wrong-label, additional-answer, empty-answer, ou chaîne vide.")

class EvaluationEntry(BaseModel):
    id: str = Field(description="Identifiant exact de la réponse.")
    result: ResultData = Field(description="Note, commentaires en français et éventuelle erreur pour cette réponse.")

# These nested definitions  do not work with the batch api, unroll them
UNROLLED_SCHEMA = {
                    "type": "ARRAY",
                    "items": {
                        "type": "OBJECT",
                        "properties": {
                            "id": {"type": "STRING", "description": "Identifiant exact de la réponse."},
                            "result": {
                                "type": "OBJECT",
                                "properties": {
                                    "score": {"type": "NUMBER", "description": "Note numérique de la réponse, sur 4 points."},
                                    "error": {"type": "STRING", "description": "Type d’erreur : wrong-label, additional-answer, empty-answer, ou chaîne vide."},
                                    "feedback": {
                                        "type": "ARRAY",
                                        "description": "Liste des commentaires destinés à l’élève, rédigés en français.",
                                        "items": {
                                            "type": "OBJECT",
                                            "properties": {
                                                "text": {"type": "STRING", "description": "Commentaire destiné à l’élève, rédigé en français."},
                                                "box_2d": {
                                                    "type": "ARRAY",
                                                    "items": {"type": "INTEGER"},
                                                    "nullable": True,
                                                    "description": "Coordonnées [ymin, xmin, ymax, xmax] normalisées entre 0 et 1000, ou null."
                                                }
                                            },
                                            "required": ["text"]
                                        }
                                    }
                                },
                                "required": ["score", "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(
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
        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"""Voici la réponse d'un élève à une question d'examen. Le JSON
ci-dessous contient des commentaires dont les rectangles de localisation
(`box_2d`) sont incorrects. Chaque commentaire doit correspondre à la
partie de la réponse où se trouve l'erreur signalée.

Par exemple, si l'élève affirme à tort qu'une fonction est continue,
les coordonnées doivent localiser le mot « continue ». Si un calcul est
faux, elles doivent localiser l'étape où apparaît l'erreur expliquée dans
le commentaire.

Analyse l'image et renvoie le même contenu JSON en corrigeant UNIQUEMENT
les coordonnées `box_2d` pour cette image. Conserve les commentaires en
français à l'identique : ne les reformule pas et ne les traduis pas.
Conserve les noms des clés JSON.
Les coordonnées doivent être [ymin, xmin, ymax, xmax], normalisées entre
0 et 1000. Si la zone est introuvable ou le rectangle invalide, renvoie
`null` pour ce rectangle.

Commentaires d'origine :
{json.dumps(localized_feedbacks, indent=2, ensure_ascii=False)}
"""



    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(
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
        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"""Cette image représente une partie de la réponse d'un élève à un examen.

Elle porte initialement le label '{label}', mais je soupçonne une erreur
de label. L'élève a peut-être lui-même écrit le mauvais label.

Analyse l'image et identifie le label de la question à laquelle cette
réponse correspond. Ne te fie pas au label écrit par l'élève : examine
le contenu de la réponse et les notations utilisées.

Renvoie UNIQUEMENT le label exact, sans le modifier ni le traduire.

Voici l'énoncé complet de l'examen :
{enonce}

Voici les labels possibles. Ta réponse doit être l'un d'entre eux :
{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(
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
        temperature=1.0,
    )
    return contents, config

def request_for_additional_answer(pdf_path, label, enonce, labels_txt):
    prompt = f"""Cette image représente une partie de la réponse d'un élève à un examen.

Elle porte initialement le label '{label}', mais je soupçonne qu'elle
contient aussi des réponses à une ou plusieurs autres questions.

Analyse l'image et identifie les labels des questions auxquelles elle
répond. Renvoie UNIQUEMENT une liste JSON contenant les labels exacts,
sans les modifier ni les traduire.

Si le bas de l'image ne contient que la première ligne d'une réponse à
une autre question, ignore cette ligne.

Voici l'énoncé complet de l'examen :
{enonce}

Voici les labels possibles. Chaque élément de ta liste doit être l'un
d'entre eux :
{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(
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
        temperature=1.0,
        response_mime_type="application/json",
    )
    return contents, config