from __future__ import annotations

import argparse
import concurrent.futures
import os
from collections.abc import Sequence
from pathlib import Path
from typing import Any

from PIL import Image, ImageDraw
from reportlab.pdfgen import canvas

from copienator import (
    CliError,
    EvaluationWorkspace,
    ExitCode,
    atomic_write_json,
    atomic_write_text,
    evaluation_parser,
    execute,
    read_json,
    utils,
    workspace_from_args,
)
from copienator.annotation_data import load_annotation_data
from copienator.annotation_documents import MAX_HEIGHT_PX
from copienator.commands import annotating, annotating_with_checks
from copienator.commands.import_annotations import clear_import_directory
from copienator.filesystem import staged_directory
from copienator.utils import natural_key


def render_item(item):
    student_id, label, content = item
    pdf_path = Path(content["pdf_path"])
    if not pdf_path.exists():
        print(f"Warning: answer PDF not found: {pdf_path}")
        return None
    base_image, _, _ = annotating.make_base_image(pdf_path)
    checkbox_renderer = annotating_with_checks.CheckboxRenderer(label)
    final_image, header_height = annotating.compose_label_image(
        base_image,
        label,
        content["result"],
        content["coordinates"][0],
        draw_callback=checkbox_renderer.callback,
        more_right=True,
        with_id=student_id,
        source_pdf=pdf_path,
    )
    if final_image is None:
        return None
    return (
        student_id,
        label,
        final_image,
        header_height,
        checkbox_renderer.checkboxes,
    )


def save_batch(batch, prefix, group_id, output_root: Path) -> None:
    output_dir = output_root / f"{prefix} G{group_id}"
    print(f"Generating group PDF: {prefix} G{group_id} ({len(batch)} elements)")
    max_width = max(item[2].width for item in batch)
    total_height = sum(item[2].height for item in batch)
    if total_height > MAX_HEIGHT_PX and len({item[0] for item in batch}) > 1:
        raise CliError(
            f"{prefix} G{group_id} would be {total_height} px high; "
            f"only a single student's block may exceed {MAX_HEIGHT_PX} px."
        )
    concatenated = Image.new("RGB", (max_width, total_height), "white")
    draw = ImageDraw.Draw(concatenated)
    checkbox_map: list[dict[str, Any]] = []
    bnote_entries: list[dict[str, Any]] = []
    current_y = 0
    previous_student = None

    for student_id, label, image, header_height, checkboxes in batch:
        concatenated.paste(image, (0, current_y))
        if student_id != previous_student:
            draw.rectangle([0, current_y, max_width, current_y + 4], fill="purple")
        previous_student = student_id
        bnote_entries.append(
            {
                "id": student_id,
                "label": label,
                "header_height": header_height,
                "layout": image.info.get('annotation_layout', annotating.layout_for(None)),
                "hmin": current_y,
                "hmax": current_y + image.height,
            }
        )
        for item in checkboxes:
            box = item.get("final_box") or item.get("rel_box")
            item["global_box"] = [
                box[0],
                box[1] + current_y,
                box[2],
                box[3] + current_y,
            ]
            item["student_id"] = student_id
            checkbox_map.append(item)
        current_y += image.height

    with staged_directory(output_dir) as staging:
        atomic_write_json(
            staging / "bnote.json",
            {"width": max_width, "height": total_height, "images": bnote_entries},
        )
        atomic_write_json(staging / "checkboxes.json", checkbox_map)
        reference = staging / "Reference.jpg"
        concatenated.save(reference, quality=90)
        pdf_path = staging / "Concat.pdf"
        pdf_canvas = canvas.Canvas(str(pdf_path), pagesize=(max_width, total_height))
        pdf_canvas.drawImage(
            str(reference),
            0,
            0,
            width=max_width,
            height=total_height,
        )
        pdf_canvas.save()


def _initial_label_groups(labels: list[str]) -> str:
    groups: dict[str, list[str]] = {}
    for label in labels:
        key = label.split(" : ")[0] if " : " in label else label
        groups.setdefault(key, []).append(label)
    return "".join(",".join(items) + "\n" for items in groups.values())


def _gemini_label_groups(
    workspace: EvaluationWorkspace, labels: list[str]
) -> list[list[str]] | None:
    source = workspace.gemini_exam_items_file
    if not source.is_file():
        return None

    groups: list[list[str]] = []
    current: list[str] = []
    try:
        source_lines = source.read_text(encoding="utf-8").splitlines()
    except (OSError, UnicodeError) as exc:
        print(f"Warning: could not read Gemini question groups from {source}: {exc}")
        return None
    for raw_line in source_lines:
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        if line == "---":
            if current:
                groups.append(current)
                current = []
            continue
        if " ### " not in line:
            continue
        label = line.split(" ### ", 1)[0].strip()
        if label and label != "CONTEXT":
            current.append(label)
    if current:
        groups.append(current)

    flattened = [label for group in groups for label in group]
    known = set(labels)
    if (
        not flattened
        or len(flattened) != len(set(flattened))
        or set(flattened) != known
    ):
        missing = sorted(known.difference(flattened), key=natural_key)
        unknown = sorted(set(flattened).difference(known), key=natural_key)
        details = []
        if missing:
            details.append("missing: " + ", ".join(missing))
        if unknown:
            details.append("unknown: " + ", ".join(unknown))
        if len(flattened) != len(set(flattened)):
            details.append("duplicate labels")
        print(
            f"Warning: ignoring incompatible Gemini question groups in {source}"
            + (f" ({'; '.join(details)})" if details else "")
        )
        return None
    return groups


def _serialize_label_groups(groups: list[list[str]]) -> str:
    return "".join(",".join(group) + "\n" for group in groups)


def _load_label_groups(
    workspace: EvaluationWorkspace, labels: list[str]
) -> list[list[str]]:
    label_groups = workspace.label_groups_file
    if not label_groups.exists():
        gemini_groups = _gemini_label_groups(workspace, labels)
        if gemini_groups is not None:
            initial_content = _serialize_label_groups(gemini_groups)
            source_description = "the groups selected in gemini_for_enonce.py"
        else:
            initial_content = _initial_label_groups(labels)
            source_description = "the label-prefix fallback"
        atomic_write_text(label_groups, initial_content)
        print(
            f"Created {label_groups} from {source_description}; "
            "review the groups before continuing."
        )
        utils.edit_file_and_enter(label_groups)
    known_labels = set(labels)
    groups: list[list[str]] = []
    unknown: set[str] = set()
    for line in label_groups.read_text(encoding="utf-8").splitlines():
        group = [label.strip() for label in line.split(",") if label.strip()]
        unknown.update(label for label in group if label not in known_labels)
        if group:
            groups.append(group)
    if unknown:
        raise CliError(
            "label_groups contains unknown labels: " + ", ".join(sorted(unknown))
        )
    return groups


def _unique_prefixes(groups: list[list[str]]) -> list[tuple[str, list[str]]]:
    used: set[str] = set()
    result: list[tuple[str, list[str]]] = []
    previous: str | None = None
    for labels in groups:
        safe_labels = [label.replace(":", "").strip() for label in labels]
        base = os.path.commonprefix(safe_labels).strip() or "Group"
        if base and previous is not None and natural_key(base) < natural_key(previous):
            candidate = f"{safe_labels[0]}+"
            if natural_key(candidate) > natural_key(previous):
                base = candidate
        base = base.removesuffix("i")
        # Keep related groups under the same short stem (I), I)-2, ...).
        # Normalize before allocating a suffix so no later rewrite can collide.
        if previous and base.startswith(previous):
            base = previous
        prefix = base
        counter = 2
        while prefix in used:
            prefix = f"{base}-{counter}"
            counter += 1
        previous = base
        used.add(prefix)
        result.append((prefix, labels))
    return result


def _existing_group_state(
    output_root: Path,
    prefix: str | None = None,
) -> tuple[set[tuple[str, str]], int]:
    existing_items: set[tuple[str, str]] = set()
    maximum_group = 0
    if not output_root.is_dir():
        return existing_items, maximum_group
    for directory in output_root.iterdir():
        if not directory.is_dir():
            continue
        if prefix is not None and not directory.name.startswith(f"{prefix} G"):
            continue
        try:
            maximum_group = max(maximum_group, int(directory.name.split(" G")[-1]))
        except ValueError:
            continue
        metadata = directory / "bnote.json"
        if not metadata.exists():
            continue
        loaded = read_json(metadata)
        if not isinstance(loaded, dict):
            raise TypeError(f"Expected a JSON object in {metadata}")
        for image in loaded.get("images", []):
            existing_items.add((str(image["id"]), str(image["label"])))
    return existing_items, maximum_group


def split_batches(rendered):
    batches = []
    heights = []

    index = 0
    while index < len(rendered):
        student_id = rendered[index][0]
        end = index + 1
        while end < len(rendered) and rendered[end][0] == student_id:
            end += 1
        student_items = rendered[index:end]
        student_height = sum(item[2].height for item in student_items)

        if student_height > MAX_HEIGHT_PX:
            # Oversized blocks remain intact in a group of their own.
            batches.append(list(student_items))
            heights.append(student_height)
            index = end
            continue

        # Reuse space in earlier groups, keeping each student's block intact.
        for batch_index, height in enumerate(heights):
            if height + student_height <= MAX_HEIGHT_PX:
                batches[batch_index].extend(student_items)
                heights[batch_index] += student_height
                break
        else:
            batches.append(list(student_items))
            heights.append(student_height)
        index = end

    return batches


def _generate_groups(
    output_root: Path,
    data,
    groups: list[list[str]],
    *,
    resume: bool,
) -> tuple[int, bool]:
    generated = 0
    problems = False
    # Resume by answer identity across all prefixes, including legacy names.
    # Changing the naming scheme must not regenerate already reviewed PDFs.
    existing_items = _existing_group_state(output_root)[0] if resume else set()
    for prefix, labels in _unique_prefixes(groups):
        _, maximum_group = (
            _existing_group_state(output_root, prefix) if resume else (set(), 0)
        )
        items = [
            (student_id, label, student_labels[label])
            for student_id, student_labels in data.items()
            for label in labels
            if label in student_labels and (student_id, label) not in existing_items
        ]
        if not items:
            continue
        items.sort(key=lambda item: (natural_key(item[0]), natural_key(item[1])))
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            rendered = list(executor.map(render_item, items))
        if any(item is None for item in rendered):
            problems = True
        successful = [item for item in rendered if item is not None]
        for index, batch in enumerate(split_batches(successful), start=1):
            save_batch(batch, prefix, maximum_group + index, output_root)
            generated += 1
    return generated, problems


def run(
    workspace: EvaluationWorkspace, *, overwrite: bool = False, refaire: bool = False,
    clear_import: bool = False,
) -> ExitCode:
    workspace.require_files("labels", "correction.json")
    workspace.require_directories("Copies", "Par label")
    labels = utils.read_all_labels(workspace.root)
    refaire_list = annotating_with_checks._load_refaire(workspace) if refaire else None
    loaded = load_annotation_data(workspace, refaire_list=refaire_list)
    groups = (
        [
            [label]
            for label in sorted(
                {label for answers in loaded.data.values() for label in answers},
                key=natural_key,
            )
        ]
        if refaire
        else _load_label_groups(workspace, labels)
    )
    for warning in loaded.warnings:
        print(f"Warning: {warning}")
    if not loaded.data:
        print("Warning: no annotation data was found.")
        return ExitCode.PARTIAL

    output_root = workspace.annotation_dir("refaire" if refaire else "grouped")
    if refaire and output_root.exists() and not overwrite:
        raise CliError(
            "BRnot already exists; use --overwrite to replace the previous redo."
        )
    if overwrite or refaire:

        class IncompleteGroupedOutput(Exception):
            pass

        try:
            with staged_directory(output_root) as staging:
                generated, problems = _generate_groups(
                    staging,
                    loaded.data,
                    groups,
                    resume=False,
                )
                if generated == 0 or problems or loaded.warnings:
                    raise IncompleteGroupedOutput
        except IncompleteGroupedOutput:
            print(
                f"Warning: grouped overwrite was incomplete; previous {output_root.name} was preserved."
            )
            return ExitCode.PARTIAL
    else:
        output_root.mkdir(parents=True, exist_ok=True)
        generated, problems = _generate_groups(
            output_root,
            loaded.data,
            groups,
            resume=True,
        )
        if problems:
            return ExitCode.PARTIAL
        if generated == 0:
            print("No new grouped annotations were required.")
    if loaded.warnings:
        return ExitCode.PARTIAL
    if clear_import:
        clear_import_directory(workspace)
    return ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    parser = evaluation_parser("Generate annotated PDFs grouped by labels.")
    parser.add_argument("--clear-import", action="store_true",
                        help="Empty IMPORT_DIR after successful annotation generation")
    parser.add_argument(
        "--overwrite", action="store_true", help="Replace annotation outputs safely"
    )
    parser.add_argument(
        "--refaire",
        action="store_true",
        help="Group only the answers in refaire.json, writing to BRnot",
    )
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    return execute(
        parser,
        argv,
        lambda args: run(
            workspace_from_args(args), overwrite=args.overwrite, refaire=args.refaire,
            clear_import=args.clear_import,
        ),
    )


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