"""pdf job (M6): annotated script PDF built from stored results — no AI (PDF-09).
Payload: {script_id, batch_id}. PyMuPDF draws vector marks over the original page JPEGs; text uses
HarfBuzz shaping with Noto fonts (Bengali conjuncts, Arabic RTL)."""
from __future__ import annotations
import io
from dataclasses import dataclass
from datetime import datetime, timezone

import pymupdf

from .. import jobs
from ..ids import new_id
from ..records import audit
from ..storage import get_storage
from .textrun import CSS, FONT_DIR, T, fmt_mark, spans

A4 = pymupdf.paper_rect("a4")          # 595 x 842 pt
MARGIN_W = 92                          # right-hand remark column (≈ 32 mm)
GAP = 3
FONT_PX = 8.5
MIN_PX = 7.0                           # never shrink margin text below this


# ------------------------------------------------------------------ data
def load(conn, script_id: str, tenant_id: str) -> dict:
    s = conn.execute("SELECT * FROM scripts WHERE id=%s AND tenant_id=%s", (script_id, tenant_id)).fetchone()
    if not s:
        raise jobs.FatalError("script not found")
    meta = conn.execute(
        """SELECT t.name AS school, t.settings_json AS tset, e.name AS exam, c.name_bn AS class_bn, c.name_en AS class_en,
                  sub.name_bn AS subject_bn, sub.name_en AS subject_en, st.name AS student, st.roll, st.section,
                  qs.total_target_marks
           FROM scripts sc JOIN tenants t ON t.id=sc.tenant_id JOIN batches b ON b.id=sc.batch_id
           JOIN question_sets qs ON qs.id=b.question_set_id LEFT JOIN exams e ON e.id=qs.exam_id
           JOIN classes c ON c.id=qs.class_id JOIN subjects sub ON sub.id=qs.subject_id
           LEFT JOIN students st ON st.id=sc.student_id WHERE sc.id=%s AND sc.tenant_id=%s""",
        (script_id, tenant_id)).fetchone()
    pages = conn.execute("SELECT * FROM pages WHERE script_id=%s AND tenant_id=%s ORDER BY sort",
                         (script_id, tenant_id)).fetchall()
    elements = conn.execute(
        """SELECT e.* FROM elements e JOIN pages p ON p.id=e.page_id
           WHERE p.script_id=%s AND e.tenant_id=%s ORDER BY p.sort, e.sort""", (script_id, tenant_id)).fetchall()
    rows = conn.execute(
        """SELECT g.id AS seg_id, g.element_ids_json, g.counted, g.not_attempted, g.attempt_no,
                  u.id AS unit_id, u.label AS unit_label, u.max_mark, u.key_points_json, u.sort AS u_sort,
                  q.label AS q_label, q.sort AS q_sort, q.max_mark AS q_max,
                  r.final_score, r.brief, r.error_items_json, r.positive_items_json, r.source
           FROM segments g JOIN scoring_units u ON u.id=g.scoring_unit_id JOIN questions q ON q.id=u.question_id
           LEFT JOIN score_results r ON r.segment_id=g.id AND r.is_current
           WHERE g.script_id=%s AND g.tenant_id=%s ORDER BY q.sort, u.sort, g.attempt_no""",
        (script_id, tenant_id)).fetchall()
    lang = (meta["tset"] or {}).get("feedback_lang", "bn")
    color = (meta["tset"] or {}).get("annotation_color", "#d0021b")
    return {"script": s, "meta": meta, "pages": pages, "elements": elements, "segments": rows,
            "lang": lang, "color": color}


# ------------------------------------------------------------------ drawing helpers
def _rgb(hex_color: str) -> tuple[float, float, float]:
    h = hex_color.lstrip("#")
    return tuple(int(h[i:i + 2], 16) / 255 for i in (0, 2, 4))


class Painter:
    def __init__(self, color: str):
        self.color = _rgb(color)
        self.hex = color
        self.arch = pymupdf.Archive(FONT_DIR)
        self._scratch = pymupdf.open().new_page(width=MARGIN_W + 20, height=2000)
        self._heights: dict[tuple[str, float], float] = {}

    def css(self, size: float = FONT_PX, color: str | None = None) -> str:
        return CSS + f"* {{color: {color or self.hex}; font-size: {size}px; line-height: 1.15;}} p {{margin:0;}}"

    def measure(self, html: str, width: float, size: float = FONT_PX) -> float:
        key = (html, width, size)
        if key not in self._heights:
            r = pymupdf.Rect(0, 0, width, 1000)
            spare, _ = self._scratch.insert_htmlbox(r, html, css=self.css(size), archive=self.arch)
            self._heights[key] = 1000 - spare if spare >= 0 else 1000
            self._scratch = pymupdf.open().new_page(width=MARGIN_W + 20, height=2000)
        return self._heights[key]

    def text(self, page, rect, html, size=FONT_PX, color=None, scale_low=0.5) -> None:
        page.insert_htmlbox(rect, html, css=self.css(size, color), archive=self.arch, scale_low=scale_low)


@dataclass
class MarginItem:
    y: float                 # desired top (page coordinates)
    html: str
    anchor: tuple[float, float] | None   # point on the image to draw a leader line to
    order: int = 5           # lower first when anchors tie (score badge = 0)
    boxed: bool = False
    size: float = FONT_PX


def place_margin(items: list[MarginItem], painter: Painter, x0: float, page_h: float,
                 more_label: str = "+{n}") -> list[tuple[MarginItem, pymupdf.Rect]]:
    """Stack remarks top-down without overlaps. If they cannot fit: (1) shrink text down to MIN_PX (7 pt, still readable in print),
    (2) fold the least important remarks (spelling/grammar, order 4) into one '+N more' note.
    Score badges (order 0) and explanations are never dropped."""
    width = A4.width - x0 - 4
    avail = page_h - 18 - 6
    items = sorted(items, key=lambda i: (i.y, i.order))

    def heights(its):
        return [painter.measure(i.html, width - (6 if i.boxed else 0), i.size) + (4 if i.boxed else 0) for i in its]

    def total(its):
        return sum(heights(its)) + GAP * max(len(its) - 1, 0)

    while total(items) > avail and min((i.size for i in items), default=MIN_PX) > MIN_PX:
        for i in items:
            i.size = max(MIN_PX, i.size - 0.5)
    folded = 0
    while total(items) > avail:
        idx = next((k for k in range(len(items) - 1, -1, -1) if items[k].order == 4), None)
        if idx is None:
            idx = next((k for k in range(len(items) - 1, -1, -1) if items[k].order == 3), None)
        if idx is None:
            break
        items.pop(idx)
        folded += 1
    if folded:
        note = MarginItem(y=page_h, html=f"<p><b>{spans(more_label.format(n=folded))}</b></p>", anchor=None,
                          order=9, size=min(i.size for i in items) if items else 6)
        items.append(note)
        while total(items) > avail and len(items) > 1:      # make room for the note itself
            idx = next((k for k in range(len(items) - 2, -1, -1) if items[k].order in (4, 3)), None)
            if idx is None:
                break
            items.pop(idx)
            folded += 1
            note.html = f"<p><b>{spans(more_label.format(n=folded))}</b></p>"
    hs = heights(items)
    tops, cursor = [], 6.0
    for it, h in zip(items, hs):
        top = max(min(it.y, page_h), cursor)
        tops.append(top)
        cursor = top + h + GAP
    overflow = cursor - GAP - (page_h - 18)
    if overflow > 0:          # pull items up, but never above the previous one
        for k in range(len(items) - 1, -1, -1):
            limit = (tops[k - 1] + hs[k - 1] + GAP) if k else 6.0
            shift = min(overflow, tops[k] - limit)
            if shift > 0:
                for j in range(k, len(items)):
                    tops[j] -= shift
                overflow -= shift
    return [(it, pymupdf.Rect(x0, top, x0 + width, top + max(h, 8))) for it, top, h in zip(items, tops, hs)]


# ------------------------------------------------------------------ cover page
def cover_pdf(d: dict, painter: Painter) -> bytes:
    lang, m, s = d["lang"], d["meta"], d["script"]
    t = T.get(lang, T["en"])
    rows, last_q = [], None
    for g in d["segments"]:
        if g["attempt_no"] > 1 and not g["counted"]:
            continue
        q = g["q_label"] if g["q_label"] != last_q else ""
        last_q = g["q_label"]
        if g["not_attempted"]:
            mark = f'<i>{spans(t["not_attempted"])}</i>'
        elif not g["counted"]:
            mark = f'{spans(fmt_mark(float(g["final_score"] or 0), lang))} <i>({spans(t["not_counted"])})</i>'
        else:
            mark = spans(fmt_mark(float(g["final_score"] or 0), lang))
        rows.append(f"<tr><td>{spans(q)}</td><td>{spans(g['unit_label'] or '—')}</td>"
                    f"<td>{spans(fmt_mark(float(g['max_mark']), lang))}</td><td>{mark}</td>"
                    f"<td>{spans(g['brief'] or '')}</td></tr>")

    def lst(key, values):
        if not values:
            return ""
        return f"<h3>{spans(t[key])}</h3><ul>" + "".join(f"<li>{spans(v)}</li>" for v in values) + "</ul>"
    status = (f'<p class="prov">{spans(t["provisional"])}</p>' if s["finalized_at"] is None
              else f'<p class="final">{spans(t["final"])}</p>')
    cls = m["class_bn"] if lang == "bn" else m["class_en"]
    subj = m["subject_bn"] if lang == "bn" else m["subject_en"]
    roll = m["roll"] or s["roll_from_filename"] or ""
    total = (f"{spans(fmt_mark(float(s['total_score'] or 0), lang))} / "
             f"{spans(fmt_mark(float(s['total_target'] or m['total_target_marks']), lang))}")
    html = f"""{status}
<h1>{spans(m['school'])}</h1>
<table class="meta">
<tr><td>{spans(t['exam'])}</td><td>{spans(m['exam'] or '')}</td><td>{spans(t['subject'])}</td><td>{spans(subj)}</td></tr>
<tr><td>{spans(t['class'])}</td><td>{spans(cls)}{(' — ' + spans(m['section'])) if m['section'] else ''}</td>
    <td>{spans(t['roll'])}</td><td>{spans(roll.translate(str.maketrans('0123456789', '০১২৩৪৫৬৭৮৯')) if lang == 'bn' else roll)}</td></tr>
<tr><td>{spans(t['name'])}</td><td colspan="3">{spans(m['student'] or '')}</td></tr>
</table>
<p class="total">{spans(t['total'])}: <b>{total}</b></p>
<table class="marks"><tr><th>{spans(t['question'])}</th><th>{spans(t['part'])}</th><th>{spans(t['max'])}</th>
<th>{spans(t['obtained'])}</th><th></th></tr>{''.join(rows)}</table>
{('<h3>' + spans(t['summary']) + '</h3><p>' + spans(s['summary']) + '</p>') if s['summary'] else ''}
{lst('strengths', s['strengths_json'])}{lst('weaknesses', s['weaknesses_json'])}{lst('suggestions', s['suggestions_json'])}
<p class="foot">{spans(t['date'] + ': ' + datetime.now(timezone.utc).strftime('%Y-%m-%d') + ' · AI-Examiner')}</p>"""
    css = CSS + f"""body {{font-size: 10px; color: #222; line-height: 1.3;}}
h1 {{font-size: 16px; margin: 4px 0 8px 0;}} h1 .bn {{font-family: bnb;}} h1 .ns {{font-family: nsb;}}
h3 .bn {{font-family: bnb;}} h3 .ns {{font-family: nsb;}} h3 {{font-size: 11px; margin: 10px 0 3px 0; color: {painter.hex};}}
table {{border-collapse: collapse; width: 100%;}} td, th {{border: 0.5px solid #999; padding: 3px 4px; vertical-align: top;}}
table.meta td {{border: none; padding: 1px 4px;}} th {{background: #eee; text-align: left;}}
.total {{font-size: 13px; margin: 8px 0;}} .prov {{color: #b00; font-size: 12px; border: 1px solid #b00; padding: 3px;}}
.final {{color: #070; font-size: 11px;}} .foot {{color: #888; font-size: 8px; margin-top: 12px;}} li {{margin-bottom: 2px;}}"""
    story = pymupdf.Story(html=html, user_css=css, archive=painter.arch)
    buf = io.BytesIO()
    writer = pymupdf.DocumentWriter(buf)
    where = A4 + (36, 36, -36, -36)
    more = 1
    while more:
        dev = writer.begin_page(A4)
        more, _ = story.place(where)
        story.draw(dev)
        writer.end_page()
    writer.close()
    return buf.getvalue()


# ------------------------------------------------------------------ answer pages
def build(d: dict, storage) -> bytes:
    painter = Painter(d["color"])
    lang = d["lang"]
    t = T.get(lang, T["en"])
    out = pymupdf.open(stream=cover_pdf(d, painter), filetype="pdf")
    el_by_id = {e["id"]: e for e in d["elements"]}
    page_by_id = {p["id"]: p for p in d["pages"]}
    els_per_page: dict[str, int] = {}
    for e in d["elements"]:
        els_per_page[e["page_id"]] = max(els_per_page.get(e["page_id"], 0), e["sort"] + 1)

    # Group annotations per page id.
    per_page: dict[str, dict] = {p["id"]: {"items": [], "ticks": [], "crosses": [], "underlines": [], "brackets": [],
                                           "stamps": []} for p in d["pages"]}

    def anchor_el(eid):
        e = el_by_id.get(eid)
        return (e, page_by_id[e["page_id"]]) if e else (None, None)

    stamped: set[str] = set()
    for g in d["segments"]:
        if g["not_attempted"] or g["final_score"] is None:
            continue
        ids = g["element_ids_json"] or []
        first, fpage = anchor_el(ids[0]) if ids else (None, None)
        last, lpage = anchor_el(ids[-1]) if ids else (None, None)
        if not first:
            continue
        label = f"{g['q_label']}{('(' + g['unit_label'] + ')') if g['unit_label'] else ''}"
        badge = f"<p><b>{spans(label)}: {spans(fmt_mark(float(g['final_score']), lang))} / {spans(fmt_mark(float(g['max_mark']), lang))}</b></p>"
        per_page[fpage["id"]]["items"].append(("badge", first, badge, 0))
        if not g["counted"] and g["q_label"] not in stamped:
            stamped.add(g["q_label"])                     # once per question, at its first part
            per_page[fpage["id"]]["stamps"].append((first, t["not_counted"]))
        # positives
        for pitem in g["positive_items_json"] or []:
            e, pg = anchor_el(pitem.get("element_id"))
            if e:
                per_page[pg["id"]]["ticks"].append(e)
        # errors: spelling/grammar grouped per line; others individually
        lang_errs: dict[str, list[dict]] = {}
        missing = []
        kp_text = {k["id"]: k["text"] for k in (g["key_points_json"] or [])}
        for err in g["error_items_json"] or []:
            if err["type"] == "missing_key_point":
                missing.append(kp_text.get(err.get("key_point_id"), err.get("remark", "")))
                continue
            e, pg = anchor_el(err.get("element_id"))
            if not e:
                continue
            if err["type"] in ("spelling", "grammar"):
                lang_errs.setdefault(e["id"], []).append(err)
            else:
                per_page[pg["id"]]["crosses"].append(e)
                per_page[pg["id"]]["brackets"].append(e)
                per_page[pg["id"]]["items"].append(("err", e, f"<p>{spans(err['remark'])}</p>", 3))
        for eid, errs in lang_errs.items():
            e, pg = anchor_el(eid)
            per_page[pg["id"]]["underlines"].append(e)
            parts = [f"{x['as_written']} → {x['correction']}" if x.get("as_written") and x.get("correction")
                     else x["remark"] for x in errs[:3]]
            more = f" (+{len(errs) - 3})" if len(errs) > 3 else ""
            per_page[pg["id"]]["items"].append(("lang", e, "<p>" + spans("; ".join(parts) + more) + "</p>", 4))
        tail = ""
        if missing:
            tail += f"<p><b>{spans(t['missing'])}:</b> {spans('; '.join(m[:60] for m in missing[:3]))}</p>"
        if g["brief"]:
            tail += f"<p><i>{spans(g['brief'])}</i></p>"
        if tail:
            per_page[lpage["id"]]["items"].append(("brief", last, tail, 8))

    storage_get = storage.get
    for p in d["pages"]:
        page = out.new_page(width=A4.width, height=A4.height)
        img = pymupdf.Rect(0, 0, A4.width - MARGIN_W, A4.height)
        w, h = p["width"], p["height"]
        scale = min(img.width / w, img.height / h)
        ir = pymupdf.Rect(0, 0, w * scale, h * scale)
        page.insert_image(ir, stream=storage_get(p["image_key"]), keep_proportion=True)
        mx0 = ir.x1 + 4
        page.draw_line((ir.x1 + 1, 0), (ir.x1 + 1, A4.height), color=(0.85, 0.85, 0.85), width=0.4)
        n_els = els_per_page.get(p["id"], 1)

        def rect_of(e):
            b = e["bbox_json"]
            if not b:      # no bbox (M3 fallback): estimate the line band from reading order
                y = ir.y0 + ir.height * (e["sort"] + 0.5) / max(n_els, 1)
                return pymupdf.Rect(ir.x0 + ir.width * 0.05, y - 6, ir.x0 + ir.width * 0.95, y + 6)
            return pymupdf.Rect(ir.x0 + b[0] * ir.width, ir.y0 + b[1] * ir.height,
                                ir.x0 + b[2] * ir.width, ir.y0 + b[3] * ir.height)

        ann = per_page[p["id"]]
        col = painter.color
        if p["kind"] == "rough":
            painter.text(page, pymupdf.Rect(8, 6, 200, 24), f"<p>{spans(t['rough'])}</p>", size=11, color="#888")
        for e in ann["underlines"]:
            r = rect_of(e)
            y = r.y1 + 1
            pts, x, up = [], r.x0, True
            while x < r.x1:                      # wavy underline
                pts.append(pymupdf.Point(x, y + (0 if up else 2)))
                x += 4; up = not up
            if len(pts) > 1:
                page.draw_polyline(pts, color=col, width=0.8)
        for e in ann["brackets"]:
            r = rect_of(e)
            page.draw_polyline([(r.x0 - 2, r.y0), (r.x0 - 5, r.y0), (r.x0 - 5, r.y1), (r.x0 - 2, r.y1)], color=col, width=1)
        for e in ann["ticks"]:
            r = rect_of(e)
            painter.text(page, pymupdf.Rect(min(r.x1 + 2, ir.x1 - 14), r.y0 - 3, ir.x1 + 2, r.y1 + 6),
                         "<p><b>✓</b></p>", size=13)
        for e in ann["crosses"]:
            r = rect_of(e)
            painter.text(page, pymupdf.Rect(min(r.x1 + 2, ir.x1 - 14), r.y0 - 3, ir.x1 + 2, r.y1 + 6),
                         "<p><b>✗</b></p>", size=13)
        for e, label in ann["stamps"]:
            r = rect_of(e)
            painter.text(page, pymupdf.Rect(r.x0 + ir.width * 0.4, r.y0 - 14, ir.x1 - 4, r.y0), f"<p>{spans(label)}</p>",
                         size=10, color="#777")

        items = []
        for kind, e, html, order in ann["items"]:
            r = rect_of(e)
            items.append(MarginItem(y=r.y0 if kind != "brief" else r.y1, html=html, order=order, boxed=(kind == "badge"),
                                    anchor=(r.x1, (r.y0 + r.y1) / 2) if kind in ("err", "lang") else None,
                                    size=FONT_PX + (1.5 if kind == "badge" else 0)))
        for it, rect in place_margin(items, painter, mx0, A4.height, t["more_errors"]):
            if it.boxed:
                page.draw_rect(rect, color=col, width=0.9)
                rect = rect + (3, 2, -3, -2)
            painter.text(page, rect, it.html, size=it.size)
            if it.anchor:
                ay = it.anchor[1]
                page.draw_line(pymupdf.Point(min(it.anchor[0] + 1, ir.x1), ay), pymupdf.Point(rect.x0 - 1, rect.y0 + 5),
                               color=col, width=0.4, dashes="[2 1] 0")
        painter.text(page, pymupdf.Rect(mx0, A4.height - 14, A4.width - 4, A4.height - 2),
                     f"<p>{spans(t['page'])} {spans(fmt_mark(p['sort'], lang))}</p>", size=7, color="#999")
    out.subset_fonts()           # each text box embeds the full Noto fonts; subsetting cuts ~9 MB to ~0.3 MB
    return out.tobytes(garbage=4, deflate=True)


# ------------------------------------------------------------------ job handler
def handle(conn, job) -> dict:
    p = job["payload_json"]
    tenant_id, script_id = job["tenant_id"], p["script_id"]
    d = load(conn, script_id, tenant_id)
    s = d["script"]
    data = build(d, get_storage())
    version = s["pdf_version"] + 1
    key = s["original_file_key"].rsplit("/", 1)[0] + f"/pdf/v{version}.pdf"
    get_storage().put(key, data, "application/pdf")
    provisional = s["finalized_at"] is None
    conn.execute("""INSERT INTO script_pdfs (id, tenant_id, script_id, version, file_key, provisional, size_bytes)
                    VALUES (%s,%s,%s,%s,%s,%s,%s)""", (new_id(), tenant_id, script_id, version, key, provisional, len(data)))
    conn.execute("""UPDATE scripts SET pdf_version=%s, status=CASE WHEN finalized_at IS NULL THEN 'pdf_generated' ELSE status END,
                    updated_at=now() WHERE id=%s AND tenant_id=%s""", (version, script_id, tenant_id))
    audit(conn, tenant_id=tenant_id, action="pdf", entity_type="script", entity_id=script_id,
          after={"version": version, "bytes": len(data), "provisional": provisional})
    return {"version": version, "bytes": len(data)}