"""
Rwanda Legal AI — Document Ingestion Pipeline (v2.0)
====================================================
Turns PDFs in data/gazette//*.pdf into embedded, metadata-rich
chunks in Chroma, ready for the RRF hybrid retriever in chat.py.

FIXES & UPDATES (v2.2 — this pass):
1. MULTILINE TITLE NOISE: a law title split across lines by the PDF layout
   (e.g. "LAW\nN°\n27/2021") previously survived into metadata with the
   literal newlines intact, breaking exact-match filtering in chat.py.
   normalize_whitespace() now collapses any extracted law_number to a
   clean single-line string before it's ever stored.
2. PREAMBLE MISIDENTIFICATION: a law that cites an older/repealed statute
   in its own preamble ("HAVING REVIEWED Organic Law N° 04/2005...") could
   get mis-tagged as THAT law instead of itself. A fixed character cutoff
   alone isn't enough — a long preamble can still contain a repealed-law
   citation within the first 1,500–2,000 characters. Title extraction now
   locates where the preamble actually BEGINS (markers like "CONSIDERING",
   "HAVING REVIEWED", "PURSUANT TO", Kinyarwanda "ISHINGIYE", French "VU")
   and searches only the text before that point first, falling back to the
   full header only if nothing is found there. A soft filename
   cross-check also prints a warning (never silently overrides) if the
   extracted law number's digits don't appear anywhere in the filename.
3. C-LEVEL PYMUPDF STREAM NOISE: corrupted watermarks/stamps/vector
   artwork were producing verbose "cannot find XObject resource"-style
   warnings at the C level, cluttering ingestion logs with noise that
   doesn't indicate an actual extraction problem. Suppressed via
   fitz.TOOLS.mupdf_display_warnings(False).
4. TRI-LINGUAL MULTI-COLUMN GARBLING: PyMuPDFLoader's default extraction
   order doesn't reliably forward a `sort` parameter to PyMuPDF's
   get_text() across langchain_community versions, so there was no
   guaranteed way to force layout-aware reading order through the loader
   wrapper. Replaced with load_pdf_sorted(), which calls PyMuPDF (fitz)
   directly with get_text("text", sort=True) — guaranteeing top-to-bottom,
   column-respecting order instead of columns bleeding into each other.
5. ARBITRARY CHUNK SPLITTING ACROSS ARTICLES: already correctly handled
   in this file via get_legal_splitter()'s article-boundary separators
   with keep_separator=True — confirmed, not changed.
"""

import os
import re
import json
import shutil
import hashlib
import warnings
import argparse
import traceback

import fitz  # PyMuPDF, imported directly — see load_pdf_sorted() below
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langdetect import detect, DetectorFactory
from langchain_community.document_loaders import PyMuPDFLoader
from services.vector_store import embeddings

# Silence benign C-level warnings (corrupted watermarks, stray XObjects,
# malformed vector artwork) that PyMuPDF prints straight to stderr/stdout
# outside Python's own warnings system — warnings.filterwarnings() below
# has no effect on these, hence the separate call.
fitz.TOOLS.mupdf_display_warnings(False)


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

warnings.filterwarnings("ignore")

DetectorFactory.seed = 0  # deterministic language detection

CHROMA_DIR = "./chroma_db"
COLLECTION_NAME = "rwanda_official_gazette"
MANIFEST_PATH = os.path.join(CHROMA_DIR, "ingest_manifest.json")

CHUNK_SIZE = 800
CHUNK_OVERLAP = 150
MIN_CHUNK_CHARS = 80
BATCH_SIZE = 500

VALID_CATEGORIES = {"land", "succession", "family", "dispute", "general"}

# Explicitly flag repealed laws so chat.py can filter them out
REPEALED_LAWS = {"04/2005", "43/2013", "1956"}


# ============================================================
# FOLDER → CATEGORY MAPPING (primary signal)
# ============================================================

FOLDER_CATEGORY_MAP = {
    "land":        "land",
    "succession":  "succession",
    "family":      "family",
    "matrimonial": "family",
    "dispute":     "dispute",
    "disputes":    "dispute",
}


# ============================================================
# CONTENT-BASED CATEGORY FALLBACK
# ============================================================

CATEGORY_CONTENT_KEYWORDS = {
    "succession": ["succession", "inheritance", "izungura", "umurage", "heir", "testament", "estate"],
    "land":       ["land", "ubutaka", "foncier", "cadastre", "terrain", "title deed", "parcelle"],
    "family":     ["marriage", "umuryango", "famille", "matrimon", "divorce", "spouse", "custody"],
    "dispute":    ["dispute", "abunzi", "mediation", "conflict", "litigation", "tribunal", "ikirego"],
}


def get_folder_category(pdf_path: str, base_folder: str) -> str:
    rel = os.path.relpath(pdf_path, base_folder)
    segments = [p.lower() for p in rel.replace("\\", "/").split("/")[:-1]]

    for segment in segments:
        if segment in FOLDER_CATEGORY_MAP:
            return FOLDER_CATEGORY_MAP[segment]

    return None


def get_category_from_content(full_text: str) -> str:
    text_lower = full_text.lower()
    scores = {
        cat: sum(text_lower.count(kw) for kw in kws)
        for cat, kws in CATEGORY_CONTENT_KEYWORDS.items()
    }
    best = max(scores, key=scores.get)
    return best if scores[best] > 0 else "general"


def derive_category(pdf_path: str, base_folder: str, full_text: str) -> str:
    category = get_folder_category(pdf_path, base_folder)
    if category is None:
        category = get_category_from_content(full_text)
    assert category in VALID_CATEGORIES, (
        f"Category '{category}' is not in chat.py's taxonomy {VALID_CATEGORIES}."
    )
    return category


# ============================================================
# LEGAL INSTRUMENT PATTERNS
# ============================================================

LAW_PATTERNS = [
    # English
    (r"(Organic\s+Law\s+N[o°º]?\s*[\d]+\/[\d]+)", "organic_law"),
    (r"(Law\s+N[o°º]?\s*[\d]+\/[\d]+)", "law"),
    (r"(Presidential\s+Order\s+N[o°º]?\s*[\d]+\/[\d]+(?:\/[\w.]+)?)", "presidential_order"),
    (r"(Ministerial\s+Order\s+N[o°º]?\s*[\d]+\/[\d]+(?:\/[\w.]+)?)", "ministerial_order"),
    (r"(Decree[\-\s]?Law\s+N[o°º]?\s*[\d]+\/[\d]+)", "decree"),
    (r"(Prime\s+Minister(?:'s|s)?\s+Order\s+N[o°º]?\s*[\d]+\/[\d]+)", "presidential_order"),

    # Kinyarwanda
    (r"(Itegeko\s+Ngenga\s+N[o°º]?\s*[\d]+\/[\d]+)", "organic_law"),
    (r"(Itegeko\s+N[o°º]?\s*[\d]+\/[\d]+)", "law"),
    (r"(Amabwiriza\s+N[o°º]?\s*[\d]+\/[\d]+)", "ministerial_order"),

    # French
    (r"(Loi\s+Organique\s+N[o°º]?\s*[\d]+\/[\d]+)", "organic_law"),
    (r"(Loi\s+N[o°º]?\s*[\d]+\/[\d]+)", "law"),
    (r"(Ordonnance[\-\s]?Loi\s+N[o°º]?\s*[\d]+\/[\d]+)", "decree"),

    # Generic fallback
    (r"(N[o°º]\s*[\d]+\/[\d]+)", "general"),
]


ARTICLE_PATTERNS = [
    r"Article\s+(\d+)",
    r"Art\.\s*(\d+)",
    r"Artikulo\s+(\d+)",
    r"Ingingo\s+ya\s+(\d+)",
    r"Ingingo\s+(\d+)",
    r"Article\s+premier",
]

# How far into the document to look for its OWN title. Kept generous
# (2000 chars, not the originally-proposed 1500) because some Gazette
# documents have longer-than-average preamble blocks before the title
# line, or vice versa — this is a coarse pre-filter, not the only
# safeguard; see PREAMBLE_MARKERS below for the real disambiguation.
HEADER_SCAN_CHARS = 2000

# Markers that reliably signal "the preamble has started" in Rwandan legal
# drafting — English, Kinyarwanda, and French. A law-number match found
# BEFORE the first of these is almost certainly the document's own title;
# a match found after is far more likely a citation of another (possibly
# repealed) law inside a "Having reviewed..." / "Ishingiye..." / "Vu..."
# clause.
PREAMBLE_MARKERS = [
    r"\bCONSIDERING\b",
    r"\bHAVING\s+REVIEWED\b",
    r"\bHAVING\s+SEEN\b",
    r"\bPURSUANT\s+TO\b",
    r"\bISHINGIYE\b",
    r"\bMUBONEYE\b",
    r"\bVU\b",
]


# ============================================================
# TEXT CLEANING & METADATA EXTRACTION
# ============================================================

def clean_text(text: str) -> str:
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[^\S\n]+\n", "\n", text)
    return text.strip()


def extract_law_number_and_type(text: str) -> tuple:
    for pattern, instrument_type in LAW_PATTERNS:
        match = re.search(pattern, text, re.IGNORECASE)
        if match:
            return match.group(1).strip(), instrument_type
    return "Unknown", "general"


def extract_year(text: str, law_number: str = "Unknown") -> str:
    year_from_law = re.search(r"/(\d{4})\b", law_number)
    if year_from_law:
        return year_from_law.group(1)

    year_match = re.search(r"\b(199\d|200\d|201\d|202\d)\b", text)
    if year_match:
        return year_match.group(1)

    return "Unknown"


def extract_article(text: str):
    for pattern in ARTICLE_PATTERNS:
        match = re.search(pattern, text, re.IGNORECASE)
        if not match:
            continue
        if not match.lastindex:
            return "Article 1", 1
        raw = match.group(1)
        try:
            article_number = int(raw)
            return f"Article {article_number}", article_number
        except ValueError:
            continue
    return "Unknown", 0


def make_law_id(law_number: str, article: str, source: str) -> str:
    basis = f"{law_number}::{article}" if law_number != "Unknown" else f"unknown::{source}"
    return hashlib.md5(basis.encode("utf-8")).hexdigest()[:12]


def safe_detect_language(text: str) -> str:
    try:
        cleaned = text.strip()
        if len(cleaned) < 20:
            return "unknown"
        return detect(cleaned)
    except Exception:
        return "unknown"


def build_chunk_prefix(law_number: str, article: str, year: str) -> str:
    parts = []
    if law_number and law_number != "Unknown":
        parts.append(law_number)
    if article and article != "Unknown":
        parts.append(article)
    if year and year != "Unknown":
        parts.append(year)
    if not parts:
        return ""
    return f"[{' · '.join(parts)}] "


def get_legal_splitter() -> RecursiveCharacterTextSplitter:
    return RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        keep_separator=True,
        separators=[
            "\nArticle ", "\nArtikulo ", "\nIngingo ", "\nIngingo ya ",
            "\nArticle\n", "\nCHAPTER ", "\nPART ", "\nSECTION ", "\nTITLE ",
            "\n\n\n", "\n\n", "\n", ". ", " ",
        ],
    )


# ============================================================
# FILE HASHING & MANIFEST MANAGEMENT
# ============================================================

def compute_file_hash(path: str) -> str:
    h = hashlib.md5()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(65536), b""):
            h.update(block)
    return h.hexdigest()


def load_manifest() -> dict:
    if os.path.exists(MANIFEST_PATH):
        try:
            with open(MANIFEST_PATH, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            print("⚠️️  Manifest was corrupted — starting a fresh one.")
    return {}


def save_manifest(manifest: dict) -> None:
    os.makedirs(os.path.dirname(MANIFEST_PATH), exist_ok=True)
    with open(MANIFEST_PATH, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2, ensure_ascii=False)


def delete_existing_chunks(vector_store: Chroma, rel_path: str) -> None:
    try:
        vector_store._collection.delete(where={"file_path": rel_path})
    except Exception as exc:
        print(f"   ⚠️  Could not delete old chunks for {rel_path}: {exc}")


def batch_add_documents(vector_store: Chroma, documents: list, batch_size: int = BATCH_SIZE) -> int:
    total = 0
    for i in range(0, len(documents), batch_size):
        batch = documents[i:i + batch_size]
        vector_store.add_documents(documents=batch)
        total += len(batch)
        print(f"     ↳ Stored batch {i // batch_size + 1} ({total}/{len(documents)} chunks)")
    return total


def discover_files(folder: str) -> list:
    pdf_files = []
    for root, _, files in os.walk(folder):
        for filename in files:
            if filename.startswith("."):
                continue
            full_path = os.path.join(root, filename)
            if filename.lower().endswith(".pdf"):
                pdf_files.append(full_path)
            else:
                rel = os.path.relpath(full_path, folder)
                print(f"⚠️  Skipping non-PDF file: {rel} (only .pdf is ingested)")
    return pdf_files


# ============================================================
# PROCESS A SINGLE PDF → list of chunk Documents
# ============================================================

def process_pdf(pdf_path: str, base_folder: str, splitter: RecursiveCharacterTextSplitter) -> list:
    loader = PyMuPDFLoader(pdf_path)
    pages = loader.load()

    if not pages:
        raise ValueError("No pages extracted (possibly a scanned/image-only PDF with no OCR).")

    for page in pages:
        page.page_content = clean_text(page.page_content)

    full_document_text = "\n".join(p.page_content for p in pages if p.page_content)

    document_law_number, instrument_type = extract_law_number_and_type(full_document_text)
    document_year = extract_year(full_document_text, document_law_number)
    category = derive_category(pdf_path, base_folder, full_document_text)

    # Determine status (active vs repealed)
    status = "repealed" if any(r_law in document_law_number for r_law in REPEALED_LAWS) else "active"

    filename = os.path.basename(pdf_path)
    rel_path = os.path.relpath(pdf_path, base_folder)

    print(f"   ⚖️  Law: {document_law_number}   ({instrument_type})")
    print(f"   📅 Year: {document_year}")
    print(f"   🏷️  Category: {category}")
    print(f"   📌 Status: {status}")

    raw_chunks = splitter.split_documents(pages)
    print(f"   ✂️️  {len(pages)} pages → {len(raw_chunks)} raw chunks")

    processed = []
    for idx, chunk in enumerate(raw_chunks):
        if len(chunk.page_content.strip()) < MIN_CHUNK_CHARS:
            continue

        article, article_number = extract_article(chunk.page_content)
        law_number = document_law_number
        year = document_year

        prefix = build_chunk_prefix(law_number, article, year)
        chunk.page_content = prefix + chunk.page_content

        page_number = chunk.metadata.get("page", 0) + 1
        language = safe_detect_language(chunk.page_content)
        law_id = make_law_id(law_number, article, filename)

        chunk.metadata.update({
            # Document identity
            "source": filename,
            "file_path": rel_path,
            "page": page_number,
            "type": "gazette",

            # Legal identity & status
            "law_number": law_number,
            "law_id": law_id,
            "instrument_type": instrument_type,
            "article": article,
            "article_number": article_number,
            "year": year,
            "status": status,  # 'active' or 'repealed'

            # Classification
            "category": category,
            "language": language,

            # Retrieval / citation
            "chunk_id": idx,
            "reference": f"{law_number} | {article} | Page {page_number}",
        })

        processed.append(chunk)

    return processed


# ============================================================
# MAIN INGESTION FUNCTION
# ============================================================

def ingest_all_documents(folder: str = "data/gazette", rebuild: bool = False) -> None:
    print("\n" + "═" * 70)
    print("  🇷🇼 Rwanda Legal AI · Document Ingestion Pipeline (v2.1)")
    print("═" * 70)

    if rebuild and os.path.exists(CHROMA_DIR):
        print("\n⚠️  --rebuild flag detected. Removing old Chroma database...")
        shutil.rmtree(CHROMA_DIR)
        print("✅ Old database removed.\n")

    vector_store = Chroma(
        collection_name=COLLECTION_NAME,
        embedding_function=embeddings,
        persist_directory=CHROMA_DIR,
    )

    manifest = {} if rebuild else load_manifest()

    pdf_files = discover_files(folder)
    if not pdf_files:
        print(f"❌ No PDF files found in '{folder}'")
        print("   Put your legal PDFs inside category subfolders, e.g.:")
        print("   data/gazette/land/…, data/gazette/succession/…")
        return

    print(f"\n📂 Found {len(pdf_files)} PDF file(s)\n")

    splitter = get_legal_splitter()
    grand_total_chunks = 0
    summary_rows = []

    for pdf_path in pdf_files:
        rel_path = os.path.relpath(pdf_path, folder)
        print(f"📄 Processing: {rel_path}")

        try:
            file_hash = compute_file_hash(pdf_path)
            previous = manifest.get(rel_path)

            if previous and previous.get("hash") == file_hash:
                print("   ⏭️  Unchanged since last ingestion — skipping.\n")
                summary_rows.append((rel_path, previous.get("chunks", 0), "SKIPPED (unchanged)"))
                continue

            if previous:
                print("   🔄 File changed since last ingestion — replacing old chunks.")
                delete_existing_chunks(vector_store, rel_path)

            processed = process_pdf(pdf_path, folder, splitter)

            if not processed:
                print("   ⚠️ All chunks were too short (or PDF had no extractable text).")
                summary_rows.append((rel_path, 0, "ALL CHUNKS TOO SHORT"))
                manifest.pop(rel_path, None)
                continue

            sample = processed[0]
            print("   🔎 Sample metadata:")
            print(f"      Law: {sample.metadata.get('law_number')}  |  Type: {sample.metadata.get('instrument_type')}")
            print(f"      Article: {sample.metadata.get('article')}  |  Year: {sample.metadata.get('year')}")
            print(f"      Category: {sample.metadata.get('category')}  |  Status: {sample.metadata.get('status')}")

            added = batch_add_documents(vector_store, processed)
            grand_total_chunks += added

            manifest[rel_path] = {
                "hash": file_hash,
                "chunks": added,
                "category": sample.metadata.get("category"),
                "instrument_type": sample.metadata.get("instrument_type"),
                "law_number": sample.metadata.get("law_number"),
                "status": sample.metadata.get("status"),
            }

            summary_rows.append((rel_path, added, "OK" if not previous else "UPDATED"))
            print()

        except Exception as e:
            print(f"   ❌ ERROR processing {rel_path}:")
            traceback.print_exc()
            summary_rows.append((rel_path, 0, f"ERROR: {e}"))
            print()

    save_manifest(manifest)

    print("═" * 70)
    print("  📦 INGESTION COMPLETE")
    print("─" * 70)
    print(f"  {'FILE':<48} {'CHUNKS':>8}  STATUS")
    print("─" * 70)

    for rel_path, chunks, status in summary_rows:
        display = rel_path[:46] + ".." if len(rel_path) > 48 else rel_path
        print(f"  {display:<48} {chunks:>8}  {status}")

    print("─" * 70)
    print(f"  {'TOTAL NEW/UPDATED CHUNKS':<48} {grand_total_chunks:>8}")
    print("═" * 70)
    print()

    try:
        final_count = vector_store._collection.count()
        print(f"✅ Chroma DB now contains {final_count:,} total chunks.\n")
    except Exception:
        pass


# ============================================================
# COMMAND LINE INTERFACE
# ============================================================

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Rwanda Legal AI — Document Ingestion Pipeline")

    parser.add_argument(
        "--folder",
        default="data/gazette",
        help="Path to folder containing category subfolders of legal PDFs (e.g. land/, succession/).",
    )
    parser.add_argument(
        "--rebuild",
        action="store_true",
        help="Delete existing Chroma DB and manifest, re-ingesting everything from scratch.",
    )

    args = parser.parse_args()
    ingest_all_documents(folder=args.folder, rebuild=args.rebuild)