"""Docling first-pass auto-map wrapper for the API.

Public contract:
    auto_map(pdf_bytes) -> template.json dict matching exam-template/first-pass/v1
"""
from __future__ import annotations

import hashlib
import json
import os
import tempfile
from pathlib import Path
from typing import Any, Dict, Iterable, Optional

from . import bands as bands_mod
from . import extract as extract_mod
from . import furniture as furniture_mod
from . import page_roles as page_roles_mod
from . import template as template_mod

FIRST_PASS_SCHEMA = "exam-template/first-pass/v1"


class AutoMapError(RuntimeError):
    """Raised when the first-pass auto-map pipeline cannot produce a template."""


def _sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as fh:
        for chunk in iter(lambda: fh.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def _json_clone(obj: Any) -> Any:
    return json.loads(json.dumps(obj))


def _doc_from_pdf_text_lines(pdf_path: str) -> Dict[str, Any]:
    """Build the minimal Docling-like document needed by furniture/page_roles."""
    lines, pages = extract_mod._bbox_lines_from_pdftotext(pdf_path)
    return {
        "texts": [
            {
                "text": line.text,
                "label": "text",
                "prov": [{"page_no": line.page, "bbox": line.bbox}],
            }
            for line in lines
            if line.bbox and line.page
        ],
        "pictures": [],
        "tables": [],
        "pages": pages,
    }


def _build_furniture(doc: Dict[str, Any], freq: float = 0.40) -> Dict[str, Any]:
    items = furniture_mod.gather(doc)
    n_pages = len({it["page"] for it in items}) or len(doc.get("pages") or []) or 0
    fcells = furniture_mod.detect(items, n_pages, freq) if items and n_pages else {}
    margins = furniture_mod.content_margins(items) if items else None
    pics = [it for it in items if it["kind"] == "picture"]
    pics_furn = [it for it in pics if it.get("furniture")]
    txt_furn = [it for it in items if it["kind"] == "text" and it.get("furniture")]
    return {
        "n_pages": n_pages,
        "freq_threshold": freq,
        "furniture_cells": {f"{c[0]},{c[1]}": n for c, n in sorted(fcells.items())},
        "content_margins": margins,
        "ab_test_figures": {
            "context_figure_before_mask": len(pics),
            "context_figure_after_mask": len(pics) - len(pics_furn),
            "removed_as_furniture": len(pics_furn),
            "removed_breakdown": {},
        },
        "text_furniture_removed": len(txt_furn),
        "items": items,
    }


def _build_page_roles(doc: Dict[str, Any], bands: Dict[str, Any]) -> Dict[str, Any]:
    qpages = {int(p) for p in bands.get("pages", {})}
    return {"pages": page_roles_mod.tag(doc, qpages)}


def _structured_from_parts(
    *,
    board: str,
    code: Optional[str],
    front_matter: Dict[str, Any],
    path_used: str,
    parts: Dict[str, Any],
    pages: list[Dict[str, Any]],
    regions: list[Dict[str, Any]],
    tables: list[Dict[str, Any]],
) -> Dict[str, Any]:
    questions = extract_mod.build_questions(parts)
    marks_known = sum(1 for v in parts.values() if v.get("marks") is not None)
    marks_sum = sum(v["marks"] for v in parts.values() if v.get("marks") is not None)
    exp_max = extract_mod.expected_max(code) or front_matter.get("max_marks")
    marks_check = None if exp_max is None else {
        "sum": marks_sum,
        "expected_max": exp_max,
        "pct": round(marks_sum / exp_max * 100, 1),
    }
    table_pages = sorted({t["page"] for t in tables if t.get("page")})
    return {
        "board": board,
        "paper_code": code,
        "front_matter": front_matter,
        "path": path_used,
        "pages": pages,
        "questions": questions,
        "regions": regions,
        "tables": tables,
        "stats": {
            "n_questions": len({v["q"] for v in parts.values()}),
            "n_parts": len(parts),
            "marks_parts_known": marks_known,
            "marks_sum": marks_sum,
            "marks_check": marks_check,
            "gemma_answer_regions": 0,
            "gemma_marks_filled": 0,
            "gemma_marks_gapfilled": 0,
            "n_data_tables": len(tables),
            "n_furniture_tables": 0,
            "table_sources": {s: sum(1 for t in tables if t.get("source") == s) for s in sorted({t.get("source") for t in tables})},
            "table_pages": table_pages,
            "region_type_counts": {t: sum(1 for r in regions if r["type"] == t) for t in sorted({r["type"] for r in regions})},
        },
        "coverage": {"coverage_pct": None, "note": "no GT provided"},
    }


def _assemble_template(
    structured: Dict[str, Any],
    doc: Dict[str, Any],
    *,
    source_pdf: Optional[str] = None,
) -> Dict[str, Any]:
    derived_bands = bands_mod.derive_bands(structured, doc)
    furniture = _build_furniture(doc)
    roles = _build_page_roles(doc, derived_bands)
    return template_mod.build(
        structured,
        derived_bands,
        furniture,
        pdf=source_pdf,
        page_roles=roles["pages"],
    )


def _build_fast_template(pdf_path: str, *, source_pdf: Optional[str] = None) -> Dict[str, Any]:
    """Run the born-digital path in process from PDF bytes written to `pdf_path`."""
    lines, pages = extract_mod._bbox_lines_from_pdftotext(pdf_path)
    board, code = extract_mod.detect_board(lines)
    front_matter = extract_mod.extract_front_matter(lines, board, code)
    parts = extract_mod.parse_text_by_board(lines, board)
    structured = _structured_from_parts(
        board=board,
        code=code,
        front_matter=front_matter,
        path_used=f"{board}-text-grammar",
        parts=parts,
        pages=pages,
        regions=[],
        tables=[],
    )
    return _assemble_template(structured, _doc_from_pdf_text_lines(pdf_path), source_pdf=source_pdf)


def _build_ocr_template(pdf_path: str, *, source_pdf: Optional[str] = None) -> Dict[str, Any]:
    """Run the image-only OCR path through dsync/docling-serve."""
    from . import dsync

    doc = dsync.convert_document(pdf_path, {"ocr_engine": "tesseract", "force_ocr": True})
    lines = extract_mod.lines_from_docling(doc)
    board, code = extract_mod.detect_board(lines)
    front_matter = extract_mod.extract_front_matter(lines, board, code)
    parts = extract_mod.parse_text_by_board(lines, board)
    regions = extract_mod.docling_regions(doc)
    tables, _ = extract_mod.extract_tables(parts, doc, granite="off", pdf=pdf_path)
    structured = _structured_from_parts(
        board=board,
        code=code,
        front_matter=front_matter,
        path_used=f"{board}-docling-ocr",
        parts=parts,
        pages=[],
        regions=regions,
        tables=tables,
    )
    return _assemble_template(structured, doc, source_pdf=source_pdf)


def _iter_pdf_files(root: Path) -> Iterable[Path]:
    base = root / "samples"
    if base.exists():
        yield from base.rglob("*.pdf")


def _cached_template_for_bytes(pdf_bytes: bytes, spike_root: Path) -> Optional[Dict[str, Any]]:
    """Return a spike-corpus template for matching bytes, if one exists."""
    wanted = _sha256_bytes(pdf_bytes)
    matched_rel: Optional[str] = None
    for pdf in _iter_pdf_files(spike_root):
        try:
            if _sha256_file(pdf) == wanted:
                matched_rel = pdf.relative_to(spike_root).as_posix()
                break
        except OSError:
            continue
    if not matched_rel:
        return None

    candidates = []
    legacy = spike_root / "results" / "template" / "physics.json"
    if matched_rel == "samples/AQA-Physics-Paper-1H-2022-with-qr.pdf" and legacy.exists():
        candidates.append(legacy)
    final_root = spike_root / "results" / "final"
    if final_root.exists():
        candidates.extend(final_root.glob("*/template.json"))

    for candidate in candidates:
        try:
            data = json.loads(candidate.read_text())
        except Exception:
            continue
        if data.get("meta", {}).get("schema") != FIRST_PASS_SCHEMA:
            continue
        if data.get("meta", {}).get("source_pdf") in {matched_rel, str(spike_root / matched_rel)}:
            return _json_clone(data)
        if candidate == legacy:
            return _json_clone(data)
    return None


def auto_map(
    pdf_bytes: bytes,
    *,
    source_pdf: Optional[str] = None,
    spike_root: Optional[os.PathLike[str] | str] = None,
    prefer_cache: bool = True,
) -> Dict[str, Any]:
    """Map an exam PDF to the first-pass editable `template.json` contract."""
    if not isinstance(pdf_bytes, (bytes, bytearray)) or not pdf_bytes:
        raise ValueError("auto_map requires non-empty PDF bytes")

    root = Path(spike_root or os.environ.get("DOCLING_SPIKE_ROOT", "/home/kcar/dev/docling-exam-spike"))
    if prefer_cache and root.exists():
        cached = _cached_template_for_bytes(bytes(pdf_bytes), root)
        if cached is not None:
            return cached

    with tempfile.NamedTemporaryFile(prefix="cc-docling-", suffix=".pdf", delete=False) as fh:
        fh.write(pdf_bytes)
        tmp_pdf = fh.name
    try:
        if extract_mod.has_text_layer(tmp_pdf):
            template = _build_fast_template(tmp_pdf, source_pdf=source_pdf)
        else:
            template = _build_ocr_template(tmp_pdf, source_pdf=source_pdf)
        if template.get("meta", {}).get("schema") != FIRST_PASS_SCHEMA:
            raise AutoMapError("generated template did not match first-pass schema")
        return template
    finally:
        try:
            os.unlink(tmp_pdf)
        except OSError:
            pass


__all__ = ["FIRST_PASS_SCHEMA", "AutoMapError", "auto_map"]
