"""
Ingestion OCR d'UN seul PDF (vision, GPU)
=========================================
Pour les PDF dont le texte est en vecteurs/images (non extractible) : on rend
chaque page en image, on lit le texte par OCR vision (EasyOCR, GPU si dispo),
puis on chunk + embed (Azure) + ajoute à la base ChromaDB.

Cible par défaut : le zonier MRH. Modifiable via argument :
    ./.venv/Scripts/python.exe ingest_ocr_single.py "MRH Train/xxx.pdf"

Prérequis :
    export AZURE_API_KEY="..."     # pour l'embedding
    pip install easyocr pymupdf    # déjà installé

Étapes :
  1) OCR -> écrit le texte dans _ocr_<nom>.txt (pour relecture)
  2) chunk + embed + ajout base (skip si déjà présent)
"""

import os
import sys
import json
import time
import logging
from pathlib import Path

os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")

import requests as http_requests
import fitz  # pymupdf
import chromadb
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# --- Config (identique au reste) ---
PROJECT_ROOT = Path(__file__).parent
DATA_DIR = Path("/data") if Path("/data").is_dir() else PROJECT_ROOT / "data"
TRAIN_DATA_DIR = Path("/train_data") if Path("/train_data").is_dir() else PROJECT_ROOT / "train_data"
COLLECTION_NAME = "rag_documents"
CHUNK_SIZE = 512
CHUNK_OVERLAP = 50
EMBEDDING_BATCH_SIZE = 16
REQUEST_TIMEOUT = 180
MAX_RETRIES = 5
OCR_DPI = 200  # qualité de rendu pour l'OCR

DEFAULT_TARGET = "MRH Train/Elaboration de zoniers en assurance MRH à partir de l'open data.pdf"

_CONFIG_PATH = DATA_DIR / "config.json"
if not _CONFIG_PATH.exists():
    _CONFIG_PATH = PROJECT_ROOT / "config.json"
with open(_CONFIG_PATH, encoding="utf-8") as _f:
    _config = json.load(_f)
EMBEDDING_ENDPOINT_URL = _config["embedding"]["endpoint_url"]
EMBEDDING_MODEL_NAME = _config["embedding"]["model"]


def ocr_pdf(pdf_path: Path) -> str:
    """Rend chaque page en image et lit le texte par OCR vision (EasyOCR)."""
    import easyocr
    import numpy as np

    # gpu=True -> utilise CUDA si dispo, sinon retombe sur CPU automatiquement
    try:
        import torch
        use_gpu = torch.cuda.is_available()
        logger.info(f"GPU CUDA disponible : {use_gpu}"
                    + (f" ({torch.cuda.get_device_name(0)})" if use_gpu else " -> OCR sur CPU"))
    except Exception:
        use_gpu = False
        logger.info("torch non détecté pour la vérif GPU -> EasyOCR choisira.")

    reader = easyocr.Reader(["fr"], gpu=use_gpu)
    doc = fitz.open(str(pdf_path))
    n = len(doc)
    pages_text = []
    for i in range(n):
        pix = doc[i].get_pixmap(dpi=OCR_DPI)
        img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
        if pix.n == 4:  # RGBA -> RGB
            img = img[:, :, :3]
        lines = reader.readtext(img, detail=0, paragraph=True)
        txt = "\n".join(lines).strip()
        if txt:
            pages_text.append(f"[Page {i+1}]\n{txt}")
        if (i + 1) % 10 == 0 or i + 1 == n:
            logger.info(f"  OCR page {i+1}/{n}")
    return "\n\n".join(pages_text)


def chunk_text(text: str, source: str) -> list[dict]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    return [{"text": c, "source": source, "chunk_index": i}
            for i, c in enumerate(splitter.split_text(text))]


def generate_embeddings(texts, api_key):
    headers = {"api-key": api_key, "Content-Type": "application/json"}
    payload = {"input": texts, "model": EMBEDDING_MODEL_NAME}
    last = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            r = http_requests.post(EMBEDDING_ENDPOINT_URL, headers=headers, json=payload, timeout=REQUEST_TIMEOUT)
            r.raise_for_status()
            return [d["embedding"] for d in r.json()["data"]]
        except Exception as e:
            last = e
            w = 2 ** attempt
            logger.warning(f"    embedding tentative {attempt}/{MAX_RETRIES} KO ({type(e).__name__}); retry {w}s")
            time.sleep(w)
    raise RuntimeError(f"Embedding échoué: {last}")


def main():
    rel = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_TARGET
    rel = rel.replace("/", os.sep)
    pdf_path = TRAIN_DATA_DIR / rel
    if not pdf_path.exists():
        logger.error(f"Introuvable: {pdf_path}")
        sys.exit(1)

    source = str(pdf_path.relative_to(TRAIN_DATA_DIR))
    logger.info(f"Cible: {source}")

    client = chromadb.PersistentClient(path=str(DATA_DIR / "chroma_db"), settings=Settings(anonymized_telemetry=False))
    col = client.get_or_create_collection(name=COLLECTION_NAME, metadata={"hnsw:space": "cosine"})

    # déjà présent ? (évite les doublons)
    existing = col.get(where={"source": source}, limit=1)
    if existing["ids"]:
        logger.warning(f"'{source}' est DÉJÀ dans la base. Rien à faire (relance après suppression si tu veux le refaire).")
        return

    # 1) OCR — réutilise le texte déjà OCRisé si présent (évite de refaire ~20 min d'OCR)
    out_txt = PROJECT_ROOT / ("_ocr_" + pdf_path.stem[:40] + ".txt")
    if out_txt.exists() and len(out_txt.read_text(encoding="utf-8").strip()) > 100:
        text = out_txt.read_text(encoding="utf-8")
        logger.info(f"=== ÉTAPE 1 : OCR déjà fait -> rechargé depuis {out_txt.name} ({len(text)} caractères) ===")
    else:
        logger.info("=== ÉTAPE 1 : OCR (vision) ===")
        text = ocr_pdf(pdf_path)
        out_txt.write_text(text, encoding="utf-8")
        logger.info(f"Texte OCR: {len(text)} caractères -> sauvegardé dans {out_txt.name} (relis-le si tu veux)")
    if len(text.strip()) < 100:
        logger.error("Très peu de texte OCR récupéré — on s'arrête. Vérifie le PDF / le moteur OCR.")
        return

    # 2) chunk + embed + add (besoin de la clé Azure ici)
    api_key = os.environ.get("AZURE_API_KEY")
    if not api_key:
        logger.error("AZURE_API_KEY non défini -> OCR fait (texte sauvé) mais embedding impossible. "
                     "Exporte la clé et relance pour finir l'ajout.")
        return

    logger.info("=== ÉTAPE 2 : chunk + embedding + ajout base ===")
    chunks = chunk_text(text, source=source)
    texts = [c["text"] for c in chunks]
    logger.info(f"{len(chunks)} chunks à embedder")

    embeddings = []
    for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
        batch = texts[i:i + EMBEDDING_BATCH_SIZE]
        logger.info(f"  batch {i // EMBEDDING_BATCH_SIZE + 1}/{(len(texts)-1)//EMBEDDING_BATCH_SIZE + 1}")
        embeddings.extend(generate_embeddings(batch, api_key))
        time.sleep(0.3)

    # IDs déterministes ET anti-collision : count() n'est PAS fiable après des suppressions
    # (le plus grand doc_N peut dépasser count()). On préfixe par la source.
    ids = [f"{source}#ocr#{i}" for i in range(len(chunks))]
    metas = [{"source": c["source"], "chunk_index": c["chunk_index"]} for c in chunks]
    col.add(ids=ids, embeddings=embeddings, documents=texts, metadatas=metas)
    logger.info("=" * 60)
    logger.info(f"AJOUTÉ : {len(chunks)} chunks pour '{source}'")
    logger.info(f"Total base : {col.count()} chunks")


if __name__ == "__main__":
    main()