#!/usr/bin/env python3
"""Extract Śrīmad-Bhāgavatam verses from HTML into structured JSON files.

Structure: Canto → Chapter → Verse
Output goes to app/lib/books/sb/

The SB HTML uses a different class mapping than BG but shares the same
ScaGoudy font encoding and the same verse structure:
  Devanagari (skip) → Transliteration pādas → Synonyms → Translation → Purport
"""

import json
import os
import re
import sys

from bs4 import BeautifulSoup, Tag
from indic_transliteration import sanscript
from indic_transliteration.sanscript import transliterate

# ---------------------------------------------------------------------------
# ScaGoudy → IAST Unicode mapping (shared with bg_extract.py)
# ---------------------------------------------------------------------------
SCAGOUDY_TO_IAST = {
    "ä": "ā",
    "Ä": "Ā",
    "å": "ṛ",
    "Å": "Ṛ",
    "ñ": "ṣ",
    "Ñ": "Ṣ",
    "ö": "ṭ",
    "Ö": "Ṭ",
    "ë": "ṇ",
    "Ë": "Ṇ",
    "ò": "ḍ",
    "Ò": "Ḍ",
    "ç": "ś",
    "Ç": "Ś",
    "ì": "ṅ",
    "ï": "ñ",
    "ü": "ū",
    "Ü": "Ū",
    "é": "ī",
    "É": "Ī",
    "à": "ṁ",
    "ù": "ḥ",
}


def decode_scagoudy(text: str) -> str:
    """Convert ScaGoudy-encoded text to proper IAST Unicode."""
    for src, dst in SCAGOUDY_TO_IAST.items():
        text = text.replace(src, dst)
    return text


def decode_html_content(tag: Tag) -> str:
    """Extract inner HTML from a tag, decode ScaGoudy, preserve formatting tags."""
    inner = "".join(str(child) for child in tag.children)
    inner = decode_scagoudy(inner)
    # Remove span tags with CSS classes (artifacts from RTF conversion)
    inner = re.sub(r'<span[^>]*class="[^"]*"[^>]*>(.*?)</span>', r"\1", inner)
    inner = re.sub(r"\s+", " ", inner).strip()
    return inner


def decode_text(tag: Tag) -> str:
    """Extract plain text from a tag and decode ScaGoudy."""
    text = tag.get_text()
    return decode_scagoudy(text).strip()


# ---------------------------------------------------------------------------
# Devanagari helpers (same as bg_extract.py)
# ---------------------------------------------------------------------------
def normalize_devanagari_anusvara(text: str) -> str:
    """Replace explicit class nasals before same-class consonants with anusvāra."""
    replacements = [
        (r"ङ्([कखगघ])", r"ं\1"),
        (r"ञ्([चछजझ])", r"ं\1"),
    ]
    for pattern, repl in replacements:
        text = re.sub(pattern, repl, text)
    return text


def join_virama_words(text: str) -> str:
    """Join Devanagari words where the previous word ends with virāma."""
    return re.sub(r"्\s+([क-ह])", r"्\1", text)


def iast_to_devanagari(iast_text: str) -> str:
    """Convert IAST transliteration to Devanagari script."""
    lines = iast_text.split("\n")
    result_lines = []
    for line in lines:
        clean = line.replace("-", "")
        clean = clean.replace("ṁ", "ṃ")
        deva = transliterate(clean, sanscript.IAST, sanscript.DEVANAGARI)
        deva = normalize_devanagari_anusvara(deva)
        deva = join_virama_words(deva)
        result_lines.append(deva)
    return "\n".join(result_lines)


def pair_padas(padas: list) -> str:
    """Pair pādas into verse lines (2 pādas per line)."""
    if not padas:
        return ""
    lines = []
    i = 0
    if len(padas) % 2 == 1:
        lines.append(padas[0])
        i = 1
    while i < len(padas) - 1:
        lines.append(padas[i] + " " + padas[i + 1])
        i += 2
    if i < len(padas):
        lines.append(padas[i])
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Synonym parser (same as bg_extract.py)
# ---------------------------------------------------------------------------
def parse_synonyms(text: str) -> list:
    """Parse synonym text into structured array."""
    result = []
    entries = text.split(";")
    for entry in entries:
        entry = entry.strip()
        if not entry:
            continue
        if "—" in entry:
            parts = entry.split("—", 1)
        elif "--" in entry:
            parts = entry.split("--", 1)
        else:
            continue
        if len(parts) == 2:
            word = parts[0].strip().strip("., ")
            meaning = parts[1].strip()
            if word and meaning:
                result.append({"word": word, "meanings": [meaning]})
    return result


def parse_synonym_html(tag: Tag) -> list:
    text = decode_text(tag)
    return parse_synonyms(text)


# ---------------------------------------------------------------------------
# CSS class helpers
# ---------------------------------------------------------------------------

# All Devanagari classes (Indevr + RM Devanagari) — content to skip
DEVANAGARI_CLASSES = frozenset(
    [
        "p23",
        "p81",
        "p91",
        "p101",
        "p102",
        "p105",
        "p113",  # Indevr
        "p126",
        "p128",
        "p129",  # RM Devanagari
    ]
)

# Transliteration pāda classes (center-aligned italic ScaGoudy)
TRANSLITERATION_CLASSES = frozenset(["p24"])

# Classes that are sometimes transliteration pādas (when they appear between
# Devanagari and synonyms) but otherwise are embedded verse quotes in purport.
# p41 and p92 are used for prose-style transliterations (common in Canto 5).
# p59 is a rare center-aligned variant also used for transliteration.
# p17 is primarily a purport class (780 uses) but occasionally holds transliteration
#   pādas (e.g. SB 5.24.7, SB 6.19.8) — the state machine ensures it's only
#   captured as transliteration when between Devanagari and synonyms.
# p114 is an indented ScaGoudy paragraph used once for a transliteration pāda.
MAYBE_TRANSLITERATION_CLASSES = frozenset(["p33", "p41", "p92", "p59", "p17", "p114"])

# Synonym classes (color #000080)
# p117 is a rare variant (2 occurrences) with same margins as p25 but no color
SYNONYM_CLASSES = frozenset(["p25", "p124", "p117"])

# Translation/purport section header classes
SECTION_HEADER_CLASSES = frozenset(["p26", "p28"])

# Translation text classes (color #800000, bold)
# p108 is maroon-colored ScaGoudy (3 occurrences) used as translation text
# p118 is a rare variant (1 occurrence) with same margins as p27 but no color
TRANSLATION_CLASSES = frozenset(["p27", "p98", "p108", "p118"])

# Purport paragraph classes — a wide set because the HTML uses many variants
PURPORT_CLASSES = frozenset(
    [
        "p29",
        "p14",
        "p17",
        "p18",
        "p30",
        "p11",
        "p45",
        "p67",
        "p77",
        "p80",
        "p117",
        "p118",
        "p119",
        "p84",
        "p79",
        "p133",
        "p135",
        "p136",
        "p62",
        "p44",
        "p83",
        "p96",  # chapter summaries (inside Summary entries)
        "p5",  # general text paragraphs (used in purport/summary context)
    ]
)

# Embedded verse quotations inside purport (center-aligned) — include in purport
PURPORT_VERSE_CLASSES = frozenset(
    [
        "p46",
        "p33",
        "p41",
        "p12",
        "p13",  # verse quotes & numbered lists in purport
        "p15",
        "p47",
        "p55",
        "p56",
        "p58",  # center-aligned purport variants
        "p66",
        "p68",  # center-aligned purport variants
        "p34",  # small center-aligned text in purport
        "p69",  # purport paragraph variant
        "p93",  # purport paragraph variant
        "p107",  # purport paragraph variant
    ]
)

# End-of-verse citation references (green, e.g. "[Bhāg. 1962:1.1.1]")
CITATION_CLASSES = frozenset(["p31"])

# Chapter-end markers ("Thus end the Bhaktivedanta purports…")
CHAPTER_END_CLASS = "p32"

# Chapter header bookmark (e.g. "SB 1.1: Questions by the Sages")
CHAPTER_BOOKMARK_CLASS = "p8"

# Chapter heading (e.g. "CHAPTER ONE\nQuestions by the Sages")
CHAPTER_HEADING_CLASS = "p19"

# Verse reference marker (e.g. "SB 1.1.1" or "SB. 6.1.1")
VERSE_REF_CLASS = "p20"

# TEXT N header
TEXT_HEADER_CLASS = "p21"

# "TEXT" label
TEXT_LABEL_CLASS = "p22"

# Canto/volume boundary markers — skip
CANTO_MARKER_CLASSES = frozenset(["p3", "p74"])

# Section title classes that are NOT verse content — skip
SKIP_CLASSES = frozenset(
    [
        "p4",
        "p9",
        "p49",
        "p50",  # canto/volume headers, dedication, glossary
        "p10",
        "p35",
        "p95",
        "p134",  # decorative / Times font
        "p48",  # glossary/author markers
        "p22",  # "TEXT" label — we use p21 for TEXT N
    ]
)


def get_css_class(tag: Tag) -> str:
    """Get the primary CSS class of a tag."""
    classes = tag.get("class", [])
    return classes[0] if classes else ""


def _tag_uses_devanagari_font(tag: Tag) -> bool:
    """Check if a tag's CSS class is a known Devanagari class."""
    return get_css_class(tag) in DEVANAGARI_CLASSES


# ---------------------------------------------------------------------------
# Reference parsing
# ---------------------------------------------------------------------------

# Matches "SB 1.1.1" or "SB. 6.1.1" — captures canto, chapter, verse
_VERSE_REF_RE = re.compile(r"SB\.?\s*(\d+)\.(\d+)\.(\d+)")
# Matches verse range like "SB 1.4.17-18" — captures start and end verse nums
_VERSE_RANGE_RE = re.compile(r"SB\.?\s*\d+\.\d+\.(\d+)\s*-\s*(\d+)")
# Matches chapter reference like "SB 1.1:" or "SB. 6.1:"
_CHAPTER_REF_RE = re.compile(r"SB\.?\s*(\d+)\.(\d+)\s*:")
# Matches summary reference like "SB 5.1 Summary"
_SUMMARY_RE = re.compile(r"SB\.?\s*(\d+)\.(\d+)\s+Summary", re.IGNORECASE)
# Matches invocation reference like "SB 2.1.invo"
_INVO_RE = re.compile(r"SB\.?\s*(\d+)\.(\d+)\.invo", re.IGNORECASE)


def _parse_verse_ref(text: str):
    """Parse a verse reference string.

    Returns (canto, chapter, verse, verse_end, clean_ref, ref_type) or None.
    ref_type is one of: 'verse', 'summary', 'invocation'
    """
    text = text.strip()

    # Check for Summary
    m = _SUMMARY_RE.search(text)
    if m:
        canto = int(m.group(1))
        chapter = int(m.group(2))
        return (canto, chapter, 0, None, f"SB {canto}.{chapter} Summary", "summary")

    # Check for Invocation
    m = _INVO_RE.search(text)
    if m:
        canto = int(m.group(1))
        chapter = int(m.group(2))
        return (
            canto,
            chapter,
            0,
            None,
            f"SB {canto}.{chapter} Invocation",
            "invocation",
        )

    # Check for verse range in the reference text
    # The ref may contain multiple individual refs then a range at the end:
    # "SB 1.4.17, SB 1.4.18, SB 1.4.17-18"
    # Or just a plain range: "SB 10.1.5-7"
    # We want the last meaningful reference.
    verse_refs = list(_VERSE_REF_RE.finditer(text))
    if not verse_refs:
        return None

    # Use the first match for canto/chapter/verse
    first = verse_refs[0]
    canto = int(first.group(1))
    chapter = int(first.group(2))
    verse = int(first.group(3))

    # Check for range
    verse_end = None
    range_match = _VERSE_RANGE_RE.search(text)
    if range_match:
        verse_end = int(range_match.group(2))
        # The first verse number from the range
        verse = int(range_match.group(1))

    # Build clean reference
    if verse_end is not None:
        clean_ref = f"SB {canto}.{chapter}.{verse}-{verse_end}"
    else:
        clean_ref = f"SB {canto}.{chapter}.{verse}"

    return (canto, chapter, verse, verse_end, clean_ref, "verse")


# ---------------------------------------------------------------------------
# Chapter title parsing from p8 tags
# ---------------------------------------------------------------------------
def _parse_chapter_bookmark(text: str):
    """Parse a chapter bookmark like 'SB 1.1: Questions by the Sages'.

    Returns (canto, chapter, title) or None.
    """
    text = text.strip()
    m = re.match(r"SB\.?\s*(\d+)\.(\d+)\s*:\s*(.+)", text)
    if m:
        canto = int(m.group(1))
        chapter = int(m.group(2))
        title = decode_scagoudy(m.group(3).strip())
        return (canto, chapter, title)
    return None


# ---------------------------------------------------------------------------
# State machine
# ---------------------------------------------------------------------------


def _finalize_verse(verse: dict):
    """Build transliteration and Devanagari from pādas. Remove internal fields."""
    padas = verse.get("_padas", [])
    verse["transliteration"] = "\n".join(padas)
    if padas:
        verse["devanagari"] = iast_to_devanagari(pair_padas(padas))
    else:
        verse["devanagari"] = ""
    if "_padas" in verse:
        del verse["_padas"]
    if "_state" in verse:
        del verse["_state"]
    if "_ref_type" in verse:
        del verse["_ref_type"]


def extract_verses(soup: BeautifulSoup, max_canto=99, max_chapter=999, max_verse=9999):
    """Extract verses from parsed SB HTML using a state machine.

    Returns (verses, chapter_info) where chapter_info is:
      { (canto, chapter): title_string }
    """
    paragraphs = soup.find_all("p")

    verses = []
    current_verse = None
    state = "IDLE"  # IDLE, DEVANAGARI, TRANSLITERATION, SYNONYMS,
    # TRANSLATION, PURPORT, SUMMARY

    chapter_titles = {}  # (canto, chapter) -> title

    for p in paragraphs:
        cls = get_css_class(p)
        text = p.get_text().strip()

        # --- Chapter bookmark (p8) ---
        if cls == CHAPTER_BOOKMARK_CLASS:
            parsed = _parse_chapter_bookmark(text)
            if parsed:
                canto, chapter, title = parsed
                chapter_titles[(canto, chapter)] = title
            # Don't change state — bookmarks appear before the chapter heading
            continue

        # --- Canto / volume markers — skip ---
        if cls in CANTO_MARKER_CLASSES or cls in SKIP_CLASSES:
            continue

        # --- Chapter heading (p19) — skip ---
        if cls == CHAPTER_HEADING_CLASS:
            continue

        # --- Verse reference (p20) ---
        if cls == VERSE_REF_CLASS:
            # Finalize previous verse
            if current_verse and current_verse.get("_ref_type") != "summary":
                _finalize_verse(current_verse)
                verses.append(current_verse)
            elif current_verse and current_verse.get("_ref_type") == "summary":
                # Discard summary entries — they're chapter preambles, not verses
                pass
            current_verse = None
            state = "IDLE"

            ref = _parse_verse_ref(text)
            if ref is None:
                continue

            canto, chapter, verse_num, verse_end, clean_ref, ref_type = ref

            if canto > max_canto:
                break
            if canto == max_canto and chapter > max_chapter:
                break
            if canto == max_canto and chapter == max_chapter and verse_num > max_verse:
                break

            if ref_type == "summary":
                # Start collecting summary text but mark it so we skip it
                current_verse = {
                    "_ref_type": "summary",
                    "_state": "SUMMARY",
                }
                state = "SUMMARY"
                continue

            current_verse = {
                "book_title": "Śrīmad-Bhāgavatam",
                "canto": canto,
                "chapter": chapter,
                "verse": verse_num,
                "reference": clean_ref,
                "devanagari": "",
                "transliteration": "",
                "_padas": [],
                "synonyms": [],
                "translation": "",
                "purport": [],
                "_state": "VERSE_START",
                "_ref_type": ref_type,
            }
            if verse_end is not None:
                current_verse["verse_end"] = verse_end
            state = "VERSE_START"
            continue

        # --- Summary state: just skip everything until next p20 ---
        if state == "SUMMARY":
            continue

        # --- TEXT N header (p21) ---
        if cls == TEXT_HEADER_CLASS and current_verse:
            continue

        # --- TEXT label (p22) ---
        if cls == TEXT_LABEL_CLASS:
            continue

        # --- Devanagari (Indevr / RM Devanagari) — SKIP, derive from transliteration ---
        if cls in DEVANAGARI_CLASSES and current_verse:
            state = "DEVANAGARI"
            current_verse["_state"] = state
            continue

        # --- Transliteration pādas (p24) ---
        if (
            cls in TRANSLITERATION_CLASSES
            and current_verse
            and state in ("VERSE_START", "DEVANAGARI", "TRANSLITERATION")
        ):
            state = "TRANSLITERATION"
            current_verse["_state"] = state
            line = decode_text(p)
            if line:
                current_verse["_padas"].append(line)
            continue

        # --- Maybe-transliteration (p33) — only if we're in transliteration context ---
        if (
            cls in MAYBE_TRANSLITERATION_CLASSES
            and current_verse
            and state in ("DEVANAGARI", "TRANSLITERATION")
        ):
            state = "TRANSLITERATION"
            current_verse["_state"] = state
            line = decode_text(p)
            if line:
                current_verse["_padas"].append(line)
            continue

        # --- Synonyms (p25, p124) ---
        if cls in SYNONYM_CLASSES and current_verse:
            state = "SYNONYMS"
            current_verse["_state"] = state
            current_verse["synonyms"] = parse_synonym_html(p)
            continue

        # --- Section headers: TRANSLATION / PURPORT (p26, p28) ---
        if cls in SECTION_HEADER_CLASSES and current_verse:
            header_text = text.strip().upper()
            if "TRANSLATION" in header_text:
                state = "TRANSLATION"
                current_verse["_state"] = state
            elif "PURPORT" in header_text:
                state = "PURPORT"
                current_verse["_state"] = state
            continue

        # --- Translation text (p27, p98) ---
        # Also handle mislabeled headers: in some verses (e.g. SB 10.4.6, 10.4.7)
        # the TRANSLATION header says "PURPORT", so state may be PURPORT but
        # translation is still empty. If we see a p27 tag and translation is empty,
        # capture it as translation regardless of state.
        if (
            cls in TRANSLATION_CLASSES
            and current_verse
            and (
                state == "TRANSLATION"
                or (state == "PURPORT" and not current_verse.get("translation"))
            )
        ):
            current_verse["translation"] = decode_html_content(p)
            state = "POST_TRANSLATION"
            current_verse["_state"] = state
            continue

        # --- p71 is blue-colored translation (miscolored — treat as translation) ---
        if cls == "p71" and current_verse and state == "TRANSLATION":
            current_verse["translation"] = decode_html_content(p)
            state = "POST_TRANSLATION"
            current_verse["_state"] = state
            continue

        # --- Purport paragraphs ---
        if state == "PURPORT" and current_verse and cls in PURPORT_CLASSES:
            content = decode_html_content(p)
            if content and "Thus end the Bhaktivedanta" not in text:
                current_verse["purport"].append(content)
            continue

        # --- Embedded verse quotations inside purport ---
        if state == "PURPORT" and current_verse and cls in PURPORT_VERSE_CLASSES:
            content = decode_html_content(p)
            if content:
                current_verse["purport"].append(content)
            continue

        # --- Citation / end-of-verse reference (p31) ---
        if cls in CITATION_CLASSES and current_verse:
            # These mark the end of a verse's content but we don't store them
            # The verse continues to the next p20
            continue

        # --- Chapter end marker (p32) ---
        if cls == CHAPTER_END_CLASS:
            if current_verse and current_verse.get("_ref_type") not in ("summary",):
                _finalize_verse(current_verse)
                verses.append(current_verse)
            current_verse = None
            state = "IDLE"
            continue

    # Don't forget the last verse
    if current_verse and current_verse.get("_ref_type") not in ("summary",):
        _finalize_verse(current_verse)
        verses.append(current_verse)

    return verses, chapter_titles


# ---------------------------------------------------------------------------
# JSON output
# ---------------------------------------------------------------------------


def write_json(data, filepath):
    """Write JSON with proper Unicode and formatting."""
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"  Written: {filepath}")


def main():
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    html_path = os.path.join(base_dir, "public", "html", "sb.html")
    output_dir = os.path.join(base_dir, "app", "lib", "books", "sb")

    print(f"Reading {html_path}...")
    with open(html_path, "r", encoding="utf-8") as f:
        soup = BeautifulSoup(f, "html.parser")

    print("Extracting all SB verses...")
    verses, chapter_titles = extract_verses(soup)

    print(f"\nExtracted {len(verses)} verses across {len(chapter_titles)} chapters")

    # Group by (canto, chapter) to build intro.json
    cantos_info = {}  # canto -> { chapter -> verse_count }
    for verse in verses:
        c = verse["canto"]
        ch = verse["chapter"]
        if c not in cantos_info:
            cantos_info[c] = {}
        if ch not in cantos_info[c]:
            cantos_info[c][ch] = 0
        cantos_info[c][ch] += 1

    # Canto names
    canto_names = {
        1: "Creation",
        2: "The Cosmic Manifestation",
        3: "The Status Quo",
        4: "The Creation of the Fourth Order",
        5: "The Creative Impetus",
        6: "Prescribed Duties for Mankind",
        7: "The Science of God",
        8: "Withdrawal of the Cosmic Creations",
        9: "Liberation",
        10: "The Summum Bonum",
    }

    # Build intro.json
    total_cantos = max(cantos_info.keys()) if cantos_info else 0
    cantos_list = []
    for canto_num in sorted(cantos_info.keys()):
        chapters_in_canto = cantos_info[canto_num]
        chapters_list = []
        for ch_num in sorted(chapters_in_canto.keys()):
            title = chapter_titles.get((canto_num, ch_num), "")
            chapters_list.append(
                {
                    "number": ch_num,
                    "title": title,
                    "totalVerses": chapters_in_canto[ch_num],
                }
            )
        cantos_list.append(
            {
                "number": canto_num,
                "title": canto_names.get(canto_num, f"Canto {canto_num}"),
                "totalChapters": len(chapters_list),
                "chapters": chapters_list,
            }
        )

    intro = {
        "book": "sb",
        "title": "Śrīmad-Bhāgavatam",
        "author": "His Divine Grace A.C. Bhaktivedanta Swami Prabhupāda",
        "totalCantos": total_cantos,
        "cantos": cantos_list,
    }
    write_json(intro, os.path.join(output_dir, "intro.json"))

    # Write verse JSON files:
    #   app/lib/books/sb/cantos/CC/chapters/CH/verses/VV.json
    for verse in verses:
        canto = str(verse["canto"]).zfill(2)
        ch = str(verse["chapter"]).zfill(2)
        vs = str(verse["verse"]).zfill(2)
        filepath = os.path.join(
            output_dir, "cantos", canto, "chapters", ch, "verses", f"{vs}.json"
        )
        write_json(verse, filepath)

    # Print summary for verification
    print("\n--- Verification ---")
    for canto_num in sorted(cantos_info.keys()):
        chapters_in_canto = cantos_info[canto_num]
        total_verses_in_canto = sum(chapters_in_canto.values())
        print(
            f"\nCanto {canto_num}: {len(chapters_in_canto)} chapters, {total_verses_in_canto} verses"
        )
        for ch_num in sorted(chapters_in_canto.keys()):
            ch_verses = [
                v for v in verses if v["canto"] == canto_num and v["chapter"] == ch_num
            ]
            if ch_verses:
                first = ch_verses[0]
                last = ch_verses[-1]
                title = chapter_titles.get((canto_num, ch_num), "")
                print(
                    f"  Ch {ch_num:>2}: {len(ch_verses):>3} verses  "
                    f"({first['reference']} → {last['reference']})  {title}"
                )

    # Quick quality check: show first verse details
    if verses:
        print("\n--- Sample (first verse) ---")
        v = verses[0]
        print(f"Reference: {v['reference']}")
        print(f"Devanagari:\n{v['devanagari']}")
        print(f"Transliteration:\n{v['transliteration']}")
        print(f"Synonyms: {len(v['synonyms'])} entries")
        if v["synonyms"]:
            print(f"  First: {v['synonyms'][0]}")
        print(f"Translation: {v['translation'][:120]}...")
        print(f"Purport: {len(v['purport'])} paragraphs")


if __name__ == "__main__":
    main()