#!/usr/bin/env python3
"""Extract structured FRCS question bank from the source .docx.

Reads word/document.xml via zipfile, parses paragraphs and their styles,
reconstructs parts -> chapters -> questions and writes src/data/frcs-bank.json.
"""

from __future__ import annotations

import json
import re
import sys
import zipfile
from pathlib import Path
from xml.etree import ElementTree as ET

HERE = Path(__file__).resolve().parent
SITE = HERE.parent
ROOT = SITE.parent
SOURCE = ROOT / "FRCS_SBA_Question_Bank_300 (1).docx"
OUTPUT = SITE / "src" / "data" / "frcs-bank.json"

W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
NS = {"w": W_NS}


def tag(name: str) -> str:
    return f"{{{W_NS}}}{name}"


def extract_paragraphs(docx_path: Path):
    """Return list of (style, text, is_bold_any) paragraphs."""
    with zipfile.ZipFile(docx_path) as zf:
        with zf.open("word/document.xml") as f:
            xml_bytes = f.read()

    root = ET.fromstring(xml_bytes)
    body = root.find(tag("body"))
    if body is None:
        return []

    paragraphs = []
    for p in body.iter(tag("p")):
        style = "Normal"
        p_pr = p.find(tag("pPr"))
        if p_pr is not None:
            pstyle = p_pr.find(tag("pStyle"))
            if pstyle is not None:
                style = pstyle.get(tag("val")) or "Normal"

        # Collect runs, note bold runs
        text_parts = []
        any_bold = False
        for r in p.findall(tag("r")):
            r_pr = r.find(tag("rPr"))
            is_bold = False
            if r_pr is not None and r_pr.find(tag("b")) is not None:
                is_bold = True
            for t in r.findall(tag("t")):
                if t.text:
                    text_parts.append(t.text)
                    if is_bold:
                        any_bold = True
            for _br in r.findall(tag("br")):
                text_parts.append("\n")
            for _tab in r.findall(tag("tab")):
                text_parts.append("\t")
        text = "".join(text_parts).strip()
        paragraphs.append({"style": style, "text": text, "bold": any_bold})
    return paragraphs


PART_RE = re.compile(r"^PART\s+([IVX]+)\s*[:\-\u2014]?\s*(.*)$", re.IGNORECASE)
CHAPTER_RE = re.compile(
    r"^(?:Chapter|Ch\.?)\s*(\d+)\s*[:\.\-\u2014]?\s*(.*)$", re.IGNORECASE
)
CHAPTER_RE_ALT = re.compile(r"^(\d+)\.\s+([A-Z][^\n]{3,})$")
QUESTION_RE = re.compile(
    r"^(?:Q(?:uestion)?\.?\s*)?(\d{1,3})[\.\)]\s+(.*)$", re.IGNORECASE
)
Q_PREFIX_RE = re.compile(r"^Q(?:uestion)?\s*(\d{1,3})\b[\.\)\:\-\s]*(.*)$", re.IGNORECASE)
OPTION_RE = re.compile(r"^([A-E])[\.\)]\s*(.+)$")
ANSWER_RE = re.compile(
    r"^(?:Correct\s+Answer|Answer)\s*[:\-\u2014]\s*([A-E])\b\.?\s*(.*)$",
    re.IGNORECASE,
)
EXPLAIN_RE = re.compile(
    r"^(?:Explanation|Rationale)\s*[:\-\u2014]\s*(.*)$", re.IGNORECASE
)
RANGE_RE = re.compile(r"Questions?\s*\d+\s*[\-\u2013\u2014]\s*\d+", re.IGNORECASE)


def classify_heading(text: str, style: str):
    """Return ('part', label) | ('chapter', label) | ('range', label) | None."""
    if not text:
        return None
    if PART_RE.match(text):
        return ("part", text)
    if CHAPTER_RE.match(text):
        return ("chapter", text)
    if style == "Heading1" and text.upper().startswith("PART"):
        return ("part", text)
    if style in ("Heading2", "Heading3") and CHAPTER_RE.match(text):
        return ("chapter", text)
    return None


def is_question_start(text: str):
    if not text:
        return None
    m = Q_PREFIX_RE.match(text)
    if m:
        return int(m.group(1)), m.group(2).strip()
    m = QUESTION_RE.match(text)
    if m:
        n = int(m.group(1))
        # Avoid catching numbered option like "1. something" if rest looks tiny
        rest = m.group(2).strip()
        if 1 <= n <= 350 and len(rest) >= 5:
            return n, rest
    return None


def parse(paragraphs):
    parts: list = []
    current_part = None
    current_chapter = None
    current_question = None
    # Order: we accumulate into last chapter
    pending_explanation_lines: list[str] = []
    state = "idle"  # idle | stem | options | explanation

    def flush_question():
        nonlocal current_question, pending_explanation_lines
        if current_question is not None:
            if pending_explanation_lines:
                expl = " ".join(pending_explanation_lines).strip()
                current_question["explanation"] = (
                    (current_question.get("explanation", "") + " " + expl).strip()
                    if current_question.get("explanation")
                    else expl
                )
            pending_explanation_lines = []
            if current_chapter is not None:
                current_chapter["questions"].append(current_question)
        current_question = None

    def ensure_part(label="PART I"):
        nonlocal current_part
        if current_part is None:
            current_part = {"title": label, "chapters": []}
            parts.append(current_part)
        return current_part

    def ensure_chapter(label="Chapter"):
        nonlocal current_chapter
        if current_chapter is None:
            ensure_part()
            current_chapter = {
                "title": label,
                "range": "",
                "questions": [],
            }
            current_part["chapters"].append(current_chapter)
        return current_chapter

    for para in paragraphs:
        text = para["text"]
        style = para["style"]
        if not text:
            continue

        # Headings
        heading = classify_heading(text, style)
        if heading:
            kind, label = heading
            flush_question()
            state = "idle"
            if kind == "part":
                current_part = {"title": label, "chapters": []}
                parts.append(current_part)
                current_chapter = None
                continue
            if kind == "chapter":
                ensure_part()
                current_chapter = {
                    "title": label,
                    "range": "",
                    "questions": [],
                }
                current_part["chapters"].append(current_chapter)
                continue

        # Range hint line (e.g., "Questions 1-10")
        if RANGE_RE.search(text) and len(text) < 60:
            if current_chapter is not None and not current_chapter["range"]:
                current_chapter["range"] = text.strip()
            continue

        # Question start
        qs = is_question_start(text)
        if qs:
            n, stem = qs
            flush_question()
            ensure_chapter()
            current_question = {
                "number": n,
                "stem": stem,
                "options": [],
                "answer": "",
                "explanation": "",
            }
            state = "stem"
            continue

        # Option
        m_opt = OPTION_RE.match(text)
        if m_opt and current_question is not None:
            letter = m_opt.group(1)
            body = m_opt.group(2).strip()
            current_question["options"].append(f"{letter}. {body}")
            state = "options"
            continue

        # Answer
        m_ans = ANSWER_RE.match(text)
        if m_ans and current_question is not None:
            current_question["answer"] = m_ans.group(1).upper()
            tail = m_ans.group(2).strip()
            if tail:
                pending_explanation_lines.append(tail)
            state = "explanation"
            continue

        # Explicit explanation/rationale marker
        m_exp = EXPLAIN_RE.match(text)
        if m_exp and current_question is not None:
            tail = m_exp.group(1).strip()
            if tail:
                pending_explanation_lines.append(tail)
            state = "explanation"
            continue

        # Fallback: append to current state
        if current_question is not None:
            if state == "stem":
                current_question["stem"] = (
                    current_question["stem"] + " " + text
                ).strip()
            elif state == "options":
                # Continuation of last option
                if current_question["options"]:
                    current_question["options"][-1] = (
                        current_question["options"][-1] + " " + text
                    )
                else:
                    current_question["stem"] = (
                        current_question["stem"] + " " + text
                    ).strip()
            elif state == "explanation":
                pending_explanation_lines.append(text)

    flush_question()
    return parts


def main():
    if not SOURCE.exists():
        print(f"ERROR: source not found: {SOURCE}", file=sys.stderr)
        sys.exit(1)

    paragraphs = extract_paragraphs(SOURCE)
    print(f"Paragraphs read: {len(paragraphs)}")

    parts = parse(paragraphs)

    # If no parts, synthesize a single default one
    if not parts:
        parts = [{"title": "PART I", "chapters": []}]

    # Clean up: ensure every chapter has a sensible range if blank
    total_questions = 0
    for part in parts:
        for ch in part["chapters"]:
            qs = ch["questions"]
            total_questions += len(qs)
            if not ch["range"] and qs:
                first = qs[0]["number"]
                last = qs[-1]["number"]
                ch["range"] = f"Questions {first}-{last}"

    # Drop empty chapters and empty parts for cleanliness
    for part in parts:
        part["chapters"] = [c for c in part["chapters"] if c["questions"]]
    parts = [p for p in parts if p["chapters"]]

    data = {
        "title": "FRCS (Tr & Orth) Basic Sciences \u2014 SBA Question Bank",
        "subtitle": "300 Single Best Answer Questions",
        "author": "Professor Mohamed Imam",
        "parts": parts,
    }

    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT.write_text(
        json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    n_parts = len(parts)
    n_chapters = sum(len(p["chapters"]) for p in parts)
    print(
        f"Extracted {n_parts} parts, {n_chapters} chapters, {total_questions} questions"
    )
    print(f"Written: {OUTPUT}")


if __name__ == "__main__":
    main()
