chunking.py

"""
Parent-Child Hierarchical Chunker
=====================================================
Parent = full judgment text (used for context retrieval)
Child  = fixed-token sub-chunks of parent (used for embedding + search)

Reads:  data/html/extracted_judgments.jsonl   (metadata)
        data/pdfs/*.pdf                        (judgment text)
Writes: data/chunks/parent_child/<neutral_citation>.json

Architecture:
  Query hits a child chunk (small, precise, embedded)
      ↓
  Child carries parent_chunk_id
      ↓
  Fetch parent for full context window
      ↓
  Send parent text to LLM for answer generation
"""

import json
import os
import re
import fitz
import logging
import tiktoken
import concurrent.futures
from pathlib import Path
from datetime import datetime

# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("data/parent_child_chunker.log", encoding="utf-8"),
    ]
)
log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
METADATA_FILE    = os.path.join("data", "html", "extracted_judgments.jsonl")
PDF_DIR          = Path("data", "pdfs")
OUTPUT_DIR       = Path("data", "chunks", "parent_child")
ERROR_LOG        = Path("data", "chunks", "parent_child_errors.jsonl")

CHILD_CHUNK_SIZE    = 512    # tokens per child chunk
CHILD_CHUNK_OVERLAP = 100    # token overlap between children
TOKENIZER_MODEL     = "cl100k_base"
MAX_WORKERS         = min(8, (os.cpu_count() or 4))

# Parent size threshold — if judgment is smaller than this, skip children
MIN_TOKENS_FOR_CHILDREN = CHILD_CHUNK_SIZE

OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# ---------------------------------------------------------------------------
# Globals (shared across threads — all read-only after init)
# ---------------------------------------------------------------------------
TOKENIZER = tiktoken.get_encoding(TOKENIZER_MODEL)

# Pre-compiled regex
RE_BACKSPACE    = re.compile(r"\x08")
RE_AUTHOR       = re.compile(r"\n\*\s*Author\n\d+\n")
RE_PAGE_NUMS    = re.compile(r"\n\s*\d{1,3}\s*\n")
RE_FOOTER       = re.compile(r"\nJudgment\s*/\s*Order of the Supreme Court\n?", re.I)
RE_HEADER       = re.compile(r"\n(Supreme Court of India|IN THE SUPREME COURT OF INDIA)\n", re.I)
RE_NEWLINES     = re.compile(r"\n{3,}")
RE_SPACES       = re.compile(r"[ \t]+")


# ---------------------------------------------------------------------------
# Step 1: Load metadata index keyed by PDF filename
# ---------------------------------------------------------------------------
def load_metadata_index(jsonl_path: str) -> dict:
    """
    Returns:
    {
      "2026_INSC_479.pdf": { ...full metadata record... },
      ...
    }
    """
    index = {}
    missing = 0

    with open(jsonl_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError:
                continue

            pdf_path = record.get("pdf_path", "")
            if pdf_path:
                filename = Path(pdf_path).name        # "2026_INSC_479.pdf"
                index[filename] = record
            else:
                # Fallback: derive from neutral citation
                nc = record.get("neutral_citation", "").strip()
                if nc:
                    filename = nc.replace(" ", "_") + ".pdf"
                    index[filename] = record
                    missing += 1

    log.info(f"Loaded {len(index)} metadata records "
             f"({missing} used neutral_citation fallback).")
    return index


# ---------------------------------------------------------------------------
# Step 2: PDF text extraction
# ---------------------------------------------------------------------------
def extract_pdf_text(pdf_path: Path) -> str:
    """Extract full text from PDF using pymupdf."""
    parts = []
    try:
        doc = fitz.open(str(pdf_path))
        for page in doc:
            parts.append(page.get_text())
        doc.close()
    except Exception as e:
        log.error(f"PDF read failed [{pdf_path.name}]: {e}")
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# Step 3: Text cleaning
# ---------------------------------------------------------------------------
def clean_text(text: str) -> str:
    """Remove PDF artifacts, headers, footers, page numbers."""
    if not text:
        return ""
    text = RE_BACKSPACE.sub("", text)
    text = RE_AUTHOR.sub("\n", text)
    text = RE_PAGE_NUMS.sub("\n\n", text)
    text = RE_FOOTER.sub("\n", text)
    text = RE_HEADER.sub("\n", text)
    text = RE_NEWLINES.sub("\n\n", text)
    text = RE_SPACES.sub(" ", text)
    return text.strip()


# ---------------------------------------------------------------------------
# Step 4: Build lean metadata (for child chunks)
# ---------------------------------------------------------------------------
def build_lean_metadata(record: dict) -> dict:
    """
    Child chunks carry only the fields needed for Qdrant payload filtering.
    Full metadata lives on the parent — fetched at answer-generation time.
    """
    return {
        "case_name":        record.get("case_name", ""),
        "neutral_citation": record.get("neutral_citation", ""),
        "date":             record.get("date", ""),
        "court":            record.get("court", "Supreme Court"),
        "case_type":        record.get("case_type", ""),
        "outcome":          record.get("outcome", ""),
        "acts":             record.get("acts", []),
        "keywords":         record.get("keywords", []),
    }


# ---------------------------------------------------------------------------
# Step 5: Build full metadata (for parent chunk)
# ---------------------------------------------------------------------------
def build_full_metadata(record: dict) -> dict:
    return {
        "case_name":        record.get("case_name", ""),
        "neutral_citation": record.get("neutral_citation", ""),
        "appeal_no":        record.get("appeal_no", ""),
        "citation":         record.get("citation", ""),
        "date":             record.get("date", ""),
        "court":            record.get("court", "Supreme Court"),
        "lower_court":      record.get("lower_court", ""),
        "jurisdiction":     record.get("jurisdiction", "India"),
        "bench":            record.get("bench", []),
        "author_judge":     record.get("author_judge", ""),
        "outcome":          record.get("outcome", ""),
        "case_type":        record.get("case_type", ""),
        "acts":             record.get("acts", []),
        "sections":         record.get("sections", []),
        "cases_cited":      record.get("cases_cited", []),
        "keywords":         record.get("keywords", []),
        "issue":            record.get("issue", ""),
        "short_summary":    record.get("short_summary", ""),
        "full_headnote":    record.get("full_headnote", ""),
        "source_url":       record.get("source_url", ""),
        "scraped_at":       record.get("scraped_at", ""),
    }


# ---------------------------------------------------------------------------
# Step 6: Core parent-child chunking logic
# ---------------------------------------------------------------------------
def build_parent_child(
    neutral_citation: str,
    cleaned_text: str,
    full_metadata: dict,
    lean_metadata: dict,
) -> dict:
    """
    Returns:
    {
      "parent":   { single parent chunk with full text + full metadata },
      "children": [ child chunks with sub-text + lean metadata ]
    }
    """
    safe_nc      = neutral_citation.replace(" ", "_")
    parent_id    = f"{safe_nc}__parent"

    tokens       = TOKENIZER.encode(cleaned_text)
    num_tokens   = len(tokens)

    # --- Parent chunk ---
    parent = {
        "chunk_id":    parent_id,
        "chunk_type":  "parent",
        "document_id": neutral_citation,
        "text":        cleaned_text,
        "token_count": num_tokens,
        "char_count":  len(cleaned_text),
        "metadata":    full_metadata,
        "chunked_at":  datetime.utcnow().isoformat(),
    }

    # --- Skip children if text is too small ---
    if num_tokens <= MIN_TOKENS_FOR_CHILDREN:
        parent["child_count"] = 0
        return {"parent": parent, "children": []}

    # --- Child chunks ---
    step     = CHILD_CHUNK_SIZE - CHILD_CHUNK_OVERLAP
    children = []
    idx      = 0

    for start in range(0, num_tokens, step):
        end          = min(start + CHILD_CHUNK_SIZE, num_tokens)
        chunk_tokens = tokens[start:end]
        chunk_text   = TOKENIZER.decode(chunk_tokens)

        children.append({
            "chunk_id":       f"{safe_nc}__child_{idx:04d}",
            "chunk_type":     "child",
            "parent_chunk_id": parent_id,
            "document_id":    neutral_citation,
            "child_index":    idx,
            "token_count":    len(chunk_tokens),
            "char_count":     len(chunk_text),
            "start_token":    start,
            "end_token":      end,
            "text":           chunk_text,
            "metadata":       lean_metadata,
        })
        idx += 1

        if end == num_tokens:
            break

    parent["child_count"] = len(children)
    return {"parent": parent, "children": children}


# ---------------------------------------------------------------------------
# Step 7: Process single PDF (called by thread pool)
# ---------------------------------------------------------------------------
def process_pdf(pdf_path: Path, metadata_index: dict) -> dict:
    """
    Returns a result dict:
    {
      "status":  "success" | "skipped" | "error",
      "file":    pdf filename,
      "message": description,
      "chunks":  { parent, children } or None
    }
    """
    filename = pdf_path.name

    # --- Idempotency: skip if already processed ---
    record           = metadata_index.get(filename)
    if not record:
        return {"status": "unmatched", "file": filename,
                "message": f"No metadata found for {filename}"}

    neutral_citation = record.get("neutral_citation", "")
    safe_nc          = neutral_citation.replace(" ", "_")
    output_file      = OUTPUT_DIR / f"{safe_nc}.json"

    if output_file.exists():
        return {"status": "skipped", "file": filename,
                "message": f"Already processed: {output_file.name}"}

    # --- Extract + clean text ---
    raw_text     = extract_pdf_text(pdf_path)
    cleaned      = clean_text(raw_text)

    if not cleaned:
        return {"status": "error", "file": filename,
                "message": "Empty text after cleaning"}

    # --- Build metadata ---
    full_meta = build_full_metadata(record)
    lean_meta = build_lean_metadata(record)

    # --- Build parent-child structure ---
    result = build_parent_child(neutral_citation, cleaned, full_meta, lean_meta)

    # --- Write output atomically ---
    # Write to temp file first, then rename — prevents corrupt files on crash
    temp_file = output_file.with_suffix(".tmp")
    try:
        with open(temp_file, "w", encoding="utf-8") as f:
            json.dump(result, f, ensure_ascii=False, indent=2)
        temp_file.rename(output_file)
    except Exception as e:
        if temp_file.exists():
            temp_file.unlink()
        return {"status": "error", "file": filename, "message": str(e)}

    return {
        "status":   "success",
        "file":     filename,
        "message":  f"{result['parent']['child_count']} children created",
        "children": result["parent"]["child_count"],
    }


# ---------------------------------------------------------------------------
# Step 8: Main pipeline
# ---------------------------------------------------------------------------
def main():
    log.info("=" * 60)
    log.info("Parent-Child Chunker — Starting")
    log.info("=" * 60)

    # Load metadata
    metadata_index = load_metadata_index(METADATA_FILE)

    # Discover PDFs
    pdf_files = sorted(PDF_DIR.glob("*.pdf"))
    log.info(f"Found {len(pdf_files)} PDFs in {PDF_DIR}")

    if not pdf_files:
        log.error("No PDFs found. Check PDF_DIR path.")
        return

    # Counters
    counts = {"success": 0, "skipped": 0, "unmatched": 0, "error": 0}
    total_children = 0
    errors = []

    # Process concurrently
    log.info(f"Processing with {MAX_WORKERS} workers...")

    with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {
            executor.submit(process_pdf, pdf_path, metadata_index): pdf_path
            for pdf_path in pdf_files
        }

        for i, future in enumerate(concurrent.futures.as_completed(futures), 1):
            pdf_path = futures[future]
            try:
                result = future.result()
                status = result["status"]
                counts[status] = counts.get(status, 0) + 1

                if status == "success":
                    total_children += result.get("children", 0)
                    if i % 50 == 0:
                        log.info(
                            f"Progress: {i}/{len(pdf_files)} | "
                            f"Success: {counts['success']} | "
                            f"Skipped: {counts['skipped']} | "
                            f"Errors: {counts['error']}"
                        )
                elif status == "error":
                    log.warning(f"[ERROR] {result['file']}: {result['message']}")
                    errors.append(result)
                elif status == "unmatched":
                    log.warning(f"[UNMATCHED] {result['file']}")
                    errors.append(result)

            except Exception as exc:
                counts["error"] += 1
                log.error(f"[EXCEPTION] {pdf_path.name}: {exc}")
                errors.append({"file": pdf_path.name, "message": str(exc)})

    # Write error log
    if errors:
        with open(ERROR_LOG, "w", encoding="utf-8") as f:
            for e in errors:
                f.write(json.dumps(e, ensure_ascii=False) + "\n")
        log.info(f"Error details → {ERROR_LOG}")

    # Final summary
    log.info("=" * 60)
    log.info("PIPELINE COMPLETE")
    log.info(f"  Successful  : {counts['success']}")
    log.info(f"  Skipped     : {counts['skipped']} (already processed)")
    log.info(f"  Unmatched   : {counts['unmatched']} (no metadata)")
    log.info(f"  Errors      : {counts['error']}")
    log.info(f"  Total children created : {total_children}")
    log.info(f"  Output dir  : {OUTPUT_DIR.resolve()}")
    log.info("=" * 60)


if __name__ == "__main__":
    main()