"""First-page name crops, numbered contact sheets, and roster-checked recognition."""
from __future__ import annotations

from pathlib import Path
import math

from google.genai import types
from PIL import Image, ImageDraw
from pydantic import BaseModel, Field, field_validator

from copienator import atomic_write_json, read_json
from copienator.annotation_layout import fingerprint
from copienator.filesystem import staged_files
from copienator.gemini_config import fresh_seed, thinking_config
from copienator.utils import natural_key

COLUMNS, ROWS = 6, 4
MARGIN_MM = 28
MISSING = "Missing"


def save_crop(page: Image.Image, pdf: Path, directory: Path, dpi: int) -> None:
    """Crop independently of the adjustable label strip, from the true top left."""
    width = min(page.width, max(1, round(MARGIN_MM * dpi / 25.4)))
    crop = page.crop((0, 0, width, max(1, page.height // 4)))
    with staged_files(directory) as staging:
        crop.save(staging / f"{pdf.stem}.jpg", "JPEG", quality=95)
        atomic_write_json(staging / f"{pdf.stem}.json", {
            "source_sha256": fingerprint(pdf), "width_mm": MARGIN_MM,
            "height_fraction": crop.height / page.height,
        })


def build_grids(directory: Path, copies_dir: Path, *, columns: int = COLUMNS,
                rows: int = ROWS, stems: set[str] | None = None) -> list[dict]:
    """Rebuild numbered sheets, excluding deleted or changed copies."""
    crops = []
    for path in sorted(directory.glob("*.jpg"), key=natural_key):
        if stems is not None and path.stem not in stems:
            continue
        pdf = copies_dir / f"{path.stem}.pdf"
        try:
            if pdf.is_file() and read_json(path.with_suffix('.json'))['source_sha256'] == fingerprint(pdf):
                crops.append(path)
        except (OSError, ValueError, KeyError, TypeError):
            continue
    grids_dir = directory / "Grids"
    grids = []
    with staged_files(grids_dir, remove=[p.name for p in grids_dir.glob('*') if p.is_file()]) as staging:
        for offset in range(0, len(crops), columns * rows):
            batch = crops[offset:offset + columns * rows]
            images = []
            for path in batch:
                with Image.open(path) as image:
                    images.append(image.convert('RGB'))
            width = max(image.width for image in images)
            height = max(image.height for image in images)
            border, header = 6, 32
            cell_width, cell_height = width + border, height + header + border
            sheet = Image.new('RGB', (columns * cell_width + border, rows * cell_height + border), 'black')
            draw = ImageDraw.Draw(sheet)
            cells = []
            for index in range(columns * rows):
                row, column = divmod(index, columns)
                x, y = border + column * cell_width, border + row * cell_height
                draw.rectangle((x, y, x + width - 1, y + height + header - 1), fill='white')
                if index < len(images):
                    sheet.paste(images[index], (x, y + header))
                    draw.text((x + 5, y + 5), f"R{row + 1} C{column + 1} | {index + 1}", fill='black')
                    cells.append({'cell': index + 1, 'row': row + 1, 'column': column + 1,
                                  'copy': batch[index].stem})
                else:
                    draw.text((x + 5, y + 5), 'EMPTY', fill='black')
            filename = f"grid_{len(grids) + 1:03d}.jpg"
            sheet.save(staging / filename, 'JPEG', quality=95)
            grids.append({'image': filename, 'columns': columns, 'rows': rows, 'cells': cells})
        atomic_write_json(staging / 'manifest.json', grids)
    return grids


class NameCell(BaseModel):
    cell: int = Field(strict=True, ge=1, le=24)
    name: str
    box_2d: list[float] | None = Field(
        default=None, description="[ymin, xmin, ymax, xmax], 0–1000 relative to the individual crop, excluding its header")

    @field_validator('box_2d', mode='before')
    @classmethod
    def valid_box(cls, value):
        # Bad localization must not discard an otherwise readable name.
        if (not isinstance(value, list) or len(value) != 4
                or any(type(v) not in (int, float) or not math.isfinite(v)
                       or not 0 <= v <= 1000 for v in value)):
            return None
        return value if value[0] < value[2] and value[1] < value[3] else None


class NameResponse(BaseModel):
    cells: list[NameCell]


def name_bottom_height_fraction(workspace, stem: str) -> float | None:
    """Read the name/class bottom on page 1, in cutleft's 0–1 height units.

    Returns None when localization is unavailable or the source/crop changed.
    This is intended as a future lower boundary for margin analysis.
    """
    directory = workspace.cutleft_dir / 'Names'
    for candidate in (directory, directory / 'Wide'):
        try:
            cached = read_json(candidate / f'{stem}_result.json')
            value = cached['name_bottom_height_fraction']
            if (cached['source_sha256'] == fingerprint(workspace.copies_dir / f'{stem}.pdf')
                    and cached['crop_sha256'] == fingerprint(candidate / f'{stem}.jpg')
                    and cached['name'] != MISSING
                    and type(value) in (int, float) and 0 < value <= 1):
                return float(value)
        except (OSError, ValueError, KeyError, TypeError):
            pass
    return None


def _identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool,
                    thinking_level: str | None = "low", wide: bool = False) -> dict[str, str]:
    """Identify pending copies before label requests; safely fall back on any failure."""
    directory = workspace.cutleft_dir / 'Names'
    if wide:
        directory /= 'Wide'
    result = {stem: MISSING for stem in groups}
    if not groups or not directory.is_dir():
        return result  # Older workspaces retain their existing label/name workflow.
    grids = build_grids(directory, workspace.copies_dir, columns=1 if wide else COLUMNS,
                        rows=4 if wide else ROWS, stems=set(groups) if wide else None)
    roster = {line.strip() for line in names_text.splitlines() if line.strip()} - {MISSING, 'Unknown', 'Continued'}
    for grid in grids:
        selected = {cell['copy']: cell for cell in grid['cells'] if cell['copy'] in groups}
        pending = {}
        for stem, cell in selected.items():
            crop_hash = fingerprint(directory / f'{stem}.jpg')
            cache_path = directory / f'{stem}_result.json'
            try:
                cached = read_json(cache_path)
                if (not overwrite and cached['crop_sha256'] == crop_hash
                        and 'name_bottom_height_fraction' in cached
                        and cached['names_text'] == names_text
                        and cached['name'] in roster | {MISSING}):
                    result[stem] = cached['name']
                    continue
            except (OSError, ValueError, KeyError, TypeError):
                pass
            pending[stem] = (cell, crop_hash, cache_path)
        if not pending:
            continue
        positions = '\n'.join(f"Cell {c['cell']}: row {c['row']}, column {c['column']}" for c in grid['cells'])
        region = 'full-width top quarters' if wide else 'top-left corners'
        prompt = f"""Identify student names from the {region} of exam first pages.
The image is a grid with exactly {grid['columns']} columns and {grid['rows']} rows. Read rows left to right,
then top to bottom. Thick black borders separate cells; headers show row, column,
and cell number. Each cell belongs to a DIFFERENT student. Never carry a name
from one cell into another. Do not include class names such as MPSI 3 in the name.
Occupied cells are exactly:
{positions}
All other cells are EMPTY padding: do not return them.
Return JSON {{"cells": [{{"cell": 1, "name": "exact authorized name", "box_2d": [100, 50, 400, 900]}}]}}.
For each identified student, box_2d must tightly enclose the handwritten name AND
any adjacent class designation, including a class written just below the name.
Use [ymin, xmin, ymax, xmax] coordinates normalized to 0–1000 relative to that
individual page crop, NOT the whole grid. Exclude the numbered header, black
borders and any padding outside the crop. Return box_2d=null for Missing names
or if the position cannot be determined reliably.
Return one entry per occupied cell, using its cell number. Match only confidently
readable handwriting to an exact authorized name below. Do not guess or assign
names by elimination. Return "Missing" for an unreadable, absent, or ambiguous name.
Authorized names:
{names_text}"""
        print(f"Identifying student names{' (full-width retry)' if wide else ''}: {grid['image']} ({len(grid['cells'])} occupied cells, model={model}, thinking={thinking_level or 'default'}).", flush=True)
        try:
            response = client.models.generate_content(
                model=model,
                contents=[types.Content(role='user', parts=[
                    types.Part.from_bytes(data=(directory / 'Grids' / grid['image']).read_bytes(), mime_type='image/jpeg'),
                    types.Part.from_text(text=prompt),
                ])],
                config=types.GenerateContentConfig(
                    http_options=types.HttpOptions(timeout=80_000, retry_options=types.HttpRetryOptions(attempts=1)),
                    automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
                    thinking_config=thinking_config(thinking_level), seed=fresh_seed(),
                    response_mime_type='application/json', response_json_schema=NameResponse.model_json_schema(),
                ),
            )
            entries = NameResponse.model_validate_json(response.text).cells
            by_cell = {}
            for entry in entries:
                by_cell.setdefault(entry.cell, []).append(entry)
            for stem, (cell, crop_hash, cache_path) in pending.items():
                candidates = by_cell.get(cell['cell'], [])
                entry = candidates[0] if len(candidates) == 1 else None
                name = entry.name if entry is not None and entry.name in roster else MISSING
                box = entry.box_2d if name != MISSING else None
                metadata = read_json(directory / f'{stem}.json')
                bottom = box[2] / 1000 * metadata['height_fraction'] if box else None
                result[stem] = name
                atomic_write_json(cache_path, {
                    'name': name, 'crop_sha256': crop_hash, 'names_text': names_text,
                    'source_sha256': metadata['source_sha256'],
                    'box_2d': box,  # 0–1000 coordinates relative to the name crop.
                    'name_bottom_height_fraction': bottom,  # Fraction of full page 1.
                })
        except Exception as exc:  # Remote failures must preserve the existing fallback.
            print(f"Name recognition failed for {grid['image']}: {type(exc).__name__}: {exc}; using label/name fallback.")
            # Do not cache transport/schema failures: a resumed run should retry.
    return result


def identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool,
                   thinking_level: str | None = "low") -> dict[str, str]:
    """Try narrow grids, then collated full-width headers before label fallback."""
    options = dict(client=client, model=model, names_text=names_text,
                   overwrite=overwrite, thinking_level=thinking_level)
    result = _identify_names(workspace, groups, **options)
    directory = workspace.cutleft_dir / 'Names'
    if not directory.is_dir():
        return result
    wide_directory = directory / 'Wide'
    pending = {}
    for stem, name in result.items():
        if name != MISSING:
            continue
        pdf = workspace.copies_dir / f'{stem}.pdf'
        try:
            source_hash = fingerprint(pdf)
            crop_path = wide_directory / f'{stem}.jpg'
            try:
                current = (crop_path.is_file() and
                           read_json(crop_path.with_suffix('.json'))['source_sha256'] == source_hash)
            except (OSError, ValueError, KeyError, TypeError):
                current = False
            if not current:
                from pdf2image import convert_from_path
                page = convert_from_path(pdf, dpi=200, first_page=1, last_page=1,
                                         use_cropbox=False)[0]
                try:
                    height = max(1, page.height // 4)
                    with staged_files(wide_directory) as staging:
                        page.crop((0, 0, page.width, height)).save(
                            staging / crop_path.name, 'JPEG', quality=95)
                        atomic_write_json(staging / f'{stem}.json', {
                            'source_sha256': source_hash,
                            'height_fraction': height / page.height,
                        })
                finally:
                    page.close()
            pending[stem] = groups[stem]
        except Exception as exc:
            print(f'Full-width name crop failed for {stem}: {type(exc).__name__}: {exc}; using label/name fallback.')
    if pending:
        result.update(_identify_names(workspace, pending, wide=True, **options))
    return result