from __future__ import annotations

import argparse
import re
import time
import typing
from collections import defaultdict
from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from google import genai
from google.genai import types
from pydantic import BaseModel, Field

import config
from copienator import (
    CliError,
    EvaluationWorkspace,
    ExitCode,
    atomic_write_json,
    execute,
    read_json,
    target_parser,
    workspace_from_target,
)
from utils import natural_key, read_all_labels

MODEL_ID = config.MODEL_FOR_LABEL_ID
api_key = config.API_KEY

my_prompt = """I'm giving you an image of the left columns of a written exam.
Students answer several exercises, which can have several questions.

The image consists of several columns, separated by vertical black
lines. The image should be read top to bottom and then left to right,
meaning first column, then second column, etc.

In their sheet, students delimit exercises and questions using
delimiters such as `Ex 1`, or `Exercice 1`, and `1)` or `a)`. You need
to give me the bounding boxes of each delimiter.

When giving the bounding box of the first question of an exercise, the
box should be large enough to contain both the exercice label
(`Exercice i`) and the question label (`1)`) parts. If they are
horizontally far apart (example : if the `1)` is to the left and the
`Exercice i` is either to the right, or in the middle) then give only
the bounding box of the question label `1)` part. You should still
label it as `Exercice i : 1)` though.

You also need to give me the student name. It should appear on the top
left of the image. Disregard any mention of `MPSI 3`, it is their
class. A list of possible student names will be given below.

You will answer with a JSON object, containing a `name` field with the
name, and a `list` field, with the list of the bounding boxes and
their labels. The box_2d should be [ymin, xmin, ymax, xmax] normalized
to 0-1000.

Here is an example :
{\"name\" : \"John Doe\", \"list\" : [{\"box_2d\": (10, 20, 30, 40), \"label\" : \"Ex 1 : 1)\"}]}

Do not provide a box_2d for the name. Only for the labels. Order the
box_2d by their position in the page, column by column : first column
(top to bottom), then second column, etc.

You may find the same label present several times, as a student either
recall the current label on a new page, or adds content to its answer
later on. Give the position of each instance of each label.

For this exam you should look for the labels given below, separated by
newlines. A student need not have answered every question, so some may
be missing.

##labels##

##wrong_labels##

Here's a list of the names of the students, pick the one that matches
the best or `\"Unknown\"` if you cannot read the name

##names##"""
my_prompt2 = """I'm giving you an image of the left columns of a written exam.
Students answer several exercises, which can have several questions.

The image consists of several columns, separated by vertical black
lines. The image should be read top to bottom and then left to right,
meaning first column, then second column, etc.

In their sheet, students delimit exercises and questions using
delimiters such as `Ex 1`, or `Exercice 1`, and `1)` or `a)`. You need
to give me the bounding boxes of each delimiter.

When giving the bounding box of the first question of an exercise, the
box should be large enough to contain both the exercice label
(`Exercice i`) and the question label (`1)`) parts.

You also need to give me the student name. It should appear on the top
left of the image. Disregard any mention of `MPSI 3`, it is their
class. A list of possible student names will be given below.

You will answer with a JSON object, containing a `name` field with the
name, and a `list` field, with the list of the bounding boxes and
their labels. The box_2d should be [ymin, xmin, ymax, xmax] normalized
to 0-1000.

Here is an example :
{\"name\" : \"John Doe\", \"list\" : [{\"box_2d\": (10, 20, 30, 40), \"label\" : \"Ex 1 : 1)\"}]}

Do not provide a box_2d for the name. Only for the labels.

You may find the same label present several times, as a student either
recall the current label on a new page, or adds content to its answer
later on. Give the position of each instance of each label.

This image is one part of a sequence (e.g., part 2 of 3) for a single
student. Here is the list of labels found in the *previous* parts of
this copy:

[
##prev_context##
]

If the first column starts with a number like =3)= or =c)=, look at
the labels in the list above. If the last relevant label was =Ex 4 :
2)=, you should label the new box =Ex 4 : 3)=.

For this exam you should look for the labels given below, separated by
newlines. A student need not have answered every question, so some may
be missing.

##labels##

##wrong_labels##

Since this copy isn't the first part of a sequence, simply set the
name to `\"Continued\"`."""

class BoxItem(BaseModel):
    box_2d: list[int] = Field(description="Bounding box coordinates (e.g., [ymin, xmin, ymax, xmax])")
    label: str = Field(description="The label associated with the specific box")

class AnnotationData(BaseModel):
    name: str = Field(description="The name identifier")
    list: typing.List[BoxItem] = Field(  # noqa: UP006 - field name shadows list
        description="List of bounding box items"
    )


def generate_request(file, labels, names, context_labels, wrong_labels):
    """Generates request for Gemini with context."""

    image_path = Path(file)

    # Format context list as a string
    context_str = ", ".join([f'"{l}"' for l in context_labels]) if context_labels else "No previous context"

    if context_labels == []:
        text = my_prompt.replace("##labels##", labels)\
                        .replace("##names##", names)
    else:
        text = my_prompt2.replace("##labels##", labels)\
                         .replace("##prev_context##", context_str)
    if wrong_labels:
        text= text.replace("##wrong_labels##\n\n", f"On a previous request, you answered with the following wrong labels : {wrong_labels}. These are wrong, since they do not exactly match any of the labels in the previous list.")
    else:
        text = text.replace("##wrong_labels##\n\n", "")


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

    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= AnnotationData.model_json_schema(),
    )
    return (contents, generate_content_config)

TARGET_INTERVAL = 3.5
Sleep = Callable[[float], None]


def selected_images(
    workspace: EvaluationWorkspace,
    targets: list[Path],
) -> tuple[list[Path], list[str]]:
    """Resolve evaluation, copy-PDF, or Cutleft-image targets."""
    workspace.require_directories("Copies", "Cutleft")
    images: list[Path] = []
    warnings: list[str] = []
    for target in targets:
        if target.is_dir():
            copy_pdfs = sorted(
                workspace.copies_dir.glob("Copie*.pdf"), key=natural_key
            )
            if not copy_pdfs:
                warnings.append(f"No Copie*.pdf files found in {workspace.copies_dir}")
            stems = [path.stem for path in copy_pdfs]
        elif target.suffix.casefold() in {".jpg", ".jpeg"}:
            if target.parent != workspace.cutleft_dir:
                raise CliError(
                    f"Image target is not in {workspace.cutleft_dir}: {target}",
                    ExitCode.INVALID_ARGUMENTS,
                )
            images.append(target)
            continue
        elif target.suffix.casefold() == ".pdf":
            stems = [target.stem]
        else:
            raise CliError(
                f"Unsupported target for label detection: {target}",
                ExitCode.INVALID_ARGUMENTS,
            )
        for stem in stems:
            found = sorted(
                workspace.cutleft_dir.glob(f"{stem}_*.jpg"), key=natural_key
            )
            if found:
                images.extend(found)
            else:
                warnings.append(
                    f"No Cutleft image variants found for {stem} in "
                    f"{workspace.cutleft_dir}"
                )
    return list(dict.fromkeys(images)), warnings


def group_images(image_files: list[Path]) -> dict[str, list[Path]]:
    groups: defaultdict[str, list[Path]] = defaultdict(list)
    for image in image_files:
        match = re.match(r"(.+)_(\d+)$", image.stem)
        groups[match.group(1) if match else image.stem].append(image)
    for files in groups.values():
        files.sort(key=natural_key)
    return dict(groups)


def _existing_context(output_json: Path) -> list[str]:
    try:
        loaded = read_json(output_json)
        if not isinstance(loaded, dict):
            return []
        return [
            str(item["label"])
            for item in loaded.get("list", [])
            if isinstance(item, dict) and "label" in item
        ]
    except (OSError, TypeError, ValueError):
        return []


def process_copy_group(
    workspace: EvaluationWorkspace,
    group_key: str,
    files: list[Path],
    *,
    client,
    labels_text: str,
    names_text: str,
    valid_labels: set[str],
    valid_names: set[str],
    overwrite: bool,
    sleep: Sleep = time.sleep,
    target_interval: float = TARGET_INTERVAL,
) -> int:
    """Process one student's image parts sequentially to preserve context."""
    accumulated_labels: list[str] = []
    generated = 0
    for image_file in files:
        started = time.monotonic()
        output_json = workspace.copies_dir / f"{image_file.stem}.json"
        if output_json.exists() and not overwrite:
            print(f"[{group_key}] Skipping {image_file.name}, output exists.")
            accumulated_labels.extend(_existing_context(output_json))
            continue

        print(
            f"[{group_key}] Processing {image_file.name} with "
            f"{len(accumulated_labels)} accumulated labels..."
        )
        attempt = 0
        wrong_labels: list[str] = []
        while True:
            if attempt > 0:
                sleep(10 * attempt)
            try:
                contents, request_config = generate_request(
                    image_file,
                    labels_text,
                    names_text,
                    accumulated_labels,
                    wrong_labels,
                )
                response = client.models.generate_content(
                    model=MODEL_ID,
                    contents=contents,
                    config=request_config,
                )
                annotation = AnnotationData.model_validate_json(response.text)
                unknown = [
                    item.label
                    for item in annotation.list
                    if item.label not in valid_labels
                ]
                if unknown:
                    print(
                        f"Error: {image_file.name} contained unknown labels: "
                        f"{unknown}"
                    )
                    wrong_labels.extend(unknown)
                    attempt += 1
                    continue
                if annotation.name not in valid_names:
                    print(
                        f"Error: {image_file.name} returned unknown name: "
                        f"{annotation.name}"
                    )
                    if attempt == 0:
                        attempt += 1
                        continue
                    annotation.name = "Unknown"

                atomic_write_json(output_json, annotation.model_dump())
                accumulated_labels.extend(box.label for box in annotation.list)
                generated += 1
                break
            except KeyboardInterrupt:
                raise
            except Exception as exc:  # noqa: BLE001 - remote API retry boundary
                print(
                    f"Error processing {image_file.name}: {exc}\n"
                    "\tIt will be retried."
                )
                attempt += 1
        sleep(max(0.0, target_interval - (time.monotonic() - started)))
    return generated


def run(
    workspace: EvaluationWorkspace,
    targets: list[Path],
    *,
    overwrite: bool = False,
    client=None,
    sleep: Sleep = time.sleep,
    max_workers: int = 12,
) -> ExitCode:
    workspace.require_files("labels")
    images, warnings = selected_images(workspace, targets)
    for warning in warnings:
        print(f"Warning: {warning}")
    if not images:
        return ExitCode.PARTIAL

    all_labels = read_all_labels(workspace.root)
    labels_text = "\n".join(all_labels) + "\n"
    names_path = workspace.names_file()
    if not names_path.is_file():
        raise CliError(f"Names file not found: {names_path}", ExitCode.INVALID_WORKSPACE)
    names_text = names_path.read_text(encoding="utf-8")
    valid_names = {
        line.strip() for line in names_text.splitlines() if line.strip()
    } | {"Unknown", "Continued"}
    if client is None:
        client = genai.Client(api_key=api_key)

    groups = group_images(images)
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [
            executor.submit(
                process_copy_group,
                workspace,
                group_key,
                files,
                client=client,
                labels_text=labels_text,
                names_text=names_text,
                valid_labels=set(all_labels),
                valid_names=valid_names,
                overwrite=overwrite,
                sleep=sleep,
            )
            for group_key, files in groups.items()
        ]
        for future in futures:
            future.result()
    return ExitCode.PARTIAL if warnings else ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    parser = target_parser("Detect handwritten question labels with Gemini")
    parser.add_argument(
        "additional_targets",
        nargs="*",
        type=Path,
        help="Additional copy PDFs or Cutleft images from the same evaluation",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        help="Regenerate JSON outputs that already exist",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()

    def handle(args: argparse.Namespace) -> ExitCode:
        workspace, target = workspace_from_target(
            args, repository=Path(__file__).resolve().parent
        )
        targets = [target]
        for additional in args.additional_targets:
            resolved = additional.expanduser().resolve()
            if not resolved.exists():
                raise CliError(
                    f"Target does not exist: {resolved}",
                    ExitCode.INVALID_WORKSPACE,
                )
            additional_workspace = EvaluationWorkspace.discover(
                resolved, repository=workspace.repository
            )
            if additional_workspace.root != workspace.root:
                raise CliError(
                    "All targets must belong to the same evaluation",
                    ExitCode.INVALID_ARGUMENTS,
                )
            targets.append(resolved)
        return run(workspace, targets, overwrite=args.overwrite)

    return execute(parser, argv, handle)


if __name__ == "__main__":
    raise SystemExit(main())