"""
embed_fallout.py — Ingestion PDFs Fallout JdR → ChromaDB (v3)

Fusion :
  - Chunking structurel par paragraphes (fallout_ingest_v2, 15/06)
  - Routing multi-collections selon README priorités (16/06)

Collections :
  fallout_lore_canon    ← lore_post_guerre/       (canonique, priorité max)
  fallout_regles        ← regles/core|supplements|fiches_de_jeu/
  fallout_lore_contexte ← lore_pre_guerre/ + aventures/ + ambiance/radio/

Chemins container : PDF_DIR=/pdfs/jdr  STATE=/pdfs/fallout_embed_state.json
Env vars : CHROMA_URL, OLLAMA_URL
"""
import json, os, re, subprocess, time, urllib.request
from pathlib import Path

CHROMA_URL  = os.getenv("CHROMA_URL", "http://chromadb:8000")
CHROMA_BASE = CHROMA_URL + "/api/v2/tenants/default_tenant/databases/default_database"
OLLAMA_URL  = os.getenv("OLLAMA_URL", "http://172.18.0.1:11434")
EMBED_MODEL = "nomic-embed-text"

PDF_DIR    = Path(os.getenv("PDF_DIR", "/pdfs/jdr"))
STATE_FILE = Path(os.getenv("STATE_FILE", "/pdfs/fallout_embed_state.json"))
BATCH_SIZE = 50

ROUTING = [
    ("lore_post_guerre",     "fallout_lore_canon",    "lore_post_guerre",  450, 80, True),
    ("regles/core",          "fallout_regles",         "regles_core",       350, 60, False),
    ("regles/supplements",   "fallout_regles",         "regles_supplement", 350, 60, False),
    ("regles/fiches_de_jeu", "fallout_regles",         "fiches_jeu",        350, 60, False),
    ("lore_pre_guerre",      "fallout_lore_contexte",  "lore_pre_guerre",   450, 80, False),
    ("aventures",            "fallout_lore_contexte",  "aventures",         450, 80, False),
    ("ambiance/radio",       "fallout_lore_contexte",  "ambiance_radio",    200, 40, False),
]

SKIP_DIRS = {
    "assets_graphiques", "battlemaps", "fonts", "photoshop_psd",
    "fonds_de_page", "images_ref", "croquis", "maps", "logos",
    "couvertures", ".claude", "archive",
}
SKIP_SLUGS = {"archive__meta_data__compressed"}
MIN_WORDS  = 40


# ── State ─────────────────────────────────────────────────────────────────────

def load_state():
    try:
        return json.loads(STATE_FILE.read_text())
    except Exception:
        return {"done": [], "failed": []}

def save_state(s):
    STATE_FILE.write_text(json.dumps(s, indent=2))


# ── ChromaDB ──────────────────────────────────────────────────────────────────

def _http(method, url, data=None, timeout=30):
    body = json.dumps(data).encode() if data is not None else None
    req = urllib.request.Request(url, data=body, method=method,
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())

def get_or_create_collection(name):
    try:
        r = _http("POST", f"{CHROMA_BASE}/collections",
                  {"name": name, "metadata": {"hnsw:space": "cosine"}})
        return r["id"]
    except urllib.error.HTTPError as e:
        if e.code == 409:
            cols = _http("GET", f"{CHROMA_BASE}/collections")
            return next(c["id"] for c in cols if c["name"] == name)
        raise

def upsert_batch(col_id, ids, embeddings, documents, metadatas):
    _http("POST", f"{CHROMA_BASE}/collections/{col_id}/upsert", {
        "ids": ids, "embeddings": embeddings,
        "documents": documents, "metadatas": metadatas,
    }, timeout=60)


# ── Embeddings ────────────────────────────────────────────────────────────────

def embed(text):
    safe = " ".join(text.split()[:280])
    r = _http("POST", f"{OLLAMA_URL}/api/embeddings",
              {"model": EMBED_MODEL, "prompt": safe}, timeout=60)
    return r["embedding"]


# ── PDF extraction ────────────────────────────────────────────────────────────

def extract_text(pdf_path):
    """pdftotext -layout en priorité, fallback pypdf, fallback fitz."""
    text = ""
    try:
        r = subprocess.run(
            ["pdftotext", "-layout", str(pdf_path), "-"],
            capture_output=True, timeout=180,
        )
        if r.returncode == 0:
            text = r.stdout.decode("utf-8", errors="replace")
    except (subprocess.TimeoutExpired, FileNotFoundError):
        pass

    if len(text.strip()) < 200:
        try:
            from pypdf import PdfReader
            reader = PdfReader(str(pdf_path))
            text = "\n".join(p.extract_text() or "" for p in reader.pages)
        except Exception:
            try:
                import fitz
                doc = fitz.open(str(pdf_path))
                text = "\n".join(p.get_text() for p in doc)
                doc.close()
            except Exception as e:
                raise RuntimeError(f"Aucune méthode d'extraction n'a fonctionné : {e}")

    return text.strip()


# ── Chunking structurel ───────────────────────────────────────────────────────

def is_section_header(line):
    l = line.strip()
    if not l or len(l) > 80:
        return False
    if l.isupper() and len(l.split()) <= 8:
        return True
    if re.match(r'^(CHAPITRE|CHAPTER|PARTIE|PART|SECTION|\d+[\.\)])\s+\S', l, re.I):
        return True
    return False

def chunk_structural(text, chunk_words=450, overlap_words=80):
    raw_blocks = re.split(r'\n{2,}', text)
    blocks = []
    for b in raw_blocks:
        b = b.strip()
        if not b:
            continue
        if len(b.split()) > 800:
            sub = [s.strip() for s in b.split('\n') if s.strip()]
            blocks.extend(sub)
        else:
            blocks.append(b)

    chunks = []
    current_blocks, current_words = [], 0

    for block in blocks:
        bwords = len(block.split())
        if bwords < MIN_WORDS and not is_section_header(block):
            current_blocks.append(block)
            current_words += bwords
            continue

        if current_words >= chunk_words:
            chunk = "\n\n".join(current_blocks).strip()
            if len(chunk.split()) >= MIN_WORDS:
                chunks.append(chunk)
            overlap_b, overlap_w = [], 0
            for ob in reversed(current_blocks):
                ow = len(ob.split())
                if overlap_w + ow <= overlap_words:
                    overlap_b.insert(0, ob)
                    overlap_w += ow
                else:
                    break
            current_blocks = overlap_b + [block]
            current_words = overlap_w + bwords
        else:
            current_blocks.append(block)
            current_words += bwords

    if current_blocks:
        chunk = "\n\n".join(current_blocks).strip()
        if len(chunk.split()) >= MIN_WORDS:
            chunks.append(chunk)

    return chunks


# ── Routing & métadonnées ─────────────────────────────────────────────────────

def route_pdf(rel_path):
    rel = rel_path.replace("\\", "/")
    for part in rel.split("/")[:-1]:
        if part.lower() in SKIP_DIRS:
            return None
    rel_low = rel.lower()
    for prefix, col, cat, cw, ov, canon in ROUTING:
        if rel_low.startswith(prefix.lower()):
            return col, cat, cw, ov, canon
    return None

def detect_langue(filename):
    n = filename.lower()
    if re.search(r'(_fr[_\-\.])|(_fr$)|(manuel.*officiel)|(bestiaire)|(armes)|'
                 r'(special.*comp)|(regles_maison)|(frequence)|(pizza)', n):
        return "fr"
    if re.search(r'(_en[_\-\.])|(_en$)|(quickstart)|(booklet)|(toolkit)|(settlers)|'
                 r'(enclave)|(rust_devil)|(winter_of_atom)|(cascadia)|(hunted)|'
                 r'(showdown)|(fully_operational)|(last_boat)', n):
        return "en"
    return "fr"

def detect_version(filename):
    m = re.search(r'v(\d+[\.\d]*)', filename.lower())
    return m.group(1) if m else ""


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    print("=== embed_fallout v3 — chunking structurel + routing multi-collections ===\n")

    all_pdfs = []
    for root, dirs, files in os.walk(PDF_DIR):
        dirs[:] = [d for d in dirs if d.lower() not in SKIP_DIRS]
        for f in sorted(files):
            if f.lower().endswith(".pdf"):
                all_pdfs.append(Path(root) / f)
    all_pdfs.sort()

    to_process, skipped = [], []
    for pdf_path in all_pdfs:
        rel  = str(pdf_path.relative_to(PDF_DIR))
        slug = rel.replace(os.sep, "__").replace(" ", "_").lower()[:-4]
        r    = route_pdf(rel)
        if r is None or slug in SKIP_SLUGS:
            skipped.append(rel)
            continue
        col, cat, cw, ov, canon = r
        to_process.append({
            "path": pdf_path, "rel": rel, "slug": slug,
            "col": col, "cw": cw, "ov": ov,
            "meta": {
                "source":    pdf_path.name,
                "categorie": cat,
                "langue":    detect_langue(pdf_path.name),
                "canonique": "true" if canon else "false",
                "version":   detect_version(pdf_path.name),
            },
        })

    print(f"PDFs trouvés : {len(all_pdfs)}")
    print(f"  A ingérer  : {len(to_process)}")
    print(f"  Ignorés    : {len(skipped)}")
    for s in skipped:
        print(f"    [SKIP] {s}")

    col_ids = {}
    print()
    for name in sorted(set(p["col"] for p in to_process)):
        col_ids[name] = get_or_create_collection(name)
        print(f"  Collection '{name}' : {col_ids[name][:8]}...")

    state = load_state()
    # Reset state pour repartir propre (collections wipées)
    state = {"done": [], "failed": []}
    save_state(state)

    total_chunks = 0

    for item in to_process:
        slug = item["slug"]
        if slug in state["done"]:
            print(f"[résumé] {item['rel']}")
            continue

        print(f"\n[->] {item['rel']}")
        print(f"     col={item['col']}  cat={item['meta']['categorie']}"
              f"  lang={item['meta']['langue']}  canon={item['meta']['canonique']}")

        try:
            text = extract_text(item["path"])
            print(f"     {len(text)} chars extraits")

            chunks = chunk_structural(text, item["cw"], item["ov"])
            print(f"     {len(chunks)} chunks (~{item['cw']}w cible, overlap {item['ov']}w)")

            ids, embeddings, documents, metadatas = [], [], [], []
            for i, chunk in enumerate(chunks):
                t0 = time.time()
                emb = embed(chunk)
                wc  = len(chunk.split())
                ids.append(f"{slug}__{i:04d}")
                embeddings.append(emb)
                documents.append(chunk)
                metadatas.append({**item["meta"],
                                  "chunk": i, "total_chunks": len(chunks),
                                  "word_count": wc})
                print(f"    chunk {i+1}/{len(chunks)} ({wc}w) {time.time()-t0:.1f}s")

                if len(ids) >= BATCH_SIZE:
                    upsert_batch(col_ids[item["col"]], ids, embeddings, documents, metadatas)
                    ids, embeddings, documents, metadatas = [], [], [], []

            if ids:
                upsert_batch(col_ids[item["col"]], ids, embeddings, documents, metadatas)

            total_chunks += len(chunks)
            state["done"].append(slug)
            state["failed"] = [f for f in state["failed"] if f.get("slug") != slug]
            print(f"  OK — {len(chunks)} chunks")

        except Exception as e:
            state["failed"].append({"slug": slug, "error": str(e)})
            print(f"  ERREUR : {e}")

        save_state(state)

    print(f"\n{'='*60}")
    print(f"TERMINÉ — {len(state['done'])}/{len(to_process)} PDFs")
    print(f"Chunks totaux : {total_chunks}")
    if state.get("failed"):
        print("Échecs :")
        for f in state["failed"]:
            print(f"  - {f['slug']}: {f['error']}")


if __name__ == "__main__":
    main()