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

from pathlib import Path

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

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": 0.25,
        })


def build_grids(directory: Path, copies_dir: Path) -> list[dict]:
    """Rebuild deterministic 6-by-4 sheets, excluding deleted or changed copies."""
    crops = []
    for path in sorted(directory.glob("*.jpg"), key=natural_key):
        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


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


def identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool,
                   thinking_level: str | None = "low") -> dict[str, str]:
    """Identify pending copies before label requests; safely fall back on any failure."""
    directory = workspace.cutleft_dir / 'Names'
    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)
    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 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'])
        prompt = f"""Identify student names from the top-left corners of exam first pages.
The image is a grid with exactly 6 columns and 4 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. Ignore class names such as MPSI 3.
Occupied cells are exactly:
{positions}
All other cells are EMPTY padding: do not return them.
Return JSON {{"cells": [{{"cell": 1, "name": "exact authorized name"}}]}}.
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: {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.name)
            for stem, (cell, crop_hash, cache_path) in pending.items():
                candidates = by_cell.get(cell['cell'], [])
                name = candidates[0] if len(candidates) == 1 and candidates[0] in roster else MISSING
                result[stem] = name
                atomic_write_json(cache_path, {'name': name, 'crop_sha256': crop_hash, 'names_text': names_text})
        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