#!/usr/bin/env python3
"""
Extract AZ-104 questions from CertyIQ PDF into data/questions.json + assets.
"""
from __future__ import annotations

import json
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

import fitz
from pypdf import PdfReader

ROOT = Path(__file__).resolve().parents[1]
SCRIPTS = Path(__file__).resolve().parent
if str(SCRIPTS) not in sys.path:
    sys.path.insert(0, str(SCRIPTS))
PDF_PATH = ROOT / "AZ-104 (634Qs).pdf"
OUT_JSON = ROOT / "data" / "questions.json"
OUT_MANIFEST = ROOT / "data" / "manifest.json"
ASSETS_DIR = ROOT / "assets" / "questions"

EXAM = "AZ-104"
CERTYIQ_BASE = "https://certyiq.com/papers/microsoft/az-104"

QUESTION_OPENER = re.compile(
    r"(?:^|\n)("
    r"Your company |You have |You need |You want |You plan |You are |You must |You configure |You create |"
    r"You download |You deploy |You manage |You sign |You store |You configure |"
    r"Does the solution meet the goal\?|"
    r"DRAG DROP\s*-\s*|HOTSPOT\s*-\s*|"
    r"Which |What should |What is |What are |What can |What would |What must |What do |"
    r"How |Each |An administrator|A company|Contoso|Given |After you|There is |For each|If you"
    r")",
    re.I | re.M,
)

QUESTION_MARKER = re.compile(r"Question:\s*(\d+)", re.I)

SCENARIO_OPENER = re.compile(
    r"^(Your company |You have |You need |You want |You plan |You are |You must |Contoso|Given )",
    re.I,
)
def find_body_start(lookback: str) -> int:
    openers = list(QUESTION_OPENER.finditer(lookback))
    if not openers:
        return max(0, len(lookback) - 2500)
    for o in openers:
        if SCENARIO_OPENER.match(o.group(1)):
            return o.start(1)
    return openers[-1].start(1)


def detect_type(body: str) -> str:
    b = body.replace("\t", " ")
    bl = b.lower()
    if re.search(r"drag\s*drop\s*-", b, re.I):
        return "drag_drop"
    if "select and place" in bl or "dragging the correct option" in bl:
        return "drag_drop"
    if "use the drop-down menus" in bl:
        return "dropdown"
    if "for each requirement" in bl or "which command should you run" in bl:
        return "dropdown"
    if re.search(r"hotspot\s*-", b, re.I) or (
        "hot area:" in bl and "drop-down menus" not in bl[:400]
    ):
        return "hotspot"
    if "does the solution meet the goal" in bl:
        return "yes_no"
    if re.search(r"which (two|three|four|five)", bl) and (
        "each correct answer" in bl or "each correct selection" in bl
    ):
        return "multi_select"
    if re.search(r"each correct selection is worth one point", bl):
        if re.search(r"which (two|three|four|five)", bl):
            return "multi_select"
        return "hotspot"
    if re.search(r"^[A-G]\.\s", b, re.M):
        return "single_choice"
    return "unknown"


def bilingual(text: str | None) -> dict:
    t = (text or "").strip()
    return {"en": t if t else None, "vi": None}


def _strip_q_markers(text: str) -> str:
    return re.sub(r"(?:<<<Q:\d+>>>\s*)+", " ", text)


def parse_options(body: str) -> list[dict]:
    body = _strip_q_markers(body)
    answer_pos = re.search(r"Answer:\s*", body, re.I)
    region = body[: answer_pos.start()] if answer_pos else body
    opts = []
    seen: set[str] = set()
    for m in re.finditer(
        r"(?:^|\n|\?\s+)\s*([A-G])\.\s*\n?\s*(.+?)(?=(?:^|\n|\?\s+)\s*[A-G]\.\s|\bAnswer:|$)",
        region,
        re.M | re.S,
    ):
        key = m.group(1)
        if key in seen:
            continue
        text = re.sub(r"\s+", " ", m.group(2).strip())
        if text and len(text) > 2:
            seen.add(key)
            opts.append({"key": key, "text": bilingual(text)})
    if not opts and re.search(r"^A\.\s*Yes", body, re.M):
        opts = [
            {"key": "A", "text": bilingual("Yes")},
            {"key": "B", "text": bilingual("No")},
        ]
    return opts


def parse_boxes(explanation: str) -> list[dict]:
    boxes = []
    for m in re.finditer(
        r"Box\s*(\d+)\s*:\s*(.+?)(?=Box\s*\d+\s*:|Reference:|$)",
        explanation,
        re.I | re.S,
    ):
        ans = re.sub(r"\s+", " ", m.group(2).strip()).rstrip(" -")
        if len(ans) > 250:
            ans = ans[:250].rsplit(" ", 1)[0] + "..."
        boxes.append(
            {
                "id": int(m.group(1)),
                "label": f"Box {m.group(1)}",
                "answer": bilingual(ans),
            }
        )
    return boxes


def parse_numbered_slots(explanation: str) -> list[dict]:
    """Drag-drop / ordered answers from '1. item' list in explanation."""
    slots = []
    for m in re.finditer(
        r"^(\d+)\.\s*\n?\s*(.+?)(?=^\d+\.\s*\n|Therefore,|Reference:|Box\s*\d|$)",
        explanation,
        re.M | re.S,
    ):
        ans = re.sub(r"\s+", " ", m.group(2).strip())
        ans = re.split(r":\s*", ans, maxsplit=1)[0] if len(ans) > 120 else ans
        if len(ans) > 200:
            ans = ans[:200].rsplit(" ", 1)[0] + "..."
        if ans and len(ans) > 3:
            slots.append(
                {
                    "id": int(m.group(1)),
                    "label": f"Slot {m.group(1)}",
                    "answer": bilingual(ans),
                }
            )
    return slots


def parse_answer_keys(raw: str, qtype: str) -> list[str]:
    raw = re.sub(r"\s+", " ", raw.strip())
    if not raw:
        return []
    if qtype == "yes_no":
        if raw.upper().startswith("YES") or raw == "A":
            return ["Yes"]
        if raw.upper().startswith("NO") or raw == "B":
            return ["No"]
    if re.fullmatch(r"[A-G](,\s*[A-G])+", raw, re.I):
        return [x.strip().upper() for x in raw.split(",")]
    if re.fullmatch(r"[A-G]+", raw, re.I):
        return list(raw.upper())
    if raw in ("Yes", "No"):
        return [raw]
    if len(raw) == 1 and raw.isalpha():
        return [raw.upper()]
    return []


def extract_yes_no_parts(body: str) -> tuple[str | None, str, str]:
    """Return (scenario, solution_line, question_line) for yes/no items."""
    b = body.replace("\t", " ")
    b = re.sub(r"Certy\s*IQ", "", b, flags=re.I)
    answer_pos = re.search(r"Answer:\s*", b, re.I)
    region = b[: answer_pos.start()] if answer_pos else b
    sol_m = re.search(r"(Solution:\s*.+?)(?=Does the solution meet the goal)", region, re.I | re.S)
    does_m = re.search(r"(Does the solution meet the goal\?)", region, re.I)
    if not does_m:
        return None, "", region.strip()
    scenario = region[: sol_m.start()].strip() if sol_m else region[: does_m.start()].strip()
    solution = sol_m.group(1).strip() if sol_m else ""
    question = region[does_m.start() :].strip()
    return scenario or None, solution, question


def extract_stem(body: str, qtype: str) -> str:
    b = body.replace("\t", " ")
    b = re.sub(r"Certy\s*IQ", "", b, flags=re.I)
    b = re.sub(r"<<<QUESTION:\d+>>>", " ", b, flags=re.I)
    b = re.sub(r"Question:\s*\d+\s*", " ", b, flags=re.I)
    b = re.sub(r"(?:<<<Q:\d+>>>\s*)+", " ", b)
    b = re.sub(r"^(HOTSPOT|DRAG DROP)\s*-\s*", "", b, flags=re.I | re.M)

    if qtype == "yes_no":
        scenario, solution, question = extract_yes_no_parts(body)
        parts = [p for p in (scenario, solution, question) if p]
        return re.sub(r"\s+", " ", " ".join(parts)).strip()

    answer_pos = re.search(r"Answer:\s*", b, re.I)
    search_region = b[: answer_pos.start()] if answer_pos else b
    opt = re.search(r"(?:^|\n|\?\s+)\s*([A-G])\.\s*\S", search_region, re.M)
    if opt:
        b = search_region[: opt.start()]
    else:
        b = search_region
    b = re.sub(r"\s+", " ", b).strip()
    b = re.sub(r"^[a-z]\.\s+", "", b)
    return b


def infer_question_id(body: str, chunk_end: int, full_text: str) -> int | None:
    window = full_text[chunk_end : chunk_end + 400]
    markers = list(QUESTION_MARKER.finditer(window))
    if markers:
        return int(markers[0].group(1))
    markers = list(QUESTION_MARKER.finditer(body[-200:]))
    if markers:
        return int(markers[-1].group(1))
    return None


QUESTION_MARKER_TOKEN = re.compile(
    r"(?:Question:\s*(\d+)\s*(?:Certy\s*IQ\s*)*|<<<QUESTION:(\d+)>>>)",
    re.I,
)


def split_multi_question_body(body: str) -> list[str]:
    """Split one text region that contains multiple Answer/Explanation blocks."""
    body = body.strip()
    if not body:
        return []

    parts = re.split(
        r"\.\s*\n(?=(?:You have |You need |You plan |You want |You are |Which |What should |What is |What are |"
        r"HOTSPOT|DRAG DROP|For each|Does the solution|An administrator|Contoso|Given ))",
        body,
        flags=re.I,
    )
    blocks: list[str] = []
    pending = ""
    for part in parts:
        piece = part.strip()
        if not piece:
            continue
        if pending and not pending.endswith("."):
            pending = f"{pending}.{piece}"
        elif pending:
            pending = f"{pending} {piece}"
        else:
            pending = piece
        if re.search(r"\bAnswer:\s*", pending, re.I) and re.search(
            r"Explanation:\s*", pending, re.I
        ):
            blocks.append(pending)
            pending = ""
    if pending and re.search(r"\bAnswer:\s*", pending, re.I):
        blocks.append(pending)

    if len(blocks) <= 1:
        answer_marks = list(re.finditer(r"\bAnswer:\s*", body, re.I))
        if len(answer_marks) <= 1:
            return [body]

        blocks = []
        prev_end = 0
        for i, am in enumerate(answer_marks):
            region = body[prev_end : am.start()]
            openers = list(QUESTION_OPENER.finditer(region))
            start = prev_end + (openers[-1].start(1) if openers else 0)
            end = answer_marks[i + 1].start() if i + 1 < len(answer_marks) else len(body)
            chunk = body[start:end].strip()
            if len(chunk) >= 50 and re.search(r"Explanation:\s*", chunk, re.I):
                blocks.append(chunk)
            prev_end = end

    return blocks


def peel_next_question(after: str) -> tuple[str, str]:
    """Split `after` into (continuation of current question, text for next question(s))."""
    after = after.strip()
    if not after:
        return "", ""

    if re.match(
        r"(You have |You plan |You need |You want |You are |HOTSPOT|DRAG DROP|For each of the following|Contoso|Given )",
        after,
        re.I,
    ):
        return "", after

    if not re.search(r"\bAnswer:\s*", after, re.I):
        return "", after

    exp = re.search(r"\bExplanation:\s*", after, re.I)
    if not exp:
        return after, ""

    rest = after[exp.end() :]
    split_re = re.compile(
        r"\.\s*\n(?=(?:You have |You plan |You need |You want |You are |"
        r"HOTSPOT|DRAG DROP|For each of the following|Contoso|Given ))",
        re.I,
    )
    m = split_re.search(rest)
    if m:
        cut = exp.end() + m.start()
        return after[:cut].strip(), after[cut:].strip()
    return after, ""


def split_question_chunks(text: str) -> list[tuple[int | None, str]]:
    """
    CertyIQ PDFs place `Question: N` between stem and options when a question spans pages.
    Each question is: text since the last split + continuation after its marker until the next
    question starts.
    """
    text = normalize_text(text)
    raw = list(QUESTION_MARKER_TOKEN.finditer(text))
    if not raw:
        return []

    groups: list[tuple[list[int], int, int]] = []
    i = 0
    while i < len(raw):
        qid = int(raw[i].group(1) or raw[i].group(2))
        start = raw[i].start()
        end = raw[i].end()
        qids = [qid]
        j = i + 1
        while j < len(raw) and raw[j].start() - raw[j - 1].start() < 200:
            qids.append(int(raw[j].group(1) or raw[j].group(2)))
            end = raw[j].end()
            j += 1
        groups.append((qids, start, end))
        i = j

    chunks: list[tuple[int | None, str]] = []
    cursor = 0

    for gi, (qids, mstart, mend) in enumerate(groups):
        next_mstart = groups[gi + 1][1] if gi + 1 < len(groups) else len(text)
        prefix = text[cursor:mstart].strip()
        after = text[mend:next_mstart].strip()

        if is_note_marker_group(qids, prefix, text, mstart):
            cursor = mend
            continue

        if len(qids) > 1:
            subs = split_multi_question_body(prefix)
            if len(subs) == len(qids):
                for qid, body in zip(qids, subs):
                    chunks.append((qid, body))
            else:
                first_id = max(1, qids[-1] - len(subs) + 1)
                for k, body in enumerate(subs):
                    chunks.append((first_id + k, body))
            cursor = mend
        else:
            qid = qids[0]
            cont, remainder = peel_next_question(after)
            body = f"{prefix}\n{cont}".strip()
            if body:
                chunks.append((qid, body))
            cursor = mend + len(after) - len(remainder)

    return chunks


def is_note_marker_group(qids: list[int], prefix: str, text: str, mstart: int) -> bool:
    """Skip metadata clusters like 'Question:1 Question:2 Question:3 distinctive result'."""
    if len(qids) < 2:
        return False
    window = text[max(0, mstart - 350) : mstart + 120].lower()
    return "identical set-up" in window or "distinctive result" in window


def build_question_body_map(text: str) -> dict[int, str]:
    """Map question id -> best extracted PDF body (marker-based split)."""
    bodies: dict[int, str] = {}
    for qid, body in split_question_chunks(text):
        if not qid or len(body) < 40:
            continue
        prev = bodies.get(qid)
        if not prev:
            bodies[qid] = body
            continue
        prev_score = _body_quality(prev)
        new_score = _body_quality(body)
        if new_score > prev_score:
            bodies[qid] = body
    return bodies


def _body_quality(body: str) -> int:
    score = len(body)
    if re.search(r"\bAnswer:\s*", body, re.I):
        score += 500
    if re.search(r"\bExplanation:\s*", body, re.I):
        score += 300
    if re.search(r"^[A-G]\.\s", body, re.M):
        score += 200
    return score


def normalize_text(text: str) -> str:
    text = text.replace("\t", " ")
    text = re.sub(
        r"Question:\s*(\d+)\s*(?:Certy\s*IQ\s*)+",
        r"\n<<<QUESTION:\1>>>\n",
        text,
        flags=re.I,
    )
    text = re.sub(r"Certy\s*IQ", "", text, flags=re.I)
    start = re.search(
        r"Your company has serval departments|Your company has several departments",
        text,
        re.I,
    )
    if start:
        text = text[start.start() :]
    return text


def parse_block(num: int, body: str) -> dict:
    qtype = detect_type(body)
    ans_m = re.search(r"Answer:\s*(.*?)\s*Explanation:", body, re.DOTALL | re.I)
    exp_m = re.search(
        r"Explanation:\s*(.*?)(?=Reference:|Note:|\Z)", body, re.DOTALL | re.I
    )
    ref_m = re.search(r"Reference:\s*(.*?)(?=Note:|Question:|\Z)", body, re.DOTALL | re.I)

    answer_raw = ans_m.group(1).strip() if ans_m else ""
    explanation = exp_m.group(1).strip() if exp_m else ""
    references = []
    if ref_m:
        refs = ref_m.group(1).strip()
        references = [u.strip() for u in re.findall(r"https?://[^\s\)\]]+", refs)]

    options = parse_options(body)
    keys = parse_answer_keys(answer_raw, qtype)
    boxes = parse_boxes(explanation)

    if qtype == "drag_drop" and not boxes:
        boxes = parse_numbered_slots(explanation)

    if not keys and boxes:
        if qtype in ("hotspot", "dropdown", "drag_drop"):
            keys = [
                (b["answer"]["en"] or "").strip()
                for b in boxes
                if b["answer"]["en"]
            ]
        elif qtype == "drag_drop":
            keys = [str(b["id"]) for b in boxes]

    if not keys and qtype == "drag_drop":
        slots = parse_numbered_slots(explanation)
        keys = [(s["answer"]["en"] or "")[:80] for s in slots if s["answer"]["en"]]

    stem = extract_stem(body, qtype)
    scenario = None
    if qtype == "yes_no":
        scenario, _, _ = extract_yes_no_parts(body)

    bl = body.lower()
    media_required = qtype in ("hotspot", "dropdown", "drag_drop") or any(
        x in bl
        for x in ("graphic", "image", "hot area", "diagram", "exhibit", "following exhibit")
    )

    note_m = re.search(
        r"Note:\s*(The question is included.+?distinctive result\.\s*)",
        body.replace("\t", " "),
        re.I,
    )

    from enrich_questions import build_interaction

    item = {
        "id": num,
        "exam": EXAM,
        "type": qtype,
        "source": {"pdfPage": None, "certyiqUrl": f"{CERTYIQ_BASE}/{num}"},
        "stem": bilingual(stem),
        "options": options,
        "answer": {
            "keys": keys,
            "text": bilingual(answer_raw if len(answer_raw) < 80 else None),
        },
        "explanation": bilingual(explanation),
        "references": references,
        "boxes": boxes,
        "media": {
            "required": media_required,
            "images": [],
            "pageRender": None,
        },
        "tags": [],
        "meta": {
            "note": note_m.group(1).strip() if note_m else None,
            "scenario": bilingual(scenario) if scenario else None,
            "translatedAt": None,
        },
    }
    interaction = build_interaction(item)
    if interaction:
        item["interaction"] = interaction
    return item


def apply_shared_scenario(questions: list[dict]) -> None:
    """Prepend last long scenario to short stems in a shared-setup group."""
    current: str | None = None
    for q in questions:
        note = (q.get("meta") or {}).get("note") or ""
        scen = (q.get("meta") or {}).get("scenario") or {}
        scen_en = (scen.get("en") or "").strip() if isinstance(scen, dict) else ""

        stem_en = (q.get("stem") or {}).get("en") or ""
        if scen_en and len(scen_en) > 80:
            current = scen_en
        elif len(stem_en) > 120 and "solution:" not in stem_en.lower():
            current = stem_en.split("Solution:")[0].strip() or stem_en[:500]

        if note and "identical set-up" in note.lower() and current:
            if len(stem_en) < 120 and current not in stem_en:
                q["stem"]["en"] = f"{current}\n\n{stem_en}".strip()

        if q.get("type") == "yes_no" and scen_en and scen_en not in stem_en:
                sol = ""
                if "Solution:" in stem_en:
                    parts = stem_en.split("Solution:", 1)
                    sol = "Solution:" + parts[1] if len(parts) > 1 else ""
                    stem_en = parts[0]
                q["stem"]["en"] = re.sub(
                    r"\s+",
                    " ",
                    f"{scen_en} {sol} {stem_en}".replace("  ", " "),
                ).strip()


def map_questions_to_pages(doc: fitz.Document) -> dict[int, int]:
    mapping = {}
    for i in range(doc.page_count):
        t = doc[i].get_text()
        for m in re.finditer(r"Question:\s*(\d+)", t):
            n = int(m.group(1))
            if n not in mapping:
                mapping[n] = i + 1
    return mapping


def save_page_media(doc: fitz.Document, qid: int | str, page_num: int) -> dict:
    """Extract embedded images only; skip blank / strip artifacts."""
    from image_utils import is_usable_diagram

    qdir = ASSETS_DIR / str(qid)
    qdir.mkdir(parents=True, exist_ok=True)
    page = doc[page_num - 1]
    images = []
    min_pixels = 120 * 120
    for idx, img in enumerate(page.get_images(full=True)):
        try:
            xref = img[0]
            pix = fitz.Pixmap(doc, xref)
            if pix.width * pix.height < min_pixels:
                continue
            if pix.n - pix.alpha > 3:
                pix = fitz.Pixmap(fitz.csRGB, pix)
            out = qdir / f"img-{idx}.png"
            pix.save(str(out))
            if is_usable_diagram(out):
                images.append(f"assets/questions/{qid}/img-{idx}.png")
            else:
                out.unlink(missing_ok=True)
        except Exception:
            continue
    return {"images": images, "pageRender": None}


def main() -> int:
    if not PDF_PATH.exists():
        print(f"PDF not found: {PDF_PATH}", file=sys.stderr)
        return 1

    print("Reading PDF text...")
    reader = PdfReader(str(PDF_PATH))
    full_text = "\n".join((p.extract_text() or "") for p in reader.pages)

    print("Opening PDF for images...")
    doc = fitz.open(str(PDF_PATH))
    page_map = map_questions_to_pages(doc)

    raw_chunks = split_question_chunks(full_text)
    print(f"Split into {len(raw_chunks)} question chunks")

    questions: list[dict] = []
    for seq, (pdf_num, body) in enumerate(raw_chunks, start=1):
        qid = pdf_num if pdf_num and 1 <= pdf_num <= 700 else seq
        q = parse_block(qid, body)
        if qid in page_map:
            q["source"]["pdfPage"] = page_map[qid]
        questions.append(q)

    # Deduplicate by id (keep block with valid answer keys, else longest stem)
    by_id: dict[int, dict] = {}
    for q in questions:
        existing = by_id.get(q["id"])
        if not existing:
            by_id[q["id"]] = q
            continue
        cur_keys = len(q.get("answer", {}).get("keys") or [])
        old_keys = len(existing.get("answer", {}).get("keys") or [])
        if cur_keys > old_keys:
            by_id[q["id"]] = q
        elif cur_keys == old_keys and len(q["stem"]["en"] or "") > len(existing["stem"]["en"] or ""):
            by_id[q["id"]] = q
    questions = [by_id[k] for k in sorted(by_id)]

    apply_shared_scenario(questions)

    print(f"Parsed {len(questions)} questions. Extracting media...")
    media_count = 0
    for q in questions:
        qid = q["id"]
        if not q["media"]["required"]:
            continue
        p = q["source"].get("pdfPage")
        if not p:
            continue
        media = save_page_media(doc, qid, p)
        q["media"]["images"] = media["images"]
        q["media"]["pageRender"] = media["pageRender"]
        media_count += 1

    doc.close()

    OUT_JSON.parent.mkdir(parents=True, exist_ok=True)
    with OUT_JSON.open("w", encoding="utf-8") as f:
        json.dump(questions, f, ensure_ascii=False, indent=2)

    type_counts: dict[str, int] = {}
    empty_keys = 0
    short_stem = 0
    for q in questions:
        type_counts[q["type"]] = type_counts.get(q["type"], 0) + 1
        if not q.get("answer", {}).get("keys"):
            empty_keys += 1
        if len((q.get("stem") or {}).get("en") or "") < 50:
            short_stem += 1

    manifest = {
        "exam": EXAM,
        "extractedAt": datetime.now(timezone.utc).isoformat(),
        "sourcePdf": PDF_PATH.name,
        "totalQuestions": len(questions),
        "typeCounts": type_counts,
        "mediaExtracted": media_count,
        "emptyAnswerKeys": empty_keys,
        "shortStems": short_stem,
        "certyiqBase": CERTYIQ_BASE,
    }
    with OUT_MANIFEST.open("w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)

    print(f"Wrote {OUT_JSON} ({len(questions)} questions)")
    print(f"Types: {type_counts}")
    print(f"Empty keys: {empty_keys}, short stems: {short_stem}")
    print(f"Media bundles: {media_count}")
    return 0


if __name__ == "__main__":
    sys.exit(main())