import re

import pymupdf

from aix import jobs
from aix.main import _handlers, run_one

from .conftest import add_script
from .helpers import make_bn_answer_pdf
from .test_pipeline_e2e import drain

PAGE1 = ["১।", "বাংলাদেশের রাজধানী ঢাকা। ঢাকা একটি পুরোনো শহর।", "এটি বুড়িগঙ্গা নদীর তীরে অবস্থিত।",
         "২", "ক) কবির নাম জসিমউদ্দিন।", "খ) উদ্দীপকে গ্রামের সমাজিক ছবি ফুটে উঠেছে।",
         "গ্রামের মানুষ কৃষিকাজ করে এবং তারা খুব পরিশ্রমি।", "নকশী কাঁথার মাঠ কাব্যে এই চিত্র দেখা যায়।"]
PAGE2 = ["৩", "ক) পদ্মা নদী।", "খ) x = 2 + 3 = 6", "অতএব উত্তর ৬ সেমি।"]


def extraction(lines, boxes):
    els = []
    for text, b in zip(lines, boxes):
        m = re.match(r"^([০-৯]+।?|[কখগঘ]\))\s*(.*)$", text)
        if m and not m.group(2):
            els.append({"type": "heading", "text": m.group(1), "confidence": 0.98, "bbox": b})
        elif m:
            els.append({"type": "line", "text": text, "confidence": 0.95, "bbox": b, "lang": ["bn"]})
        else:
            els.append({"type": "line", "text": text, "confidence": 0.95, "bbox": b, "lang": ["bn"]})
    return {"kind": "answer", "legibility": 7, "elements": els}


def scoring(world):
    rev = {v: k for k, v in world.items()}

    def refs(text, uid):
        block = re.search(rf"=== Unit {uid}.*?\n(.*?)(?:\n===|\Z)", text, re.S).group(1)
        return re.findall(r"\[(p\d+-l\d+)\]", block)

    def build(call):
        text = call["parts"][0].text
        units = []
        for uid in dict.fromkeys(re.findall(r"=== Unit (\S+)", text)):
            r = refs(text, uid)
            name = rev[uid]
            u = {"unit_id": uid, "correctness": 0.9, "completeness": 0.8, "positive_items": [], "error_items": [],
                 "key_points": [], "brief": "মূল বিষয় ঠিক আছে, তবে ব্যাখ্যা সংক্ষিপ্ত।", "detailed": "বিস্তারিত"}
            if name == "u1":
                u["positive_items"] = [{"line_ref": r[1], "key_point_id": "KP1"}]
                u["brief"] = "সঠিক উত্তর, অতিরিক্ত তথ্যও দিয়েছে।"
            if name == "u2a":
                u["error_items"] = [{"type": "spelling", "line_ref": r[0], "as_written": "জসিমউদ্দিন",
                                     "correction": "জসীমউদ্‌দীন", "remark": "বানান ভুল"}]
            if name == "u2b":
                u["error_items"] = [
                    {"type": "spelling", "line_ref": r[0], "as_written": "সমাজিক", "correction": "সামাজিক", "remark": "বানান"},
                    {"type": "spelling", "line_ref": r[1], "as_written": "পরিশ্রমি", "correction": "পরিশ্রমী", "remark": "বানান"},
                    {"type": "missing_key_point", "key_point_id": "KP1", "remark": "মূল বিষয় নেই"}]
                u["correctness"], u["completeness"] = 0.6, 0.5
                u["positive_items"] = [{"line_ref": r[2]}]
            if name == "u3b":
                u["error_items"] = [{"type": "calculation", "line_ref": r[0], "as_written": "2 + 3 = 6",
                                     "correction": "5", "remark": "যোগফল ভুল: ২+৩=৫"}]
                u["correctness"] = 0.3
            units.append(u)
        return {"units": units, "strengths": ["প্রশ্ন বুঝে উত্তর দিয়েছে", "উদাহরণ ব্যবহার করেছে", "হাতের লেখা পরিষ্কার"],
                "weaknesses": ["বানান ভুল", "হিসাবে অসাবধানতা", "সৃজনশীল খ-এ বিশ্লেষণ কম"],
                "suggestions": ["বানান অনুশীলন করো", "হিসাব শেষে যাচাই করো", "উদ্দীপকের সাথে মিল দেখাও"],
                "summary": "ভালো প্রস্তুতি, তবে বানান ও হিসাবে মনোযোগ দরকার।"}
    return build


def test_annotated_pdf(conn, storage, world, fake, tmp_path):
    pdf_bytes, boxes = make_bn_answer_pdf([PAGE1, PAGE2])
    sid = add_script(conn, storage, world, pdf_bytes)
    handlers = {k: v for k, v in _handlers().items() if k != "noop"}
    jobs.enqueue(conn, "rasterize", {"script_id": sid, "batch_id": world["batch"]}, tenant_id=world["tenant"])
    assert run_one(conn, handlers)
    conn.execute("""UPDATE jobs j SET priority = 1000 - p.sort FROM pages p
                    WHERE j.type='extract_page' AND p.id = j.payload_json->>'page_id'""")
    fake.push("ExtractionPage", extraction(PAGE1, boxes[0]))
    fake.push("ExtractionPage", extraction(PAGE2, boxes[1]))
    fake.push("ScriptFinding", scoring(world))
    drain(conn, handlers)
    failed = conn.execute("SELECT type, last_error FROM jobs WHERE status='failed'").fetchall()
    assert not failed, failed

    s = conn.execute("SELECT status, pdf_version FROM scripts WHERE id=%s", (sid,)).fetchone()
    assert s["status"] == "pdf_generated" and s["pdf_version"] == 1
    row = conn.execute("SELECT file_key, size_bytes, provisional FROM script_pdfs WHERE script_id=%s", (sid,)).fetchone()
    assert row["provisional"]
    data = storage.get(row["file_key"])
    doc = pymupdf.open(stream=data, filetype="pdf")
    assert doc.page_count == 3                         # cover + 2 answer pages
    # Shaped Bengali doesn't round-trip through the PDF text layer, so check Latin/ASCII content only.
    assert "Test School" in doc[0].get_text() and "AI-Examiner" in doc[0].get_text()
    assert len(data) < 2_000_000
    for i in range(3):
        doc[i].get_pixmap(dpi=110).save(f"/tmp/pdf_page{i}.png")
    (tmp_path / "script.pdf").write_bytes(data)
    open("/tmp/sample_annotated_script.pdf", "wb").write(data)


def test_margin_overflow_stays_on_page():
    from aix.pipeline.pdf import A4, MARGIN_W, MarginItem, Painter, place_margin
    painter = Painter("#d0021b")
    items = [MarginItem(y=700 + i, html=f"<p>বানান ভুল নম্বর {i}: সমাজিক → সামাজিক</p>", anchor=(400, 700 + i), order=4)
             for i in range(40)]
    items.insert(0, MarginItem(y=50, html="<p><b>১: ২ / ২</b></p>", anchor=None, order=0, boxed=True))
    placed = place_margin(items, painter, A4.width - MARGIN_W + 4, A4.height, "আরও {n}টি ভুল")
    tops = [r.y0 for _, r in placed]
    assert all(r.y1 <= A4.height - 2 for _, r in placed)
    assert tops == sorted(tops)                         # still in reading order, no overlap reordering
    assert all(placed[i][1].y0 >= placed[i - 1][1].y1 for i in range(1, len(placed)))   # no overlaps
    assert placed[0][0].order == 0                                    # score badge kept
    assert "আরও" in placed[-1][0].html                                # folded remarks summarized