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

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

MAX_HEIGHT_PX = 25000


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,
    )
    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)
    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,
                "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 _load_label_groups(workspace: EvaluationWorkspace, labels: list[str]) -> list[list[str]]:
    label_groups = workspace.root / "label_groups"
    if not label_groups.exists():
        atomic_write_text(label_groups, _initial_label_groups(labels))
        print(f"Created {label_groups}; 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
        prefix = base.removesuffix("i")
        counter = 2
        while prefix in used:
            prefix = f"{base}-{counter}"
            counter += 1
        if counter == 2 and previous and previous in prefix:
            prefix = f"{previous}-{counter}"
        elif counter == 2:
            previous = prefix
        used.add(prefix)
        result.append((prefix, labels))
    return result


def _existing_group_state(
    output_root: Path,
    prefix: str,
) -> 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() or 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):
    def split(maximum_height: float):
        batches = []
        current = []
        current_height = 0
        previous_student = None
        for item in rendered:
            student_id = item[0]
            image_height = item[2].height
            if (
                current
                and current_height + image_height > maximum_height
                and student_id != previous_student
            ):
                batches.append(current)
                current = []
                current_height = 0
            current.append(item)
            current_height += image_height
            previous_student = student_id
        if current:
            batches.append(current)
        return batches

    strict = split(MAX_HEIGHT_PX)
    relaxed = split(1.1 * MAX_HEIGHT_PX)
    return relaxed if len(relaxed) < len(strict) else strict


def _generate_groups(
    output_root: Path,
    data,
    groups: list[list[str]],
    *,
    resume: bool,
) -> tuple[int, bool]:
    generated = 0
    problems = False
    for prefix, labels in _unique_prefixes(groups):
        existing_items, 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) -> ExitCode:
    workspace.require_files("labels", "correction.json")
    workspace.require_directories("Copies", "Par label")
    labels = utils.read_all_labels(workspace.root)
    groups = _load_label_groups(workspace, labels)
    loaded = load_annotation_data(workspace)
    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("grouped")
    if overwrite:
        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("Warning: grouped overwrite was incomplete; previous BGnot 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.")
    return ExitCode.PARTIAL if loaded.warnings else ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    parser = evaluation_parser("Generate annotated PDFs grouped by labels.")
    parser.add_argument("--overwrite", action="store_true", help="Replace BGnot safely")
    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),
    )


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