"""Reproducible edits to reviewed answer bodies, never to headers or supplements."""
from __future__ import annotations

import re
from copy import deepcopy
from hashlib import sha256
from pathlib import Path

from PIL import Image, ImageDraw

from copienator import CliError, atomic_write_json, read_json, utils
from copienator.commands.resolve_manual import ManualInstruction, parse_instruction_text

DRAFT = "annotation_modifications.txt"
APPLIED = "annotation_modifications_applied.txt"
BASELINE = "annotation_baseline.json"
MODES = ("BGnot", "Bnot")


def read_instructions(text: str, mode: str) -> list[ManualInstruction]:
    declared = re.findall(r"^### Source:\s*(\S+)\s*$", text, re.MULTILINE)
    if declared != [mode] or mode not in MODES:
        raise CliError(f"Le fichier doit contenir exactement une ligne '### Source: {mode}'.")
    return parse_instruction_text(text)


def normalize_copy(root: Path, value: str) -> str:
    number = value.strip().removeprefix("Copie")
    if not number.isdecimal():
        raise CliError("Numéro de copie invalide.")
    matches = [p.name.removeprefix("Copie") for p in (root / "Copies").glob("Copie*")
               if p.is_dir() and p.name.removeprefix("Copie").isdecimal()
               and int(p.name.removeprefix("Copie")) == int(number)]
    if len(matches) != 1:
        raise CliError(f"Copie {value} introuvable ou ambiguë.")
    return matches[0]


def baseline(root: Path) -> dict:
    path = root / BASELINE
    if not path.is_file():
        raise CliError(f"Relancez « Lire les annotations manuscrites » pour {root.name} : "
                       "la séparation exacte entre en-tête et réponse est indisponible.")
    return read_json(path)


def load_image(directory: Path, entry: dict) -> Image.Image:
    filename = entry["file"]
    if Path(filename).name != filename:
        raise CliError("Chemin d’image de référence invalide.")
    with Image.open(directory / filename) as image:
        result = image.convert("RGB").copy()
    height = entry["header_height"]
    if not isinstance(height, int) or not 0 <= height < result.height:
        raise CliError("Séparation en-tête/réponse invalide ; relancez la lecture des annotations.")
    return result


def body_image(directory: Path, label: str) -> Image.Image:
    entry = baseline(directory)["images"].get(label)
    if entry is None:
        raise CliError(f"Pas de réponse annotée pour {label} dans {directory.name}.")
    image = load_image(directory, entry)
    return image.crop((0, entry["header_height"], image.width, image.height))


def validate_instructions(root: Path, mode: str, instructions: list[ManualInstruction]) -> None:
    labels = set(utils.read_all_labels(root))
    sources = set()
    targets = set()
    for item in instructions:
        for label in (item.old_label, item.new_label):
            if label not in labels or Path(label).name != label or label in {".", ".."}:
                raise CliError(f"Label inconnu ou invalide : {label}")
        if item.old_label == item.new_label:
            raise CliError("La source et la cible doivent être différentes.")
        if normalize_copy(root, item.copy_id) != item.copy_id:
            raise CliError(f"Utilisez Copie{normalize_copy(root, item.copy_id)} dans le fichier.")
        key = (item.copy_id, item.old_label)
        if key in sources:
            raise CliError("Une seule instruction par label source et copie est autorisée.")
        sources.add(key)
        targets.add((item.copy_id, item.new_label))
        directory = root / mode / f"Copie{item.copy_id}"
        source = body_image(directory, item.old_label)
        if item.cut:
            y = round(source.height * item.cut[0] / 100)
            if not 0 < y < source.height:
                raise CliError("Chaque partie de la coupe doit contenir au moins un pixel.")
        target = baseline(directory)["images"].get(item.new_label)
        if target:
            load_image(directory, target)
        # Also report damaged current outputs instead of silently replacing them.
        for label in (item.old_label, item.new_label):
            path = directory / f"{label}.jpg"
            if path.exists():
                with Image.open(path) as image:
                    image.load()
    if sources & targets:
        raise CliError("Une cible ne peut pas également être source dans le même fichier.")


def stack(images: list[Image.Image]) -> Image.Image | None:
    images = [image for image in images if image.height]
    if not images:
        return None
    result = Image.new("RGB", (max(i.width for i in images), sum(i.height for i in images)), "white")
    y = 0
    for image in images:
        result.paste(image, (0, y))
        y += image.height
    return result


def edited_bodies(directory: Path, meta: dict, instructions: list[ManualInstruction]):
    headers, bodies = {}, {}
    for label, entry in meta["images"].items():
        image = load_image(directory, entry)
        h = entry["header_height"]
        headers[label] = image.crop((0, 0, image.width, h))
        bodies[label] = image.crop((0, h, image.width, image.height))
    originals = bodies.copy()
    changed = set()
    for item in instructions:
        source, target = item.old_label, item.new_label
        if source not in originals:
            raise CliError(f"Source annotée absente : {source}. Les modifications n’ont pas été appliquées.")
        piece = originals[source]
        if item.cut:
            percent, keep = item.cut
            y = round(piece.height * percent / 100)
            if not 0 < y < piece.height:
                raise CliError(f"Coupe invalide pour {source}.")
            top, bottom = piece.crop((0, 0, piece.width, y)), piece.crop((0, y, piece.width, piece.height))
            bodies[source], piece = (top, bottom) if keep == 1 else (bottom, top)
            changed.add(source)
        elif item.operator.startswith("x"):
            bodies.pop(source, None)
            changed.add(source)
        if item.should_merge:
            parts = ([bodies[target]] if target in bodies else [])
            parts.insert(0 if item.pipe_first else len(parts), piece)
            bodies[target] = stack(parts)
            changed.add(target)
        elif item.should_copy:
            bodies[target] = piece
            changed.add(target)
        elif item.operator.endswith("x"):
            bodies.pop(target, None)
            changed.add(target)
    return headers, bodies, changed


def output_removals(directory: Path) -> tuple[str, ...]:
    info = read_json(directory / "info.json") if (directory / "info.json").exists() else {}
    return ("Concat.jpg", "Concat_F.jpg", "Concat_F.pdf", "touched.json", "answer_labels.json",
            *(f"{label}.jpg" for label, entry in info.items()
              if entry.get("not_empty") and Path(label).name == label))


def rebuild(root: Path, mode: str, directory: Path, staging: Path, meta: dict,
            instructions: list[ManualInstruction], scores: dict) -> None:
    # Imported lazily: both readers use this module while saving their baseline.
    from copienator.commands import annotating
    from copienator.commands.reading_grouped_annotations import (
        get_extra_pdfs_as_images,
        save_paginated_pdf,
    )

    headers, bodies, changed = edited_bodies(directory, meta, instructions)
    info = deepcopy(meta["info"])
    filtered = set(meta["filtered"]) | changed
    all_labels = utils.read_all_labels(root)
    images, groups = [], []
    seen_contexts = set()
    order = sorted(bodies, key=utils.natural_key)
    for label in order:
        if label not in bodies:
            continue
        body = bodies[label]
        header = headers.get(label)
        if header is None:
            # No prior answer means no prior header. Do not invent a score or
            # resurrect Gemini feedback from correction.json.
            header = Image.new("RGB", (max(body.width, 800), 80), "white")
            ImageDraw.Draw(header).text((20, 20), f"{label}    Note : {scores.get(label, '')}",
                                        fill="black", font_size=30)
        image = stack([header, body])
        image.save(staging / f"{label}.jpg")
        images.append(image)
        if label in filtered:
            extras = get_extra_pdfs_as_images(root, label, annotating, all_labels,
                                            seen_contexts=seen_contexts) if mode == "BGnot" else []
            groups.append([*extras, image])
    for label, entry in info.items():
        entry["not_empty"] = label in bodies
        if label in changed and label in bodies:
            entry["present"] = True
        entry["touched"] = mode == "BGnot" and label in bodies and label in filtered
        entry["score"] = scores.get(label, entry["score"])
    atomic_write_json(staging / "info.json", info)
    if images:
        stack(images).save(staging / "Concat.jpg")
    if groups:
        if mode == "BGnot":
            save_paginated_pdf(groups, staging / "Concat_F.pdf")
        else:
            stack([image for group in groups for image in group]).save(staging / "Concat_F.jpg")


def stage_baseline(root: Path, mode: str, output: Path, staging: Path,
                   images: dict[str, Image.Image], heights: dict[str, int],
                   info: dict, filtered: list[str], scores: dict) -> None:
    """Save exact lossless reviewed layers, then replay only applied edits."""
    if mode not in MODES:
        return
    applied = output.parent / APPLIED
    instructions = read_instructions(applied.read_text(encoding="utf-8"), mode) if applied.is_file() else []
    selected = [i for i in instructions if f"Copie{i.copy_id}" == output.name]
    if not images.keys() <= heights.keys():
        if selected:
            raise CliError("Relancez la lecture principale des annotations avant de réappliquer les modifications.")
        # Legacy selective redo can still preserve its old images. A full read is
        # required before body editing; never guess the old header boundaries.
        return
    meta = {"version": 1, "images": {}, "info": info, "filtered": filtered}
    for label, image in images.items():
        filename = ".annotation-base-" + sha256(label.encode()).hexdigest() + ".png"
        image.save(staging / filename)
        meta["images"][label] = {"file": filename, "header_height": heights[label]}
    atomic_write_json(staging / BASELINE, meta)
    if selected:
        # Remove freshly staged aggregates before rebuilding; empty filtered
        # output must not retain the reader's previous version.
        for name in ("Concat.jpg", "Concat_F.jpg", "Concat_F.pdf", *(f"{label}.jpg" for label in images)):
            (staging / name).unlink(missing_ok=True)
        rebuild(root, mode, staging, staging, meta, selected, scores)