#!/usr/bin/env python3
"""Pull Hebrew judgments from local parquet files into the tau-rag corpus.

Sources (auto-detected — first-found wins):
  • LawDBHeb/cases_clean.parquet         — 750k clean Hebrew judgments
  • LawDBHeb/cases_law14_waiver.parquet  — 10 judgments with rich metadata
  • TAU_API/.../uploads/cases_all.parquet  (alt)

Output:  tau_rag/runtime/parquet_cases.jsonl  (one JSON record per line)

The autoload hook in fastapi_app.py will pick this file up on startup,
classify the domain of each doc, and inject into the live pipeline so
they're searchable + structured-viewable just like the rest of the corpus.

Usage:
    python3 -m tau_rag.scripts.ingest_parquet_cases [--n 1000] [--source PATH]

Default: 1000 random samples from cases_clean + all 10 from waiver.
"""
from __future__ import annotations

import argparse
import json
import random
import re
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

import pyarrow.parquet as pq


# ============================================================================
# Locate parquet sources
# ============================================================================

# Auto-discovery: search common iCloud / Drive / local locations.
# The user's machine may have any of these, so we probe in order.
HOME = Path.home()
CANDIDATE_DIRS = [
    # macOS iCloud Drive
    HOME / "Library/Mobile Documents/com~apple~CloudDocs/LawDBHeb",
    HOME / "Library/Mobile Documents/com~apple~CloudDocs",
    # Sandbox path (CI / Cowork)
    Path("/sessions/kind-affectionate-mccarthy/mnt/com~apple~CloudDocs/LawDBHeb"),
    Path("/sessions/kind-affectionate-mccarthy/mnt/com~apple~CloudDocs"),
    # Other common Mac locations
    HOME / "Documents/LawDBHeb",
    HOME / "Downloads",
    HOME / "Documents",
    # Linux / generic
    Path("/data/LawDBHeb"),
    Path("/data"),
]

CLEAN_FILENAMES = ["cases_clean.parquet", "cases_all.parquet"]
RICH_FILENAMES = ["cases_law14_waiver.parquet"]


def _find_first(filenames):
    """Walk CANDIDATE_DIRS looking for any of the given filenames.
    Recurses one level into each so iCloud subfolders are caught."""
    for base in CANDIDATE_DIRS:
        if not base.exists():
            continue
        # Direct hit
        for f in filenames:
            p = base / f
            if p.exists():
                return p
        # Glob one level deep — catches LawDBHeb/cases_clean.parquet when
        # base is iCloudDocs root, etc.
        for f in filenames:
            try:
                hits = list(base.glob(f"**/{f}"))
                # Prefer paths under LawDBHeb/
                hits.sort(key=lambda p: ("LawDBHeb" not in str(p), len(str(p))))
                if hits:
                    return hits[0]
            except (OSError, PermissionError):
                continue
    return None


CANDIDATE_CLEAN = [_find_first(CLEAN_FILENAMES) or Path("cases_clean.parquet")]
CANDIDATE_RICH = [_find_first(RICH_FILENAMES) or Path("cases_law14_waiver.parquet")]

OUT_PATH = (Path(__file__).resolve().parent.parent / "runtime"
             / "parquet_cases.jsonl")


# ============================================================================
# Cleaning
# ============================================================================

# Strip HTML comments / Word CSS / inline style blocks that pollute many of
# the rows (especially in cases_law14_waiver). We're permissive: anything
# that looks structural and Hebrew-rare gets flattened.
HTML_NOISE = re.compile(
    r"<!--.*?-->"                       # HTML comments
    r"|<style[^>]*>.*?</style>"          # inline <style>
    r"|<[^>]+>"                          # any HTML tag
    r"|@font-face\s*\{[^}]*\}"           # CSS @font-face
    r"|/\*.*?\*/"                        # CSS comments
    r"|\bmso-[a-z\-]+:\s*[^;]+;"         # MS Office mso-* properties
    r"|\bpanose-1:\s*[^;\n]+;?",
    re.DOTALL | re.IGNORECASE,
)
WHITESPACE = re.compile(r"\n\s*\n\s*(\s*\n)+", re.MULTILINE)
LEADING_WS = re.compile(r"^[ \t]+", re.MULTILINE)
TRAILING_WS = re.compile(r"[ \t]+$", re.MULTILINE)


def clean_text(raw: str) -> str:
    """Strip HTML/CSS noise; collapse runs of empty lines; trim."""
    if not raw:
        return ""
    s = HTML_NOISE.sub("", raw)
    s = LEADING_WS.sub("", s)
    s = TRAILING_WS.sub("", s)
    s = WHITESPACE.sub("\n\n", s)
    return s.strip()


# ============================================================================
# Metadata extraction
# ============================================================================

CITATION_RE = re.compile(
    r'(בג"?ץ|ע"?א|ע"?פ|רע"?א|רע"?פ|דנ"?א|דנ"?פ|בש"?פ|בש"?א|ה"?פ)'
    r'\s*(\d+\s*[/\-]\s*\d+)'
)


def extract_citation(html_title: str) -> Optional[str]:
    if not html_title:
        return None
    m = CITATION_RE.search(html_title)
    if m:
        return f"{m.group(1)} {m.group(2).replace(' ', '')}"
    return None


def case_id_from_stable_id(stable_id: str) -> str:
    """cases_clean uses composite stable_ids; pull a stable hash-friendly suffix."""
    if not stable_id:
        return ""
    # Take the first segment that looks like a doc filename
    parts = stable_id.split("::")
    for p in parts:
        if "." in p and len(p) <= 60:
            return re.sub(r"[^A-Za-z0-9_\-.]", "_", p)
    return re.sub(r"[^A-Za-z0-9_\-.]", "_", parts[0])


def build_doc_from_clean(row) -> Optional[Dict[str, Any]]:
    """One row from cases_clean.parquet → {id, text, metadata} or None
    if the row should be skipped (empty text, etc)."""
    raw_text = row.get("text") or ""
    text = clean_text(raw_text)
    if len(text) < 200:
        return None
    sid = row.get("stable_id") or ""
    case_id = case_id_from_stable_id(sid)
    if not case_id:
        return None
    citation = extract_citation(row.get("html_title") or "")
    inyan = row.get("meta_inyan_nm") or ""
    return {
        "id": f"parquet/{case_id}",
        "text": text,
        "metadata": {
            "type": "caselaw",
            "source": "lawdbheb_clean",
            "title": row.get("html_title") or citation or case_id,
            "citation": citation,
            "court": _court_from_inyan(inyan),
            "language": "he",
            "stable_id": sid,
        },
    }


def build_doc_from_rich(row) -> Optional[Dict[str, Any]]:
    """One row from cases_law14_waiver.parquet (rich metadata)."""
    raw_text = row.get("text") or ""
    text = clean_text(raw_text)
    if len(text) < 200:
        return None
    case_id = str(row.get("CaseId") or row.get("CaseNum") or row.get("Id") or "")
    if not case_id:
        return None

    citation = (extract_citation(row.get("html_title") or "")
                 or extract_citation(row.get("CaseName") or ""))

    # judges + side names — flatten arrays
    judges = row.get("meta_judge")
    judges_list = list(judges) if judges is not None and len(judges) else []
    sides = row.get("meta_side_nm")
    sides_list = list(sides) if sides is not None and len(sides) else []

    verdict_dt = row.get("meta_verdict_dt") or row.get("VerdictDt")
    verdict_dt_str = str(verdict_dt)[:10] if verdict_dt else None

    return {
        "id": f"parquet_rich/{case_id}",
        "text": text,
        "metadata": {
            "type": "caselaw",
            "source": "lawdbheb_waiver",
            "title": row.get("CaseName") or row.get("html_title") or f"תיק {case_id}",
            "citation": citation,
            "court": row.get("meta_court_nm") or _court_from_inyan(row.get("meta_inyan_nm") or ""),
            "verdict_dt": verdict_dt_str,
            "judges": judges_list,
            "parties": sides_list,
            "verdict_type": row.get("meta_verdict_ty"),
            "case_num": row.get("CaseNum"),
            "language": "he",
        },
    }


def _court_from_inyan(inyan: str) -> Optional[str]:
    """Map the Hebrew 'inyan' (court-type) field to a full court name."""
    inyan = (inyan or "").strip()
    if not inyan:
        return None
    if inyan in ('בג"ץ', "בגץ"):
        return "בית המשפט העליון בשבתו כבית משפט גבוה לצדק"
    if inyan in ('ע"א', "עא") or inyan in ('רע"א', "רעא") or inyan in ('דנ"א', "דנא"):
        return "בית המשפט העליון"
    if inyan in ('ע"פ', "עפ") or inyan in ('רע"פ', "רעפ"):
        return "בית המשפט העליון"
    return None


# ============================================================================
# Domain classification — best-effort, runs only if classifier importable
# ============================================================================

def classify_domain(text: str) -> Optional[Dict[str, Any]]:
    try:
        # The script may be invoked as a module from anywhere. Make sure
        # the parent of tau_rag/ is on sys.path.
        repo_root = Path(__file__).resolve().parent.parent.parent
        if str(repo_root) not in sys.path:
            sys.path.insert(0, str(repo_root))
        from tau_rag.domain_classifier import classify
        res = classify(text)
        if res.get("top"):
            return {
                "domain": res["top"],
                "domain_scores": res.get("scores", {}),
            }
    except Exception:
        pass
    return None


# ============================================================================
# Main
# ============================================================================

def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
    ap.add_argument("--n", type=int, default=1000,
                     help="how many rows to sample from cases_clean (default 1000)")
    ap.add_argument("--clean-source", default=str(CANDIDATE_CLEAN[0]),
                     help="path to the bulk parquet file")
    ap.add_argument("--rich-source", default=str(CANDIDATE_RICH[0]),
                     help="path to the rich-metadata parquet file (10 rows)")
    ap.add_argument("--out", default=str(OUT_PATH),
                     help="output JSONL path")
    ap.add_argument("--seed", type=int, default=42, help="sample seed")
    ap.add_argument("--min-text-chars", type=int, default=300,
                     help="skip rows with fewer than N chars after cleaning")
    ap.add_argument("--append", action="store_true",
                     help="append to existing JSONL instead of overwriting")
    args = ap.parse_args()

    rng = random.Random(args.seed)
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # Track existing IDs when appending so we don't double-load the same doc
    existing_ids: set = set()
    if args.append and out_path.exists():
        with out_path.open("r", encoding="utf-8") as f:
            for line in f:
                try:
                    rec = json.loads(line)
                    existing_ids.add(rec.get("id"))
                except Exception:
                    pass
        print(f"[append] {len(existing_ids)} existing docs in {out_path.name}")

    n_written = 0
    n_skipped = 0
    n_classified = 0
    n_dup = 0
    open_mode = "a" if args.append else "w"

    with out_path.open(open_mode, encoding="utf-8") as fout:

        # ---- Source 1: cases_clean.parquet (bulk, simple metadata) ----
        clean_path = Path(args.clean_source)
        if clean_path.exists():
            print(f"[1/2] Reading {clean_path.name} ({args.n} samples)...",
                  flush=True)
            pf = pq.ParquetFile(str(clean_path))
            total = pf.metadata.num_rows
            # Sample row indices uniformly
            target_indices = sorted(rng.sample(range(total), min(args.n, total)))
            ti_set = set(target_indices)

            seen = 0
            written = 0
            # Iterate in batches to avoid loading 1.3GB at once
            for batch in pf.iter_batches(batch_size=4_000):
                df = batch.to_pandas()
                for i in range(len(df)):
                    if (seen + i) not in ti_set:
                        continue
                    row = df.iloc[i].to_dict()
                    doc = build_doc_from_clean(row)
                    if not doc:
                        n_skipped += 1
                        continue
                    if len(doc["text"]) < args.min_text_chars:
                        n_skipped += 1
                        continue
                    if doc["id"] in existing_ids:
                        n_dup += 1
                        continue
                    # Classify domain
                    cls = classify_domain(doc["text"])
                    if cls:
                        doc["metadata"].update(cls)
                        n_classified += 1
                    fout.write(json.dumps(doc, ensure_ascii=False) + "\n")
                    existing_ids.add(doc["id"])
                    written += 1
                    n_written += 1
                seen += len(df)
                if written >= args.n:
                    break
            print(f"      wrote {written} clean docs", flush=True)
        else:
            print(f"[skip] clean source not found: {clean_path}", flush=True)

        # ---- Source 2: cases_law14_waiver.parquet (rich metadata) ----
        rich_path = Path(args.rich_source)
        if rich_path.exists():
            print(f"[2/2] Reading {rich_path.name} (all rows, rich metadata)...",
                  flush=True)
            pf = pq.ParquetFile(str(rich_path))
            df = pf.read().to_pandas()
            written = 0
            for i in range(len(df)):
                row = df.iloc[i].to_dict()
                doc = build_doc_from_rich(row)
                if not doc:
                    n_skipped += 1
                    continue
                if len(doc["text"]) < args.min_text_chars:
                    n_skipped += 1
                    continue
                if doc["id"] in existing_ids:
                    n_dup += 1
                    continue
                cls = classify_domain(doc["text"])
                if cls:
                    doc["metadata"].update(cls)
                    n_classified += 1
                fout.write(json.dumps(doc, ensure_ascii=False) + "\n")
                existing_ids.add(doc["id"])
                written += 1
                n_written += 1
            print(f"      wrote {written} rich docs", flush=True)
        else:
            print(f"[skip] rich source not found: {rich_path}", flush=True)

    print()
    mode_msg = "Appended" if args.append else "Wrote"
    print(f"✅ {mode_msg} {n_written} docs to {out_path}")
    print(f"   ({n_classified} classified into a legal domain · "
           f"{n_skipped} skipped — too short / empty"
           + (f" · {n_dup} duplicates skipped" if args.append else "") + ")")
    if args.append:
        print(f"   total docs in file now: {len(existing_ids)}")
    print()
    print("Next: restart the server (lsof -ti:8000 | xargs kill -9; make run-local).")
    print("The autoload hook will pick up parquet_cases.jsonl automatically.")


if __name__ == "__main__":
    main()