"""
Ingest Training Data Script — Version parallélisée
=====================================================

Trois phases :
  1. Extraction + chunking de tous les fichiers  (ThreadPoolExecutor, I/O + CPU)
  2. Embedding via Azure OpenAI                  (ThreadPoolExecutor, I/O pur → gros gain)
  3. Insertion bulk dans ChromaDB                (1 seul appel par tranche de 5000)

Usage:
    python ingest_train_data.py              # ingestion normale (skip si déjà fait)
    python ingest_train_data.py --reset      # supprime la collection et recommence

Paramètres de parallélisme (ajuster selon les limites de l'API Azure) :
    EXTRACTION_WORKERS  : workers pour la lecture/parsing des PDFs
    EMBEDDING_WORKERS   : workers pour les appels HTTP à l'API d'embedding
    EMBEDDING_BATCH_SIZE: chunks par requête (max ~2048 pour Azure OpenAI)
"""

import argparse
import json
import logging
import os
import re
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import numpy as np

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

import requests as http_requests
import chromadb
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pypdf import PdfReader

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

# ──────────────────────────────────────────────────────────────────────────────
# Configuration test de la configuration
# ──────────────────────────────────────────────────────────────────────────────

PROJECT_ROOT   = Path(__file__).parent
DATA_DIR       = Path("/data_sementic")       if Path("/data_sementic").is_dir()       else PROJECT_ROOT / "data_sementic"
TRAIN_DATA_DIR = Path("/train_data") if Path("/train_data").is_dir() else PROJECT_ROOT / "train_data"

CHROMA_PERSIST_DIR = str(DATA_DIR / "chroma_db")
COLLECTION_NAME    = "rag_documents"

# ── Méthode de chunking ────────────────────────────────────────────────────────
# "semantic"   : coupe là où le sens change (recommandé, nécessite des appels embedding)
# "recursive"  : coupe par taille fixe (fallback, sans appel API supplémentaire)
CHUNKING_METHOD = "semantic"

# Paramètres du semantic chunking
SEMANTIC_BREAKPOINT_PERCENTILE = 95   # coupe où la distance dépasse le 95e percentile
SEMANTIC_BUFFER_SIZE           = 1    # phrases de contexte de chaque côté pour l'embedding
SEMANTIC_MIN_CHUNK_CHARS       = 150  # fusionne les chunks trop petits
SEMANTIC_MAX_CHUNK_CHARS       = 2000 # re-découpe les chunks trop grands

# Paramètres du chunking récursif (fallback ou méthode choisie)
CHUNK_SIZE    = 1000
CHUNK_OVERLAP = 150

# Parallélisme
EMBEDDING_BATCH_SIZE = 64   # chunks par requête API (était 16)
EMBEDDING_WORKERS    = 6    # appels API simultanés — réduire si rate-limit Azure
EXTRACTION_WORKERS   = 8    # workers pour l'extraction/parsing PDF

# Insertion ChromaDB par tranches pour éviter les problèmes mémoire
CHROMA_INSERT_BATCH  = 5_000

# Retry sur les appels embedding
EMBEDDING_MAX_RETRIES = 3
EMBEDDING_RETRY_DELAY = 2.0  # secondes entre deux tentatives

# ──────────────────────────────────────────────────────────────────────────────
# Chargement de la config
# ──────────────────────────────────────────────────────────────────────────────

_CONFIG_PATH = DATA_DIR / "config.json"
if not _CONFIG_PATH.exists():
    _CONFIG_PATH = PROJECT_ROOT / "config.json"
    logger.warning(f"Pas de config.json dans {DATA_DIR} — utilisation de 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"]

AZURE_API_KEY = os.environ.get("AZURE_API_KEY")
if not AZURE_API_KEY:
    logger.error("AZURE_API_KEY n'est pas défini. Définissez la variable d'environnement avant de lancer.")
    sys.exit(1)

# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────

def extract_text_from_pdf(pdf_path: Path) -> str:
    reader = PdfReader(str(pdf_path))
    pages_text = []
    for page_num, page in enumerate(reader.pages, start=1):
        text = page.extract_text()
        if text and text.strip():
            pages_text.append(f"[Page {page_num}]\n{text.strip()}")
    return "\n\n".join(pages_text)


def chunk_text(text: str, source: str) -> list[dict]:
    if CHUNKING_METHOD == "semantic":
        return _semantic_chunk(text, source)
    return _recursive_chunk(text, source)


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


def _split_sentences(text: str) -> list[str]:
    """Découpe le texte en phrases sur la ponctuation de fin."""
    parts = re.split(r'(?<=[.!?])\s+', text.strip())
    return [p.strip() for p in parts if len(p.strip()) > 10]


def _embed_in_batches(texts: list[str]) -> list[list[float]]:
    """Embedding par batches avec retry, compatible avec les gros documents."""
    all_embeddings: list = []
    for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
        batch = texts[i : i + EMBEDDING_BATCH_SIZE]
        for attempt in range(1, EMBEDDING_MAX_RETRIES + 1):
            try:
                all_embeddings.extend(generate_embeddings(batch))
                break
            except Exception as e:
                if attempt == EMBEDDING_MAX_RETRIES:
                    raise
                time.sleep(EMBEDDING_RETRY_DELAY)
    return all_embeddings


def _cosine_distance(a: list[float], b: list[float]) -> float:
    va, vb = np.array(a, dtype=np.float32), np.array(b, dtype=np.float32)
    return 1.0 - float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb) + 1e-10))


def _enforce_size_limits(chunks: list[str]) -> list[str]:
    """Fusionne les chunks trop courts, re-découpe les trop longs."""
    # Fusion des chunks sous le minimum
    merged: list[str] = []
    buf = ""
    for chunk in chunks:
        if buf and len(buf) + len(chunk) < SEMANTIC_MIN_CHUNK_CHARS:
            buf += " " + chunk
        else:
            if buf:
                merged.append(buf)
            buf = chunk
    if buf:
        merged.append(buf)

    # Re-découpage des chunks dépassant le maximum
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=SEMANTIC_MAX_CHUNK_CHARS,
        chunk_overlap=50,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    result: list[str] = []
    for chunk in merged:
        if len(chunk) > SEMANTIC_MAX_CHUNK_CHARS:
            result.extend(splitter.split_text(chunk))
        else:
            result.append(chunk)
    return result


def _semantic_chunk(text: str, source: str) -> list[dict]:
    """
    Semantic chunking :
    1. Découpe en phrases
    2. Crée des fenêtres de contexte (buffer de phrases adjacentes)
    3. Embed les fenêtres pour mesurer les ruptures sémantiques
    4. Coupe aux distances > percentile de seuil
    """
    sentences = _split_sentences(text)
    if len(sentences) <= 2:
        return _recursive_chunk(text, source)

    # Fenêtres : chaque phrase + BUFFER phrases de contexte de chaque côté
    windows = []
    for i in range(len(sentences)):
        start = max(0, i - SEMANTIC_BUFFER_SIZE)
        end   = min(len(sentences), i + SEMANTIC_BUFFER_SIZE + 1)
        windows.append(" ".join(sentences[start:end]))

    try:
        embeddings = _embed_in_batches(windows)
    except Exception as e:
        logger.warning(f"Semantic chunking échoué pour {source} → fallback récursif : {e}")
        return _recursive_chunk(text, source)

    # Distance cosine entre fenêtres consécutives
    distances = [
        _cosine_distance(embeddings[i], embeddings[i + 1])
        for i in range(len(embeddings) - 1)
    ]

    # Points de coupure au-dessus du seuil percentile
    threshold   = float(np.percentile(distances, SEMANTIC_BREAKPOINT_PERCENTILE))
    breakpoints = {i + 1 for i, d in enumerate(distances) if d > threshold}

    # Regroupement des phrases en chunks
    raw_chunks: list[str] = []
    current: list[str] = []
    for i, sentence in enumerate(sentences):
        if i in breakpoints and current:
            raw_chunks.append(" ".join(current))
            current = [sentence]
        else:
            current.append(sentence)
    if current:
        raw_chunks.append(" ".join(current))

    final_chunks = _enforce_size_limits(raw_chunks)
    return [{"text": c, "source": source, "chunk_index": i} for i, c in enumerate(final_chunks)]


def generate_embeddings(texts: list[str]) -> list[list[float]]:
    """Un seul appel HTTP — appelé en parallèle par plusieurs workers."""
    headers = {"api-key": AZURE_API_KEY, "Content-Type": "application/json"}
    payload = {"input": texts, "model": EMBEDDING_MODEL_NAME}
    resp = http_requests.post(EMBEDDING_ENDPOINT_URL, headers=headers, json=payload, timeout=120)
    resp.raise_for_status()
    return [item["embedding"] for item in resp.json()["data"]]


# ──────────────────────────────────────────────────────────────────────────────
# Workers (appelés depuis les ThreadPoolExecutors)
# ──────────────────────────────────────────────────────────────────────────────

def _process_file(file_path: Path) -> tuple[str, list[dict]]:
    """Phase 1 — Extrait et chunke un fichier. Retourne (source, chunks)."""
    source = str(file_path.relative_to(TRAIN_DATA_DIR))
    try:
        text = extract_text_from_pdf(file_path) if file_path.suffix.lower() == ".pdf" \
               else file_path.read_text(encoding="utf-8")
    except Exception as e:
        logger.warning(f"  Ignoré (erreur lecture) {source}: {e}")
        return source, []
    if not text.strip():
        logger.warning(f"  Ignoré (pas de texte) {source}")
        return source, []
    chunks = chunk_text(text, source=source)
    return source, chunks


def _embed_batch(batch_idx: int, texts: list[str]) -> tuple[int, list[list[float]]]:
    """Phase 2 — Embedde un batch avec retry. Retourne (batch_idx, embeddings)."""
    for attempt in range(1, EMBEDDING_MAX_RETRIES + 1):
        try:
            return batch_idx, generate_embeddings(texts)
        except Exception as e:
            if attempt == EMBEDDING_MAX_RETRIES:
                logger.error(f"  Batch {batch_idx} échoué après {EMBEDDING_MAX_RETRIES} tentatives : {e}")
                raise
            logger.warning(f"  Batch {batch_idx} — tentative {attempt} échouée ({e}), retry dans {EMBEDDING_RETRY_DELAY}s...")
            time.sleep(EMBEDDING_RETRY_DELAY)


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

def main():
    parser = argparse.ArgumentParser(description="Ingestion parallélisée des documents d'entraînement")
    parser.add_argument(
        "--reset", action="store_true",
        help="Supprime la collection ChromaDB existante avant de réingérer (nécessaire après changement de chunking)",
    )
    args = parser.parse_args()

    if not TRAIN_DATA_DIR.exists():
        logger.error(f"Dossier train_data introuvable : {TRAIN_DATA_DIR}")
        sys.exit(1)

    # ── ChromaDB ──────────────────────────────────────────────────────────────
    chroma_client = chromadb.PersistentClient(
        path=CHROMA_PERSIST_DIR,
        settings=Settings(anonymized_telemetry=False),
    )
    if args.reset:
        try:
            chroma_client.delete_collection(COLLECTION_NAME)
            logger.info("Collection supprimée (--reset).")
        except Exception:
            pass
    collection = chroma_client.get_or_create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )
    if collection.count() > 0 and not args.reset:
        logger.info(f"Collection déjà peuplée ({collection.count()} chunks). Utilisez --reset pour réingérer.")
        sys.exit(0)

    # ── Découverte des fichiers ───────────────────────────────────────────────
    files = sorted(list(TRAIN_DATA_DIR.rglob("*.pdf")) + list(TRAIN_DATA_DIR.rglob("*.txt")))
    if not files:
        logger.error(f"Aucun fichier PDF/TXT trouvé dans {TRAIN_DATA_DIR}")
        sys.exit(1)

    t_start = time.perf_counter()
    logger.info(f"{'='*60}")
    logger.info(f"Fichiers trouvés    : {len(files)}")
    logger.info(f"Méthode de chunking : {CHUNKING_METHOD.upper()}")
    if CHUNKING_METHOD == "semantic":
        logger.info(f"  Seuil percentile  : {SEMANTIC_BREAKPOINT_PERCENTILE}")
        logger.info(f"  Buffer phrases    : {SEMANTIC_BUFFER_SIZE}")
        logger.info(f"  Taille min/max    : {SEMANTIC_MIN_CHUNK_CHARS} / {SEMANTIC_MAX_CHUNK_CHARS} chars")
    else:
        logger.info(f"  Chunk size/overlap: {CHUNK_SIZE} / {CHUNK_OVERLAP}")
    logger.info(f"Batch size embedding: {EMBEDDING_BATCH_SIZE}")
    logger.info(f"Workers extraction  : {EXTRACTION_WORKERS}")
    logger.info(f"Workers embedding   : {EMBEDDING_WORKERS}")
    logger.info(f"{'='*60}")

    # ── Phase 1 : Extraction + chunking en parallèle ─────────────────────────
    logger.info(f"\nPhase 1/3 — Extraction et chunking ({EXTRACTION_WORKERS} workers)...")
    all_chunks: list[dict] = []
    errors_extraction = 0

    with ThreadPoolExecutor(max_workers=EXTRACTION_WORKERS) as executor:
        futures = {executor.submit(_process_file, f): f for f in files}
        done = 0
        for future in as_completed(futures):
            done += 1
            try:
                source, chunks = future.result()
                all_chunks.extend(chunks)
                logger.info(f"  [{done}/{len(files)}] {source} → {len(chunks)} chunks")
            except Exception as e:
                errors_extraction += 1
                logger.error(f"  [{done}/{len(files)}] Erreur inattendue : {e}")

    t_phase1 = time.perf_counter() - t_start
    logger.info(f"Phase 1 terminée en {t_phase1:.1f}s — {len(all_chunks)} chunks ({errors_extraction} erreurs)")

    if not all_chunks:
        logger.error("Aucun chunk produit. Vérifiez les fichiers dans train_data/.")
        sys.exit(1)

    # ── Phase 2 : Embedding en parallèle ─────────────────────────────────────
    logger.info(f"\nPhase 2/3 — Embedding ({len(all_chunks)} chunks, batch={EMBEDDING_BATCH_SIZE}, workers={EMBEDDING_WORKERS})...")

    texts  = [c["text"] for c in all_chunks]
    batches = [(i, texts[i : i + EMBEDDING_BATCH_SIZE]) for i in range(0, len(texts), EMBEDDING_BATCH_SIZE)]
    embeddings_map: dict[int, list] = {}
    errors_embedding = 0

    t_phase2_start = time.perf_counter()
    with ThreadPoolExecutor(max_workers=EMBEDDING_WORKERS) as executor:
        futures = {executor.submit(_embed_batch, idx, batch): idx for idx, batch in batches}
        done = 0
        for future in as_completed(futures):
            done += 1
            try:
                batch_idx, embeddings = future.result()
                embeddings_map[batch_idx] = embeddings
            except Exception as e:
                errors_embedding += 1
                logger.error(f"  Batch échoué, {errors_embedding} erreurs totales")
            if done % 20 == 0 or done == len(batches):
                elapsed = time.perf_counter() - t_phase2_start
                rate = done / elapsed if elapsed > 0 else 0
                eta = (len(batches) - done) / rate if rate > 0 else 0
                logger.info(f"  {done}/{len(batches)} batches — {elapsed:.0f}s écoulés, ETA ~{eta:.0f}s")

    t_phase2 = time.perf_counter() - t_phase2_start
    logger.info(f"Phase 2 terminée en {t_phase2:.1f}s ({errors_embedding} erreurs)")

    if errors_embedding > 0:
        logger.warning(f"{errors_embedding} batches ont échoué — les chunks correspondants seront absents de la base.")

    # Reconstruction dans l'ordre original
    flat_embeddings: list = []
    failed_indices: set[int] = set()
    valid_chunks: list[dict] = []

    for batch_start in range(0, len(texts), EMBEDDING_BATCH_SIZE):
        if batch_start in embeddings_map:
            flat_embeddings.extend(embeddings_map[batch_start])
            batch_end = min(batch_start + EMBEDDING_BATCH_SIZE, len(all_chunks))
            valid_chunks.extend(all_chunks[batch_start:batch_end])
        else:
            failed_indices.add(batch_start)

    # ── Phase 3 : Insertion bulk dans ChromaDB ───────────────────────────────
    logger.info(f"\nPhase 3/3 — Insertion dans ChromaDB ({len(valid_chunks)} chunks)...")
    t_phase3_start = time.perf_counter()

    for i in range(0, len(valid_chunks), CHROMA_INSERT_BATCH):
        batch_chunks     = valid_chunks[i : i + CHROMA_INSERT_BATCH]
        batch_embeddings = flat_embeddings[i : i + CHROMA_INSERT_BATCH]
        collection.add(
            ids        = [f"doc_{i + j}" for j in range(len(batch_chunks))],
            embeddings = batch_embeddings,
            documents  = [c["text"]         for c in batch_chunks],
            metadatas  = [{"source": c["source"], "chunk_index": c["chunk_index"]} for c in batch_chunks],
        )
        logger.info(f"  Inséré {min(i + CHROMA_INSERT_BATCH, len(valid_chunks))}/{len(valid_chunks)} chunks")

    t_phase3 = time.perf_counter() - t_phase3_start
    t_total   = time.perf_counter() - t_start

    logger.info(f"\n{'='*60}")
    logger.info(f"Ingestion terminée en {t_total:.1f}s ({t_total/60:.1f} min)")
    logger.info(f"  Phase 1 (extraction) : {t_phase1:.1f}s")
    logger.info(f"  Phase 2 (embedding)  : {t_phase2:.1f}s")
    logger.info(f"  Phase 3 (ChromaDB)   : {t_phase3:.1f}s")
    logger.info(f"Chunks dans la base    : {collection.count()}")
    logger.info(f"{'='*60}")


if __name__ == "__main__":
    main()