extract_tipitaka.py

#!/usr/bin/env python3
"""
Tipitaka Data Extraction Pipeline
==================================
Extracts text, alignments, and dictionaries from the Tipitaka APK SQLite databases
into JSON files suitable for Neo4j graph DB + Firebase storage import.

Output structure:
  output/texts/          - One JSON per text layer with segments and character spans
  output/alignments/     - Segment-to-segment mappings (commentary ↔ root)
  output/dictionaries/   - Merged dictionary data (DPD, PTS, Myanmar)
  output/metadata/       - Master index of all texts and relationships
"""

import sqlite3
import json
import os
import re
import zlib
import html
from collections import defaultdict
from difflib import SequenceMatcher
import sys
import time

# ============================================================
# CONFIGURATION
# ============================================================

DB_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "db_extracted")
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")

# Collection code mapping
COLLECTION_MAP = {
    "11": {"collection": "vinaya", "pitaka": "vinaya", "layer": "mula", "name": "Vinaya Piṭaka"},
    "12": {"collection": "digha_nikaya", "pitaka": "sutta", "layer": "mula", "name": "Dīgha Nikāya"},
    "13": {"collection": "majjhima_nikaya", "pitaka": "sutta", "layer": "mula", "name": "Majjhima Nikāya"},
    "14": {"collection": "samyutta_nikaya", "pitaka": "sutta", "layer": "mula", "name": "Saṃyutta Nikāya"},
    "15": {"collection": "anguttara_nikaya", "pitaka": "sutta", "layer": "mula", "name": "Aṅguttara Nikāya"},
    "16": {"collection": "khuddaka_nikaya", "pitaka": "sutta", "layer": "mula", "name": "Khuddaka Nikāya"},
    "17": {"collection": "abhidhamma", "pitaka": "abhidhamma", "layer": "mula", "name": "Abhidhamma Piṭaka"},
    "21": {"collection": "vinaya", "pitaka": "vinaya", "layer": "atthakatha", "name": "Vinaya Aṭṭhakathā"},
    "22": {"collection": "digha_nikaya", "pitaka": "sutta", "layer": "atthakatha", "name": "DN Aṭṭhakathā"},
    "23": {"collection": "majjhima_nikaya", "pitaka": "sutta", "layer": "atthakatha", "name": "MN Aṭṭhakathā"},
    "24": {"collection": "samyutta_nikaya", "pitaka": "sutta", "layer": "atthakatha", "name": "SN Aṭṭhakathā"},
    "25": {"collection": "anguttara_nikaya", "pitaka": "sutta", "layer": "atthakatha", "name": "AN Aṭṭhakathā"},
    "26": {"collection": "khuddaka_nikaya", "pitaka": "sutta", "layer": "atthakatha", "name": "KN Aṭṭhakathā"},
    "27": {"collection": "abhidhamma", "pitaka": "abhidhamma", "layer": "atthakatha", "name": "Abhidhamma Aṭṭhakathā"},
    "31": {"collection": "vinaya", "pitaka": "vinaya", "layer": "tika", "name": "Vinaya Ṭīkā"},
    "32": {"collection": "digha_nikaya", "pitaka": "sutta", "layer": "tika", "name": "DN Ṭīkā"},
    "33": {"collection": "majjhima_nikaya", "pitaka": "sutta", "layer": "tika", "name": "MN Ṭīkā"},
    "34": {"collection": "samyutta_nikaya", "pitaka": "sutta", "layer": "tika", "name": "SN Ṭīkā"},
    "35": {"collection": "anguttara_nikaya", "pitaka": "sutta", "layer": "tika", "name": "AN Ṭīkā"},
    "36": {"collection": "khuddaka_nikaya", "pitaka": "sutta", "layer": "tika", "name": "KN Ṭīkā"},
    "37": {"collection": "abhidhamma", "pitaka": "abhidhamma", "layer": "tika", "name": "Abhidhamma Ṭīkā"},
    "48": {"collection": "other", "pitaka": "other", "layer": "other", "name": "Additional Texts"},
}

# Layer type suffixes in filenames
LAYER_SUFFIX_MAP = {
    "mul": "mula",
    "att": "atthakatha",
    "tik": "tika",
    "nrf": "other",
}


def ensure_dirs():
    """Create output directory structure."""
    for subdir in ["texts", "alignments", "dictionaries", "metadata"]:
        os.makedirs(os.path.join(OUTPUT_DIR, subdir), exist_ok=True)


def get_db(name):
    """Get a SQLite connection."""
    return sqlite3.connect(os.path.join(DB_DIR, name))


def natural_sort_key(path):
    """Sort key that handles k1, k2, ..., k10, k100 correctly."""
    match = re.search(r'k(\d+)$', path)
    if match:
        return int(match.group(1))
    return 0


def parse_fts_path(path):
    """Parse FTS path like '11@vin01m.mul0@k5' into components."""
    parts = path.split("@")
    if len(parts) != 3:
        return None
    prefix = parts[0]
    bookcode = parts[1]  # e.g., vin01m.mul0
    paragraph = parts[2]  # e.g., k5

    # Extract chapter number from bookcode
    # vin01m.mul0 → base=vin01m, chapter=0
    # vin02m1.mul5 → base=vin02m1, chapter=5
    dot_parts = bookcode.split(".")
    if len(dot_parts) != 2:
        return None

    base_name = dot_parts[0]  # e.g., vin01m or vin02m1
    type_with_chapter = dot_parts[1]  # e.g., mul0 or att5

    # Split type and chapter number
    match = re.match(r'([a-z]+)(\d+)$', type_with_chapter)
    if not match:
        return None

    layer_type = match.group(1)  # mul, att, tik, nrf
    chapter_num = int(match.group(2))

    # Derive text_id: the base book identifier without volume subdivisions
    # vin01m → text_id = vin01m
    # vin02m1 → text_id = vin02m (strip trailing volume number for grouping? No, keep as-is)
    text_id = base_name

    return {
        "prefix": prefix,
        "bookcode": bookcode,
        "text_id": text_id,
        "base_name": base_name,
        "layer_type": layer_type,
        "chapter": chapter_num,
        "paragraph": paragraph,
        "original_path": path,
    }


def strip_html(html_text):
    """Remove HTML tags and clean up text."""
    text = re.sub(r'<[^>]+>', '', html_text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text


# ============================================================
# STEP 1: EXTRACT TEXTS
# ============================================================

def extract_titles_from_cstpali():
    """Extract book titles from cstpali HTML <title> tags."""
    print("  Extracting titles from cstpali.db...")
    conn = get_db("cstpali.db")
    titles = {}

    rows = conn.execute(
        "SELECT filename, content FROM html_files WHERE filename LIKE 'book/%.html'"
    ).fetchall()

    for filename, content_blob in rows:
        try:
            html_content = zlib.decompress(content_blob).decode("utf-8")
            title_match = re.search(r'<title>(.*?)</title>', html_content)
            if title_match:
                title = html.unescape(title_match.group(1).strip())
                # filename like 'book/vin01m.mul.html'
                # Extract key: vin01m
                key = filename.replace("book/", "").replace(".html", "")
                # Remove .mul, .att, .tik, .nrf suffix
                key = re.sub(r'\.(mul|att|tik|nrf)(\.vol\d+)?$', '', key)
                titles[key] = title
        except Exception:
            pass

    conn.close()
    return titles


def extract_css_classes_from_cstpali():
    """Extract paragraph CSS classes from chapter HTML files."""
    print("  Extracting CSS classes from chapter HTML...")
    conn = get_db("cstpali.db")
    css_map = {}  # key: "bookcode@paragraph" → class

    rows = conn.execute(
        "SELECT filename, content FROM html_files WHERE filename LIKE 'chapter/%.html' AND filename NOT LIKE 'chapter/my.%'"
    ).fetchall()

    for filename, content_blob in rows:
        try:
            html_content = zlib.decompress(content_blob).decode("utf-8")
            # Extract chapter key from filename
            # chapter/vin01m.mul0.html → vin01m.mul0
            chapter_key = filename.replace("chapter/", "").replace(".html", "")

            # Find paragraphs with IDs and classes
            for match in re.finditer(
                r'<p\s+(?:class="([^"]*)")?\s*id="(k\d+)"', html_content
            ):
                css_class = match.group(1) or ""
                para_id = match.group(2)
                css_map[f"{chapter_key}@{para_id}"] = css_class

            # Also match id before class
            for match in re.finditer(
                r'<p[^>]*id="(k\d+)"[^>]*class="([^"]*)"', html_content
            ):
                para_id = match.group(1)
                css_class = match.group(2)
                css_map[f"{chapter_key}@{para_id}"] = css_class

        except Exception:
            pass

    conn.close()
    print(f"  Found CSS classes for {len(css_map)} paragraphs")
    return css_map


def extract_texts():
    """Extract all texts from FTS database into per-text JSON files."""
    print("\n=== STEP 1: Extracting texts ===")

    titles = extract_titles_from_cstpali()
    css_map = extract_css_classes_from_cstpali()

    print("  Loading FTS data...")
    conn = get_db("fts_tipitaka.db")
    rows = conn.execute("SELECT path, cont FROM pn").fetchall()
    conn.close()
    print(f"  Loaded {len(rows)} segments from FTS")

    # Group segments by text_id
    texts = defaultdict(list)
    skipped = 0

    for path, content in rows:
        parsed = parse_fts_path(path)
        if not parsed:
            skipped += 1
            continue
        texts[parsed["text_id"]].append({
            "parsed": parsed,
            "content": content,
        })

    if skipped:
        print(f"  Skipped {skipped} unparseable paths")

    print(f"  Found {len(texts)} unique text layers")

    # Process each text
    text_metadata = []

    for text_id in sorted(texts.keys()):
        segments_raw = texts[text_id]

        # Sort segments by chapter, then by paragraph number
        segments_raw.sort(
            key=lambda s: (s["parsed"]["chapter"], natural_sort_key(s["parsed"]["paragraph"]))
        )

        # Determine metadata from first segment's prefix
        first_prefix = segments_raw[0]["parsed"]["prefix"]
        collection_info = COLLECTION_MAP.get(first_prefix, {
            "collection": "unknown",
            "pitaka": "unknown",
            "layer": "unknown",
            "name": "Unknown",
        })

        layer_type = segments_raw[0]["parsed"]["layer_type"]
        layer = LAYER_SUFFIX_MAP.get(layer_type, layer_type)

        # Get title
        title_key = text_id
        title_pali = titles.get(title_key, "")
        title_breadcrumb = title_pali

        # Extract just the Pali title (last part after >)
        if " > " in title_pali:
            title_pali = title_pali.split(" > ")[-1].strip()

        # Build full_text and segments with spans
        full_text_parts = []
        segments = []
        current_offset = 0

        for seg in segments_raw:
            content = seg["content"]
            parsed = seg["parsed"]

            span_start = current_offset
            span_end = current_offset + len(content)

            # Look up CSS class
            css_key = f"{parsed['bookcode']}@{parsed['paragraph']}"
            css_class = css_map.get(css_key, "")

            segment_id = f"{text_id}_ch{parsed['chapter']}_{parsed['paragraph']}"

            segments.append({
                "id": segment_id,
                "chapter": parsed["chapter"],
                "paragraph": parsed["paragraph"],
                "span_start": span_start,
                "span_end": span_end,
                "content": content,
                "css_class": css_class,
                "original_path": parsed["original_path"],
            })

            full_text_parts.append(content)
            current_offset = span_end + 1  # +1 for newline separator

        full_text = "\n".join(full_text_parts)

        # Build text JSON
        text_json = {
            "id": text_id,
            "title_pali": title_pali,
            "title_breadcrumb": title_breadcrumb,
            "collection": collection_info["collection"],
            "pitaka": collection_info["pitaka"],
            "layer": layer,
            "layer_type": layer_type,
            "fts_prefix": first_prefix,
            "source_filename": f"book/{text_id.rstrip('0123456789') if re.match(r'.*[mat]\d+$', text_id) else text_id}.{layer_type}.html",
            "total_segments": len(segments),
            "total_characters": len(full_text),
            "full_text": full_text,
            "segments": segments,
        }

        # Write JSON file
        output_path = os.path.join(OUTPUT_DIR, "texts", f"{text_id}.json")
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(text_json, f, ensure_ascii=False, indent=2)

        text_metadata.append({
            "id": text_id,
            "title_pali": title_pali,
            "title_breadcrumb": title_breadcrumb,
            "collection": collection_info["collection"],
            "pitaka": collection_info["pitaka"],
            "layer": layer,
            "fts_prefix": first_prefix,
            "total_segments": len(segments),
            "total_characters": len(full_text),
        })

    print(f"  Wrote {len(text_metadata)} text JSON files to output/texts/")
    return text_metadata


# ============================================================
# STEP 2: EXTRACT ALIGNMENTS
# ============================================================

def extract_bold_quotes_from_chapter(conn, chapter_filename):
    """Extract bold quoted text and their paragraph IDs from a chapter HTML."""
    row = conn.execute(
        "SELECT content FROM html_files WHERE filename=?", (chapter_filename,)
    ).fetchone()

    if not row:
        return []

    try:
        html = zlib.decompress(row[0]).decode("utf-8")
    except Exception:
        return []

    quotes = []

    # Find each paragraph, then find bold quotes within it
    for para_match in re.finditer(
        r'<p[^>]*id="(k\d+)"[^>]*>(.*?)</p>', html, re.DOTALL
    ):
        para_id = para_match.group(1)
        para_content = para_match.group(2)

        # Find all bold quotes in this paragraph
        for bold_match in re.finditer(r'<b class="bld">(.*?)</b>', para_content, re.DOTALL):
            bold_html = bold_match.group(1)
            bold_text = strip_html(bold_html).strip()

            if len(bold_text) >= 3:  # Skip very short quotes
                quotes.append({
                    "paragraph": para_id,
                    "quoted_text": bold_text,
                })

    return quotes


def build_segment_index(root_segments):
    """Build a fast lookup index from root text for substring matching."""
    # Build concatenated text with segment boundary tracking
    full_text = ""
    boundaries = []  # (start_in_full, end_in_full, segment)

    for seg in root_segments:
        start = len(full_text)
        full_text += seg["content"].lower() + "\n"
        end = len(full_text) - 1  # before the newline
        boundaries.append((start, end, seg))

    return full_text, boundaries


def find_segment_at_position(pos, boundaries):
    """Binary search to find which segment contains position `pos`."""
    lo, hi = 0, len(boundaries) - 1
    while lo <= hi:
        mid = (lo + hi) // 2
        start, end, seg = boundaries[mid]
        if pos < start:
            hi = mid - 1
        elif pos > end:
            lo = mid + 1
        else:
            return seg
    return None


def find_matching_segment(quoted_text, root_full_text, root_boundaries):
    """Find the best matching root segment for a quoted text using fast substring search."""
    quoted_clean = quoted_text.lower().strip("''\"' ")

    if len(quoted_clean) < 3:
        return None

    # Fast exact substring match in the concatenated text
    pos = root_full_text.find(quoted_clean)
    if pos >= 0:
        seg = find_segment_at_position(pos, root_boundaries)
        if seg:
            # Calculate position within the segment
            seg_start_in_full = None
            for bstart, bend, bseg in root_boundaries:
                if bseg["id"] == seg["id"]:
                    seg_start_in_full = bstart
                    break

            local_pos = pos - (seg_start_in_full or 0)
            return {
                "segment_id": seg["id"],
                "span_start": seg["span_start"] + max(0, local_pos),
                "span_end": seg["span_start"] + max(0, local_pos) + len(quoted_clean),
                "confidence": 1.0,
                "method": "exact",
            }

    return None


def derive_root_text_id(commentary_text_id):
    """
    Derive the root text ID from a commentary text ID.
    vin01a → vin01m (commentary → root)
    vin01t1 → vin01a (tika → commentary) and vin01m (tika → root)
    s0101a → s0101m
    """
    # Replace the letter before any trailing digits
    # vin01a → replace 'a' with 'm'
    # vin01t1 → replace 't' with 'a' or 'm'

    # Pattern: base + single letter (a/t) + optional digit
    match = re.match(r'^(.*?)([at])(\d*)$', commentary_text_id)
    if not match:
        return None, None

    base = match.group(1)  # e.g., 'vin01'
    letter = match.group(2)  # 'a' or 't'
    suffix = match.group(3)  # '' or '1', '2'

    if letter == 'a':
        # Commentary → Root
        return f"{base}m", "mula_to_atthakatha"
    elif letter == 't':
        # Sub-commentary → Commentary
        return f"{base}a", "atthakatha_to_tika"

    return None, None


def extract_alignments(text_metadata):
    """Build alignment files by matching bold quotes in commentary to root text."""
    print("\n=== STEP 2: Extracting alignments ===")

    conn_cstpali = get_db("cstpali.db")

    # Get all chapter files available
    chapter_files = {}
    rows = conn_cstpali.execute(
        "SELECT filename FROM html_files WHERE filename LIKE 'chapter/%.html' AND filename NOT LIKE 'chapter/my.%'"
    ).fetchall()
    for (fname,) in rows:
        chapter_files[fname] = True

    # Load text IDs that are commentary or tika
    commentary_texts = [
        m for m in text_metadata
        if m["layer"] in ("atthakatha", "tika")
    ]

    print(f"  Processing {len(commentary_texts)} commentary/tika texts...")

    alignment_count = 0
    alignment_files_count = 0

    for ct in commentary_texts:
        commentary_id = ct["id"]
        root_id, alignment_type = derive_root_text_id(commentary_id)

        if not root_id:
            continue

        # Load root text segments
        root_path = os.path.join(OUTPUT_DIR, "texts", f"{root_id}.json")
        if not os.path.exists(root_path):
            continue

        with open(root_path, "r", encoding="utf-8") as f:
            root_data = json.load(f)

        root_segments = root_data["segments"]

        # Build fast search index for root text
        root_full_text, root_boundaries = build_segment_index(root_segments)

        # Load commentary text to get chapter info
        commentary_path = os.path.join(OUTPUT_DIR, "texts", f"{commentary_id}.json")
        if not os.path.exists(commentary_path):
            continue

        with open(commentary_path, "r", encoding="utf-8") as f:
            commentary_data = json.load(f)

        # Find all chapter files for this commentary
        # Derive chapter filenames from segments
        chapter_nums = set()
        for seg in commentary_data["segments"]:
            chapter_nums.add(seg["chapter"])

        # Get layer type from the commentary
        layer_type = commentary_data.get("layer_type", "att")

        all_alignments = []

        for ch_num in sorted(chapter_nums):
            # Build chapter filename
            chapter_fname = f"chapter/{commentary_id}.{layer_type}{ch_num}.html"

            if chapter_fname not in chapter_files:
                continue

            quotes = extract_bold_quotes_from_chapter(conn_cstpali, chapter_fname)

            # Get commentary segments for this chapter (to record target spans)
            commentary_chapter_segs = [
                s for s in commentary_data["segments"] if s["chapter"] == ch_num
            ]
            commentary_seg_map = {s["paragraph"]: s for s in commentary_chapter_segs}

            for quote in quotes:
                match = find_matching_segment(quote["quoted_text"], root_full_text, root_boundaries)

                if match:
                    target_seg = commentary_seg_map.get(quote["paragraph"])
                    if not target_seg:
                        continue

                    alignment = {
                        "id": f"align_{commentary_id}_{len(all_alignments):04d}",
                        "source_segment_id": match["segment_id"],
                        "target_segment_id": target_seg["id"],
                        "quoted_text": quote["quoted_text"],
                        "match_confidence": match["confidence"],
                        "match_method": match["method"],
                        "source_span_start": match["span_start"],
                        "source_span_end": match["span_end"],
                        "target_span_start": target_seg["span_start"],
                        "target_span_end": target_seg["span_end"],
                    }
                    all_alignments.append(alignment)

        if all_alignments:
            # Derive alignment filename
            # vin01a → vin01_mula_attha
            base = re.match(r'^(.*?)[amt]', commentary_id)
            base_name = base.group(1) if base else commentary_id
            align_suffix = "mula_attha" if alignment_type == "mula_to_atthakatha" else "attha_tika"

            alignment_json = {
                "source_text_id": root_id,
                "target_text_id": commentary_id,
                "alignment_type": alignment_type,
                "total_alignments": len(all_alignments),
                "alignments": all_alignments,
            }

            output_path = os.path.join(
                OUTPUT_DIR, "alignments", f"{base_name}_{align_suffix}.json"
            )
            with open(output_path, "w", encoding="utf-8") as f:
                json.dump(alignment_json, f, ensure_ascii=False, indent=2)

            alignment_count += len(all_alignments)
            alignment_files_count += 1

        # Progress
        if alignment_files_count % 20 == 0 and alignment_files_count > 0:
            print(f"    Processed {alignment_files_count} alignment files...")

    conn_cstpali.close()
    print(f"  Wrote {alignment_files_count} alignment files with {alignment_count} total alignments")
    return alignment_count


# ============================================================
# STEP 3: EXTRACT DICTIONARIES
# ============================================================

def get_all_sharded_tables(conn, exclude=("synonyms",)):
    """Get all word tables from a sharded dictionary DB (excluding system/special tables)."""
    tables = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite%'"
    ).fetchall()
    return [t[0] for t in tables if t[0] not in exclude]


def extract_dpd():
    """Extract Digital Pali Dictionary with parsed metadata."""
    print("  Extracting DPD (Digital Pali Dictionary)...")
    conn = get_db("dpd_tipitakapali.db")
    tables = get_all_sharded_tables(conn)

    entries = []
    for table in tables:
        rows = conn.execute(f'SELECT word, defi FROM "{table}"').fetchall()
        for word, defi in rows:
            entry = {
                "word": word,
                "definition_html": defi,
            }

            # Parse embedded JSON metadata: @{...}@
            json_match = re.search(r'@(\{.*?\})@', defi)
            if json_match:
                try:
                    meta = json.loads(json_match.group(1))
                    entry["dpd_id"] = meta.get("id")
                    entry["lemma"] = meta.get("lemma")
                    entry["family_sets"] = meta.get("family_sets", [])
                except json.JSONDecodeError:
                    pass

            # Extract POS and meaning from HTML
            # Try the dpd div first, then the bold meaning pattern
            meaning_match = re.search(r'<div class=dpd><p>(.*?)</p>', defi, re.DOTALL)
            if meaning_match:
                meaning_html = meaning_match.group(1)
                # Extract bold meaning specifically
                bold_meaning = re.search(r'<b>(.*?)</b>', meaning_html)
                if bold_meaning:
                    entry["meaning"] = strip_html(bold_meaning.group(1)).strip()

                meaning_text = strip_html(meaning_html)
                # POS is usually first word(s) like "masc." "nt." "verb." "adj."
                pos_match = re.match(r'^((?:\w+\.\s*(?:\(.*?\)\s*)?)+)', meaning_text)
                if pos_match:
                    entry["pos"] = pos_match.group(1).strip()

            # Extract Sanskrit cognate
            sanskrit_match = re.search(r'<th>Sanskrit<td><i>(.*?)</i>', defi)
            if sanskrit_match:
                entry["etymology_sanskrit"] = sanskrit_match.group(1)

            # Plain text definition
            entry["definition_plain"] = strip_html(defi)[:500]

            entries.append(entry)

    conn.close()

    output_path = os.path.join(OUTPUT_DIR, "dictionaries", "dpd.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

    print(f"    → {len(entries)} DPD entries")
    return len(entries)


def extract_pts():
    """Extract PTS Pali-English Dictionary."""
    print("  Extracting PTS Pali-English Dictionary...")
    conn = get_db("ptsped2015ed_tipitakapali.db")
    tables = get_all_sharded_tables(conn)

    entries = []
    for table in tables:
        if table == "synonyms":
            continue
        rows = conn.execute(f'SELECT word, defi FROM "{table}"').fetchall()
        for word, defi in rows:
            entry = {
                "word": word,
                "definition_html": defi,
                "definition_plain": strip_html(defi)[:500],
            }

            # Extract etymology [a + baddha]
            etym_match = re.search(r'\[([^\]]+)\]', strip_html(defi))
            if etym_match:
                entry["etymology"] = etym_match.group(1)

            entries.append(entry)

    conn.close()

    output_path = os.path.join(OUTPUT_DIR, "dictionaries", "pts.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

    print(f"    → {len(entries)} PTS entries")
    return len(entries)


def extract_abhidhan():
    """Extract Myanmar/Burmese Abhidhana Dictionary."""
    print("  Extracting Abhidhana (Myanmar) Dictionary...")
    conn = get_db("abhidhan_tipitakapali.db")
    tables = get_all_sharded_tables(conn)

    entries = []
    for table in tables:
        if table == "synonyms":
            continue
        rows = conn.execute(f'SELECT word, defi FROM "{table}"').fetchall()
        for word, defi in rows:
            entries.append({
                "word": word,
                "definition_html": defi,
                "definition_plain": strip_html(defi)[:500],
            })

    conn.close()

    output_path = os.path.join(OUTPUT_DIR, "dictionaries", "abhidhan.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

    print(f"    → {len(entries)} Abhidhana entries")
    return len(entries)


def extract_inflections():
    """Extract inflection/declension tables."""
    print("  Extracting inflections...")
    conn = get_db("dpd_inflection_tipitakapali.db")
    tables = get_all_sharded_tables(conn)

    entries = []
    for table in tables:
        rows = conn.execute(f'SELECT word, defi FROM "{table}"').fetchall()
        for word, defi in rows:
            entry = {
                "lemma": word,
                "inflection_html": defi,
            }

            # Extract pattern info
            heading_match = re.search(r"<p class='heading'>(.*?)</p>", defi)
            if heading_match:
                heading = strip_html(heading_match.group(1))
                entry["heading"] = heading

                # Extract gender and pattern
                gender_match = re.search(r'a (masc|fem|nt)', heading)
                if gender_match:
                    entry["gender"] = gender_match.group(1)

                pattern_match = re.search(r'like (\w+)', heading)
                if pattern_match:
                    entry["declension_pattern"] = pattern_match.group(1)

            entries.append(entry)

    conn.close()

    output_path = os.path.join(OUTPUT_DIR, "dictionaries", "inflections.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

    print(f"    → {len(entries)} inflection entries")
    return len(entries)


def extract_synonyms():
    """Extract synonym mappings."""
    print("  Extracting synonyms...")
    conn = get_db("dpd_synonyms_tipitakapali.db")
    tables = get_all_sharded_tables(conn)

    entries = []
    for table in tables:
        try:
            # Synonym tables have (word, synonym) columns
            rows = conn.execute(f'SELECT word, synonym FROM "{table}"').fetchall()
            for word, synonym in rows:
                entries.append({
                    "word": word,
                    "synonym": synonym,
                })
        except Exception:
            # Some tables might have different schema (defi column)
            try:
                rows = conn.execute(f'SELECT word, defi FROM "{table}"').fetchall()
                for word, defi in rows:
                    entries.append({
                        "word": word,
                        "synonym": defi,
                    })
            except Exception:
                pass

    conn.close()

    output_path = os.path.join(OUTPUT_DIR, "dictionaries", "synonyms.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

    print(f"    → {len(entries)} synonym mappings")
    return len(entries)


def extract_compound_splits():
    """Extract compound word splitting data."""
    print("  Extracting compound splits...")
    conn = get_db("dpd_splitter_tipitakapali.db")
    tables = get_all_sharded_tables(conn)

    entries = []
    for table in tables:
        if table == "synonyms":
            continue
        try:
            rows = conn.execute(f'SELECT word, defi FROM "{table}"').fetchall()
            for word, defi in rows:
                entry = {
                    "compound_word": word,
                }

                # Parse defi: "abhiramati+iti@abhiramatī+iti"
                if "@" in defi:
                    parts = defi.split("@")
                    entry["components"] = parts[0]
                    entry["alternative"] = parts[1] if len(parts) > 1 else None
                else:
                    entry["components"] = defi

                entries.append(entry)
        except Exception:
            pass

    conn.close()

    output_path = os.path.join(OUTPUT_DIR, "dictionaries", "compound_splits.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(entries, f, ensure_ascii=False, indent=2)

    print(f"    → {len(entries)} compound split entries")
    return len(entries)


def extract_dictionaries():
    """Extract all dictionary data."""
    print("\n=== STEP 3: Extracting dictionaries ===")

    stats = {}
    stats["dpd"] = extract_dpd()
    stats["pts"] = extract_pts()
    stats["abhidhan"] = extract_abhidhan()
    stats["inflections"] = extract_inflections()
    stats["synonyms"] = extract_synonyms()
    stats["compound_splits"] = extract_compound_splits()

    return stats


# ============================================================
# STEP 4: GENERATE METADATA
# ============================================================

def generate_metadata(text_metadata, alignment_count, dict_stats):
    """Generate master metadata/index file."""
    print("\n=== STEP 4: Generating metadata ===")

    # Group texts by collection and layer
    collections = defaultdict(lambda: {"mula": [], "atthakatha": [], "tika": [], "other": []})

    for tm in text_metadata:
        layer = tm["layer"]
        collection = tm["collection"]
        collections[collection][layer].append(tm["id"])

    # Build relationships (which root → commentary → tika)
    relationships = []
    for tm in text_metadata:
        if tm["layer"] in ("atthakatha", "tika"):
            root_id, rel_type = derive_root_text_id(tm["id"])
            if root_id:
                relationships.append({
                    "source": root_id,
                    "target": tm["id"],
                    "type": rel_type,
                })

    metadata = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "summary": {
            "total_texts": len(text_metadata),
            "total_alignments": alignment_count,
            "dictionary_entries": dict_stats,
        },
        "collection_codes": COLLECTION_MAP,
        "collections": dict(collections),
        "texts": text_metadata,
        "relationships": relationships,
    }

    output_path = os.path.join(OUTPUT_DIR, "metadata", "collections.json")
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(metadata, f, ensure_ascii=False, indent=2)

    print(f"  Wrote metadata with {len(relationships)} relationships")


# ============================================================
# MAIN
# ============================================================

def load_text_metadata_from_output():
    """Load text metadata from already-extracted text files."""
    texts_dir = os.path.join(OUTPUT_DIR, "texts")
    metadata = []
    for fname in sorted(os.listdir(texts_dir)):
        if not fname.endswith(".json"):
            continue
        fpath = os.path.join(texts_dir, fname)
        with open(fpath, "r", encoding="utf-8") as f:
            data = json.load(f)
        metadata.append({
            "id": data["id"],
            "title_pali": data.get("title_pali", ""),
            "title_breadcrumb": data.get("title_breadcrumb", ""),
            "collection": data["collection"],
            "pitaka": data["pitaka"],
            "layer": data["layer"],
            "fts_prefix": data["fts_prefix"],
            "total_segments": data["total_segments"],
            "total_characters": data["total_characters"],
        })
    return metadata


def main():
    start_time = time.time()
    print("=" * 60)
    print("Tipitaka Data Extraction Pipeline")
    print("=" * 60)

    ensure_dirs()

    # Check what steps to run
    step = sys.argv[1] if len(sys.argv) > 1 else "all"

    if step in ("all", "texts"):
        text_metadata = extract_texts()
    else:
        print("\n  Skipping text extraction (already done), loading metadata...")
        text_metadata = load_text_metadata_from_output()
        print(f"  Loaded metadata for {len(text_metadata)} texts")

    if step in ("all", "texts", "alignments"):
        alignment_count = extract_alignments(text_metadata)
    else:
        alignment_count = 0

    if step in ("all", "texts", "dictionaries"):
        dict_stats = extract_dictionaries()
    else:
        dict_stats = {}

    if step in ("all", "texts", "metadata"):
        generate_metadata(text_metadata, alignment_count, dict_stats)

    elapsed = time.time() - start_time
    print(f"\n{'=' * 60}")
    print(f"DONE in {elapsed:.1f}s")
    print(f"Output directory: {OUTPUT_DIR}")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()