#!/usr/bin/env python3
"""
Extract CertyIQ PDF questions into data/{exam}/questions-review.json.

Each entry is one Answer + Explanation block from the PDF, in document order.
Fields are kept as plain text for manual review before importing to questions.json.
"""
from __future__ import annotations

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

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))

from exam_config import ExamProfile, get_profile, normalize_pdf_text

MARKER_Q = re.compile(
    r"<<<Q:(\d+)>>>|<<<QUESTION:(\d+)>>>|Question:\s*(\d+)",
    re.I,
)
ANSWER_BLOCK = re.compile(
    r"\bAnswer:\s*(?P<ans>.*?)\s*Explanation:\s*",
    re.I | re.S,
)
MARKER_REFERENCE = re.compile(r"\nReference:\s*", re.I)
MARKER_NOTE = re.compile(r"\nNote:\s*", re.I)
NEXT_QUESTION = re.compile(
    r"\n(?:"
    r"Your company |You have |You need |You want |You plan |You are |You must |"
    r"Which |What should |What is |What are |How |Each |An administrator|Contoso|Given |"
    r"HOTSPOT|DRAG DROP|For each of the following|Does the solution|"
    r"A company|There is an? |After you|If you"
    r")",
    re.I,
)
OPTION_LINE = re.compile(
    r"(?:^|\n|\?\s+)\s*([A-G])\.\s*\n?\s*(.+?)(?=(?:^|\n|\?\s+)\s*[A-G]\.\s|\bAnswer:|\Z)",
    re.M | re.S,
)


def strip_q_markers(text: str) -> str:
    """Remove page-break Question:N tokens; options follow the last marker before Answer."""
    return re.sub(r"(?:<<<Q:\d+>>>\s*)+", " ", text)
SHARED_SETUP_NOTE = re.compile(
    r"Note:\s*The question is included.+?distinctive result\.\s*",
    re.I | re.S,
)


def parse_qids(text: str) -> list[int]:
    ids: list[int] = []
    for m in MARKER_Q.finditer(text):
        ids.append(int(m.group(1) or m.group(2) or m.group(3)))
    return ids


def trim_explanation_tail(exp: str) -> tuple[str, str | None, str | None]:
    cut = len(exp)
    ref_m = MARKER_REFERENCE.search(exp)
    if ref_m:
        cut = min(cut, ref_m.start())
    note_m = MARKER_NOTE.search(exp)
    if note_m:
        cut = min(cut, note_m.start())
    nxt_m = NEXT_QUESTION.search(exp)
    if nxt_m:
        cut = min(cut, nxt_m.start())

    explanation = exp[:cut].strip()
    reference = None
    if ref_m and ref_m.start() < len(exp):
        ref_end = note_m.start() if note_m and note_m.start() > ref_m.start() else len(exp)
        nxt2 = NEXT_QUESTION.search(exp, ref_m.end())
        if nxt2:
            ref_end = min(ref_end, nxt2.start())
        reference = exp[ref_m.end() : ref_end].strip()

    note = None
    if note_m:
        note_end = len(exp)
        nxt3 = NEXT_QUESTION.search(exp, note_m.end())
        if nxt3:
            note_end = nxt3.start()
        note = exp[note_m.end() : note_end].strip()
        note = re.sub(r"<<<Q:\d+>>>.*", "", note, flags=re.S).strip() or None

    return explanation, reference, note


def advance_end(exp_raw: str, explanation: str, reference: str | None, note: str | None) -> int:
    pos = len(explanation)
    if reference and MARKER_REFERENCE.search(exp_raw):
        pos = max(pos, exp_raw.find(reference) + len(reference))
    if note and MARKER_NOTE.search(exp_raw):
        idx = exp_raw.find(note)
        if idx >= 0:
            pos = max(pos, idx + len(note))
    nxt = NEXT_QUESTION.search(exp_raw, pos)
    return nxt.start() if nxt else len(exp_raw)


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


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: list[dict] = []
    seen: set[str] = set()
    for m in OPTION_LINE.finditer(region):
        key = m.group(1)
        if key in seen:
            continue
        text = re.sub(r"\s+", " ", m.group(2).strip())
        if text and len(text) > 1:
            seen.add(key)
            opts.append({"key": key, "text": text})
    if not opts and re.search(r"^A\.\s*Yes", region, re.M):
        opts = [{"key": "A", "text": "Yes"}, {"key": "B", "text": "No"}]
    return opts


def extract_stem(body: str) -> str:
    b = strip_q_markers(body)
    b = re.sub(r"^(HOTSPOT|DRAG DROP)\s*-\s*", "", b, flags=re.I | re.M)
    answer_pos = re.search(r"Answer:\s*", b, re.I)
    region = b[: answer_pos.start()] if answer_pos else b
    opt = re.search(r"(?:^|\n|\?\s+)\s*([A-G])\.\s*\S", region, re.M)
    stem = region[: opt.start()].strip() if opt else region.strip()
    return re.sub(r"\s+", " ", stem)


def normalize_answer_raw(raw: str) -> str:
    from clean_text import clean_extracted_text

    raw = clean_extracted_text(raw or "") or ""
    raw = re.sub(r"Question:\s*\d+", "", raw, flags=re.I)
    raw = re.sub(r"\s+", " ", raw.strip())
    if not raw:
        return ""
    # Keep first line only for letter answers (PDF sometimes appends junk)
    first = raw.split("\n")[0].strip()
    m = re.match(r"^([A-G](?:,\s*[A-G])*)", first, re.I)
    if m:
        return m.group(1).upper().replace(" ", "")
    if re.match(r"^(Yes|No)\b", first, re.I):
        return first.split()[0].capitalize()
    if re.match(r"^\d+\.", first):
        return first
    return first if len(first) <= 80 else raw[:80]


def parse_answer_keys(raw: str, qtype: str) -> list[str]:
    raw = normalize_answer_raw(raw)
    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) and len(raw) <= 7:
        return list(raw.upper())
    if raw in ("Yes", "No"):
        return [raw]
    if re.match(r"^\d+\.", raw):
        return [raw]
    if len(raw) <= 80:
        return [raw]
    return []


def repair_block_from_marker_text(item: dict) -> None:
    """Re-parse question fields from markerText when body is missing or broken."""
    marker = (item.get("markerText") or "").strip()
    if not marker or not re.search(r"\bAnswer:\s*", marker, re.I):
        return

    body = marker
    body = re.sub(r"^Answer:\s*", "", body, flags=re.I)
    body = re.sub(r"(?:<<<Q:\d+>>>\s*)+", " ", body)
    body = re.sub(r"(?:<<<QUESTION:\d+>>>\s*)+", " ", body, flags=re.I)
    body = body.strip()

    answers = list(re.finditer(r"\bAnswer:\s*", body, re.I))
    if answers:
        after = body[answers[-1].end() :]
        nxt = NEXT_QUESTION.search(after)
        if nxt:
            body = body[: answers[-1].end() + nxt.start()].strip()

    from clean_text import clean_extracted_text

    body = clean_extracted_text(body) or body
    if len(body) < 40:
        return

    qtype = detect_type(body)
    options = parse_options(body)
    ans_m = re.search(
        r"\bAnswer:\s*(?P<ans>.*?)\s*Explanation:\s*",
        body,
        re.I | re.S,
    )
    answer_raw = normalize_answer_raw(ans_m.group("ans")) if ans_m else None
    keys = parse_answer_keys(answer_raw or "", qtype) if answer_raw else []
    new_stem = extract_stem(body)

    cur_stem = (item.get("stem") or "").strip()
    cur_qt = (item.get("questionText") or "").strip()
    if len(new_stem) <= len(cur_stem) + 15 and cur_qt:
        return

    item["questionText"] = body
    item["stem"] = new_stem
    item["options"] = options
    item["type"] = qtype
    if answer_raw and keys:
        item["answer"] = {"raw": answer_raw, "keys": keys}
    if item.get("flags"):
        item["flags"]["missingOptions"] = not options and qtype in (
            "single_choice",
            "multi_select",
            "yes_no",
        )
        item["flags"]["missingAnswer"] = not keys and qtype not in DEFERRED_TYPES


def map_pdf_ids_by_position(text: str, blocks: list[dict], profile: ExamProfile) -> None:
    """Assign pdfId from Question:N markers in each block's raw PDF span."""
    from document_ids import assign_and_deduplicate_review_blocks

    merged = assign_and_deduplicate_review_blocks(
        blocks,
        certyiq_base=profile.certyiq_base,
        preface_seq_max=profile.preface_seq_max,
    )
    for item in merged:
        repair_block_from_marker_text(item)
    blocks[:] = merged


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 enrich_blocks_from_marker_split(text: str, blocks: list[dict], profile: ExamProfile) -> None:
    """Replace weak bodies with marker-split chunks from extract_pdf."""
    from clean_text import clean_extracted_text
    from extract_pdf import build_question_body_map, normalize_text as normalize_extract

    normalized = normalize_extract(text)
    for pat in profile.content_start_patterns:
        m = re.search(pat, normalized, re.I)
        if m:
            normalized = normalized[m.start() :]
            break
    bodies = build_question_body_map(normalized)
    by_pdf: dict[int, dict] = {}
    for b in blocks:
        pid = b.get("pdfId")
        if pid is not None:
            by_pdf[int(pid)] = b

    for pid, alt in bodies.items():
        item = by_pdf.get(pid)
        if not item:
            continue
        cur = item.get("questionText") or ""
        if _body_quality(alt) <= _body_quality(cur):
            continue
        alt = clean_extracted_text(alt) or alt
        item["questionText"] = alt
        item["type"] = detect_type(alt)
        item["stem"] = extract_stem(alt)
        item["options"] = parse_options(alt)
        keys = parse_answer_keys(
            item.get("answer", {}).get("raw") or "", item["type"]
        )
        item["answer"]["keys"] = keys
        qtype = item["type"]
        item["flags"]["missingOptions"] = not item["options"] and qtype in (
            "single_choice",
            "multi_select",
            "yes_no",
        )


def extract_questions(text: str, profile: ExamProfile) -> list[dict]:
    blocks: list[dict] = []
    cursor = 0
    seq = 0

    for m in ANSWER_BLOCK.finditer(text):
        gap = text[cursor : m.start()]
        gap_ids = parse_qids(gap)

        body_raw = SHARED_SETUP_NOTE.sub("", gap)
        body_raw = re.sub(r"(?:<<<Q:\d+>>>\s*)+", " ", body_raw).strip()

        answer_raw = normalize_answer_raw(m.group("ans"))
        exp_start = m.end()
        next_m = ANSWER_BLOCK.search(text, exp_start)
        exp_slice = text[exp_start : next_m.start() if next_m else len(text)]
        marker_text = text[cursor : exp_start + len(exp_slice)]
        markers_in_explanation = parse_qids(exp_slice)
        explanation, reference, note = trim_explanation_tail(exp_slice)
        from clean_text import clean_extracted_text

        explanation = clean_extracted_text(explanation)
        reference = clean_extracted_text(reference)
        note = clean_extracted_text(note)
        body_raw = clean_extracted_text(body_raw) or ""

        seq += 1
        qtype = detect_type(body_raw)
        options = parse_options(body_raw)
        keys = parse_answer_keys(answer_raw, qtype)
        refs_urls = re.findall(r"https?://[^\s\)\]]+", reference or "")

        blocks.append(
            {
                "seq": seq,
                "pdfId": None,
                "exam": profile.code,
                "type": qtype,
                "source": {
                    "certyiqUrl": None,
                },
                "markersInGap": gap_ids,
                "markersInExplanation": markers_in_explanation,
                "markerText": marker_text,
                "questionText": body_raw,
                "stem": extract_stem(body_raw),
                "options": options,
                "answer": {
                    "raw": answer_raw or None,
                    "keys": keys,
                },
                "explanation": explanation or None,
                "reference": reference,
                "references": refs_urls,
                "note": note,
                "flags": {
                    "missingPdfId": True,
                    "missingAnswer": not answer_raw
                    and qtype not in ("drag_drop", "hotspot", "dropdown"),
                    "missingExplanation": not explanation,
                    "missingOptions": not options
                    and qtype in ("single_choice", "multi_select", "yes_no"),
                },
            }
        )

        cursor = exp_start + advance_end(exp_slice, explanation, reference, note)

    map_pdf_ids_by_position(text, blocks, profile)
    enrich_blocks_from_marker_split(text, blocks, profile)
    return blocks


def recover_missing_blocks(text: str, blocks: list[dict], profile: ExamProfile) -> int:
    """Add review blocks for PDF bodies missing from Answer-block split."""
    from clean_text import clean_extracted_text
    from document_ids import infer_pdf_id_from_text
    from extract_pdf import build_question_body_map, normalize_text as normalize_extract

    normalized = normalize_extract(text)
    for pat in profile.content_start_patterns:
        m = re.search(pat, normalized, re.I)
        if m:
            normalized = normalized[m.start() :]
            break
    bodies = build_question_body_map(normalized)
    existing = {int(b["pdfId"]) for b in blocks if b.get("pdfId") is not None}
    added = 0
    max_seq = max((int(b.get("seq") or 0) for b in blocks), default=0)

    for _map_key, body in sorted(bodies.items()):
        if not body or len(body) < 60:
            continue
        if not re.search(r"\bAnswer:\s*", body, re.I):
            continue

        inferred = infer_pdf_id_from_text(body)
        if inferred is None or inferred in existing:
            continue

        body = clean_extracted_text(body) or body
        ans_m = re.search(
            r"\bAnswer:\s*(?P<ans>.*?)\s*Explanation:\s*",
            body,
            re.I | re.S,
        )
        exp_m = re.search(r"\bExplanation:\s*", body, re.I)
        answer_raw = ""
        exp_slice = ""
        if not ans_m and re.search(r"\bAnswer:\s*", body, re.I) and exp_m:
            ans_start = re.search(r"\bAnswer:\s*", body, re.I)
            answer_raw = normalize_answer_raw(
                body[ans_start.end() : exp_m.start()]
            )
            exp_slice = body[exp_m.end() :]
        elif ans_m:
            answer_raw = normalize_answer_raw(ans_m.group("ans"))
            exp_slice = body[ans_m.end() :]
        elif re.search(r"\bAnswer:\s*", body, re.I):
            continue
        else:
            continue

        explanation, reference, note = trim_explanation_tail(exp_slice) if exp_slice else ("", None, None)
        explanation = clean_extracted_text(explanation)
        reference = clean_extracted_text(reference)
        note = clean_extracted_text(note)
        qtype = detect_type(body)
        options = parse_options(body)
        keys = parse_answer_keys(answer_raw, qtype)
        refs_urls = re.findall(r"https?://[^\s\)\]]+", reference or "")

        max_seq += 1
        blocks.append(
            {
                "seq": max_seq,
                "pdfId": inferred,
                "exam": profile.code,
                "type": qtype,
                "source": {"certyiqUrl": f"{profile.certyiq_base}/{inferred}"},
                "markersInGap": [],
                "markersInExplanation": parse_qids(exp_slice) if exp_slice else [],
                "markerText": body,
                "questionText": body,
                "stem": extract_stem(body),
                "options": options,
                "answer": {"raw": answer_raw or None, "keys": keys},
                "explanation": explanation or None,
                "reference": reference,
                "references": refs_urls,
                "note": note,
                "flags": {
                    "missingPdfId": False,
                    "missingAnswer": not answer_raw
                    and qtype not in DEFERRED_TYPES,
                    "missingExplanation": not explanation,
                    "missingOptions": not options
                    and qtype in ("single_choice", "multi_select", "yes_no"),
                    "recoveredFromMarkerSplit": True,
                },
            }
        )
        existing.add(inferred)
        added += 1

    map_pdf_ids_by_position(text, blocks, profile)
    return added


DEFERRED_TYPES = frozenset({"hotspot", "dropdown", "drag_drop"})


def main() -> int:
    parser = argparse.ArgumentParser(description="Extract PDF to questions-review.json")
    parser.add_argument(
        "--exam",
        default="az-104",
        choices=sorted(__import__("exam_config").EXAMS.keys()),
    )
    args = parser.parse_args()
    profile = get_profile(args.exam)

    if not profile.pdf_path.exists():
        print(f"PDF not found: {profile.pdf_path}", file=sys.stderr)
        return 1

    print(f"Reading PDF ({profile.code})...")
    reader = PdfReader(str(profile.pdf_path))
    full_text = "\n".join((p.extract_text() or "") for p in reader.pages)

    print("Extracting blocks (Answer + Explanation)...")
    normalized = normalize_pdf_text(full_text, profile)
    blocks = extract_questions(normalized, profile)
    recovered_count = recover_missing_blocks(full_text, blocks, profile)
    print(f"  Blocks: {len(blocks)} (marker-recovered: {recovered_count})")

    pdf_ids = {b["pdfId"] for b in blocks if b["pdfId"]}
    missing = sorted(set(range(1, profile.expected_questions + 1)) - pdf_ids)
    dup_ids = sorted(k for k in pdf_ids if sum(1 for b in blocks if b["pdfId"] == k) > 1)

    flags_summary = {
        "missingPdfId": sum(1 for b in blocks if b["flags"]["missingPdfId"]),
        "missingAnswer": sum(1 for b in blocks if b["flags"]["missingAnswer"]),
        "missingExplanation": sum(1 for b in blocks if b["flags"]["missingExplanation"]),
        "missingOptions": sum(1 for b in blocks if b["flags"]["missingOptions"]),
    }

    type_counts: dict[str, int] = {}
    for b in blocks:
        type_counts[b["type"]] = type_counts.get(b["type"], 0) + 1

    profile.data_dir.mkdir(parents=True, exist_ok=True)
    profile.questions_review_json.write_text(
        json.dumps(blocks, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    manifest = {
        "exam": profile.code,
        "extractedAt": datetime.now(timezone.utc).isoformat(),
        "sourcePdf": profile.pdf_name,
        "outputFile": profile.questions_review_json.name,
        "totalBlocks": len(blocks),
        "uniquePdfIds": len(pdf_ids),
        "expectedQuestions": profile.expected_questions,
        "missingPdfIds": missing,
        "missingCount": len(missing),
        "duplicatePdfIds": dup_ids[:50],
        "duplicateCount": len(dup_ids),
        "typeCounts": type_counts,
        "flags": flags_summary,
        "notes": [
            "seq = order in PDF (1-based). Use for review when pdfId is wrong.",
            "pdfId = first Question:N after Answer: in markerText (handles page-break markers).",
            "markersInGap = Question:N markers between previous answer and this one.",
        ],
    }
    profile.review_manifest_json.write_text(
        json.dumps(manifest, indent=2), encoding="utf-8"
    )

    print(f"Wrote {profile.questions_review_json}")
    print(f"  Blocks: {len(blocks)}, unique pdfId: {len(pdf_ids)}, missing ids: {len(missing)}")
    print(f"  Types: {type_counts}")
    print(f"  Flags: {flags_summary}")
    print(f"  Sample seq 1-5 pdfIds:", [b["pdfId"] for b in blocks[:5]])
    for b in blocks:
        if b.get("pdfId") in (18, 19, 20, 21):
            print(
                f"  pdfId {b['pdfId']} seq {b['seq']} "
                f"ans={b['answer']['raw']!r} stem={b['stem'][:60]}..."
            )
    return 0


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