"""extract_page job (M3): one AI pass per page -> elements, kind, legibility, cover fields, flags.
Payload: {page_id, script_id, batch_id}."""
from __future__ import annotations
import pathlib

from .. import ai, config, jobs
from ..db import Jsonb
from ..ids import new_id
from ..providers.base import ImagePart, TextPart
from ..records import audit, raise_flag, record_usage
from ..schemas.extraction import ExtractionPage
from ..storage import get_storage
from .labels import normalize_digits

SYSTEM = (pathlib.Path(__file__).parent.parent / "prompts" / "extract_system.md").read_text(encoding="utf-8")
PROMPT_VERSION = "extract-v1"


def _known_labels(conn, script_id: str, tenant_id: str) -> list[str]:
    rows = conn.execute(
        """SELECT q.label AS q, u.label AS u FROM scripts s
           JOIN batches b ON b.id=s.batch_id JOIN questions q ON q.question_set_id=b.question_set_id
           JOIN scoring_units u ON u.question_id=q.id
           WHERE s.id=%s AND s.tenant_id=%s ORDER BY q.sort, u.sort""", (script_id, tenant_id)).fetchall()
    return [f"{r['q']}({r['u']})" if r["u"] else r["q"] for r in rows]


def _prev_tail(conn, page, tenant_id: str) -> str:
    rows = conn.execute(
        """SELECT e.text FROM elements e JOIN pages p ON p.id=e.page_id
           WHERE p.script_id=%s AND p.sort=%s AND e.tenant_id=%s AND e.text IS NOT NULL
           ORDER BY e.sort DESC LIMIT 3""", (page["script_id"], page["sort"] - 1, tenant_id)).fetchall()
    return "\n".join(r["text"] for r in reversed(rows)) or "(not available)"


def handle(conn, job) -> dict:
    p = job["payload_json"]
    tenant_id, page_id = job["tenant_id"], p["page_id"]
    page = conn.execute("SELECT * FROM pages WHERE id=%s AND tenant_id=%s", (page_id, tenant_id)).fetchone()
    if not page:
        raise jobs.FatalError("page not found")
    if page["status"] == "extracted":          # idempotent re-run
        _maybe_enqueue_assemble(conn, job, page["script_id"])
        return {"skipped": "already extracted"}

    labels = _known_labels(conn, page["script_id"], tenant_id)
    user = TextPart(f"Known question/part labels in this paper: {', '.join(labels) or '(none)'}\n"
                    f"Previous page ended with:\n{_prev_tail(conn, page, tenant_id)}\n"
                    f"Transcribe the attached page.")
    image = ImagePart(get_storage().get(page["ai_image_key"] or page["image_key"]))
    res = ai.call(conn, stage="extract", tenant_id=tenant_id, system=SYSTEM, parts=[user, image],
                  schema=ExtractionPage, max_tokens=16000)
    out: ExtractionPage = res.parsed
    threshold = float(config.get(conn, "extract.confidence_threshold", tenant_id))

    kind = out.kind
    if kind == "blank" and out.elements:
        kind = "answer"                         # never silently blank a page that has content
    if kind == "blank":
        raise_flag(conn, tenant_id=tenant_id, batch_id=p["batch_id"], type_="blank_disagreement",
                   script_id=page["script_id"], page_id=page_id,
                   payload={"ink_ratio": float(page["ink_ratio"] or 0)})
    if kind == "unreadable":
        raise_flag(conn, tenant_id=tenant_id, batch_id=p["batch_id"], type_="unreadable_page",
                   script_id=page["script_id"], page_id=page_id, payload={"notes": out.notes})

    conn.execute("DELETE FROM elements WHERE page_id=%s AND tenant_id=%s", (page_id, tenant_id))
    low = []
    for i, el in enumerate(out.elements):
        flags = list(el.flags)
        if kind == "rough" and "rough" not in flags:
            flags.append("rough")
        eid = new_id()
        conn.execute(
            """INSERT INTO elements (id, tenant_id, page_id, sort, type, bbox_json, text, latex, description,
                                     table_json, lang_json, confidence, flags_json, candidates_json, q_label)
               VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
            (eid, tenant_id, page_id, i, el.type, Jsonb(el.bbox) if el.bbox else None, el.text, el.latex,
             el.description, Jsonb(el.table) if el.table else None, Jsonb(el.lang), round(el.confidence, 3),
             Jsonb(flags), Jsonb(el.candidates) if el.candidates else None,
             el.text if el.type == "heading" else None))
        if el.confidence < threshold or "uncertain" in flags or "illegible" in flags:
            low.append(eid)
    if low:
        raise_flag(conn, tenant_id=tenant_id, batch_id=p["batch_id"], type_="low_confidence_extraction",
                   script_id=page["script_id"], page_id=page_id, payload={"element_ids": low[:50], "count": len(low)})

    if kind == "cover" and out.cover and out.cover.roll:
        roll = "".join(ch for ch in normalize_digits(out.cover.roll) if ch.isdigit()).lstrip("0")
        script = conn.execute("SELECT roll_from_filename FROM scripts WHERE id=%s AND tenant_id=%s",
                              (page["script_id"], tenant_id)).fetchone()
        conn.execute("UPDATE scripts SET roll_from_cover=%s WHERE id=%s AND tenant_id=%s",
                     (roll, page["script_id"], tenant_id))
        if script["roll_from_filename"] and roll and roll != script["roll_from_filename"].lstrip("0"):
            raise_flag(conn, tenant_id=tenant_id, batch_id=p["batch_id"], type_="cover_roll_mismatch",
                       script_id=page["script_id"], page_id=page_id,
                       payload={"filename": script["roll_from_filename"], "cover": roll})

    conf = min((e.confidence for e in out.elements), default=1.0)
    conn.execute(
        """UPDATE pages SET kind=%s, legibility=%s, confidence=%s, status='extracted', extracted_at=now(),
                           model_info_json=%s, tokens_in=%s, tokens_out=%s, cost_micro=%s, notes=%s
           WHERE id=%s AND tenant_id=%s""",
        (kind, out.legibility, round(conf, 3), Jsonb({**res.model_info(), "prompt": PROMPT_VERSION, **res.extra}),
         res.tokens_in, res.tokens_out, res.cost_micro, out.notes, page_id, tenant_id))
    record_usage(conn, tenant_id=tenant_id, op_type="extract_page",
                 credits_milli=int(config.get(conn, "credits.extract_page", tenant_id)), result=res,
                 batch_id=p["batch_id"], script_id=page["script_id"], page_id=page_id)
    audit(conn, tenant_id=tenant_id, action="extract_page", entity_type="page", entity_id=page_id,
          after={"kind": kind, "elements": len(out.elements)}, model_info=res.model_info(), tokens=res.tokens())
    _maybe_enqueue_assemble(conn, job, page["script_id"])
    return {"elements": len(out.elements), "kind": kind}


def _maybe_enqueue_assemble(conn, job, script_id: str) -> None:
    tenant_id = job["tenant_id"]
    # Row lock serializes concurrent "last page" checks for the same script.
    conn.execute("SELECT id FROM scripts WHERE id=%s AND tenant_id=%s FOR UPDATE", (script_id, tenant_id))
    left = conn.execute("""SELECT count(*) AS n FROM pages WHERE script_id=%s AND tenant_id=%s
                           AND status NOT IN ('extracted','skipped')""", (script_id, tenant_id)).fetchone()["n"]
    if left == 0:
        conn.execute("UPDATE scripts SET status='extracted', updated_at=now() WHERE id=%s AND tenant_id=%s",
                     (script_id, tenant_id))
        jobs.enqueue(conn, "assemble", {"script_id": script_id, "batch_id": job["payload_json"]["batch_id"]},
                     tenant_id=tenant_id, priority=job["priority"], dedupe_key=f"assemble:{script_id}",
                     correlation_id=job["correlation_id"])