app.py

"""
RAG Chat API - Gustave Eiffel Hackathon 2026
=============================================
Version améliorée :
- Chunking fixe conservé
- Recherche vectorielle conservée
- Recherche BM25 locale ajoutée
- Reranking déterministe sobre ajouté
- Query rewriting conditionnel par LLM
- Expansion déterministe des acronymes actuariels
- Seuil de distance adaptatif
- Top-K maximum pour limiter tokens, coût et CO2
Architecture:
    User Query
    → Optional LLM Query Rewriting
    → Deterministic Query Expansion
    → Embedding
    → Hybrid Vector + BM25 Search
    → Cheap Deterministic Reranking
    → Adaptive Distance Filtering
    → Context Retrieval
    → LLM Generation
    → Response
"""

import os
import json
import logging
import math
import re
import time
import unicodedata
from collections import Counter, defaultdict
from pathlib import Path
from typing import Optional

# Must be set before chromadb is imported so the module never registers its
# posthog telemetry hook.
os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")

import requests as http_requests
import gradio as gr
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import chromadb
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pypdf import PdfReader

from llm import call_llm as call_llm_with_metrics


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Suppress non-fatal chromadb telemetry errors
logging.getLogger("chromadb.telemetry.product.posthog").setLevel(logging.CRITICAL)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

# Resolve the data directory: /data when running inside HF Spaces bucket mount,
# ./data for local development.
DATA_DIR = Path("/data") if Path("/data").is_dir() else Path("./data")

CHROMA_PERSIST_DIR = str(DATA_DIR / "chroma_db")
TRAIN_DOCS_DIR = Path("./train_data")
COLLECTION_NAME = "rag_documents"

# Chunking fixe de base
CHUNK_SIZE = 512
CHUNK_OVERLAP = 50

# Ancien top_k par défaut
TOP_K_RESULTS = 3


# ---------------------------------------------------------------------------
# Retrieval Configuration
# ---------------------------------------------------------------------------

# Nombre de candidats récupérés par la recherche vectorielle.
# On reste volontairement bas pour limiter la latence.
RETRIEVAL_CANDIDATES = 8

# Nombre de candidats récupérés par BM25.
# BM25 est local et ne consomme pas de tokens, mais on limite quand même
# pour garder un reranking rapide.
BM25_CANDIDATES = 8

# Nombre maximum de candidats hybrides rerankés après déduplication.
RERANK_CANDIDATES_LIMIT = 12

# Poids du reranker déterministe.
# Aucun modèle supplémentaire : coût API nul, CO2 quasi nul côté LLM.
RERANK_VECTOR_WEIGHT = 0.60
RERANK_BM25_WEIGHT = 0.30
RERANK_TERM_COVERAGE_WEIGHT = 0.10

# Seuil minimal du score reranké pour accepter un chunk lexicalement pertinent.
# Garde-fou contre les chunks récupérés par mots-clés trop faibles.
MIN_RERANK_SCORE = 0.12
MIN_TERM_COVERAGE = 0.15

# Seuil de distance par défaut.
# Plus le seuil est bas, plus on est strict.
DISTANCE_THRESHOLD = 0.60

# Nombre maximum de chunks envoyés au LLM.
# 2 réduit les tokens / coût / CO2 par rapport à 3.
MAX_CONTEXT_CHUNKS = 2

# Si True, quand aucun chunk ne passe les filtres hybrides,
# on envoie quand même le meilleur chunk.
# Ici on garde False pour éviter les mauvais contextes.
FALLBACK_TO_BEST_CHUNK = False


# ---------------------------------------------------------------------------
# Query Rewriting Configuration
# ---------------------------------------------------------------------------

# Active/désactive le query rewriting LLM.
QUERY_REWRITE_ENABLED = True

# On ne reformule que les questions courtes/ambiguës pour limiter coût, tokens et CO2.
QUERY_REWRITE_MAX_WORDS = 4

# Nombre maximal de tokens générés par le LLM pour la reformulation.
QUERY_REWRITE_MAX_TOKENS = 30


# ---------------------------------------------------------------------------
# Load Config
# ---------------------------------------------------------------------------

_CONFIG_PATH = DATA_DIR / "config.json"
if not _CONFIG_PATH.exists():
    _CONFIG_PATH = Path(__file__).parent / "config.json"
    logger.warning(
        f"No config.json found in {DATA_DIR} — falling back to root config.json. "
        "Copy config.json to the data directory and fill in your values for production use."
    )

with open(_CONFIG_PATH, encoding="utf-8") as _f:
    _config = json.load(_f)

# Embedding model
EMBEDDING_ENDPOINT_URL = _config["embedding"]["endpoint_url"]
EMBEDDING_MODEL_NAME = _config["embedding"]["model"]

# LLM
LLM_ENDPOINT_URL = _config["llm"]["endpoint_url"]
LLM_MODEL_NAME = _config["llm"]["model"]
LLM_MAX_TOKENS = _config["llm"].get("max_completion_tokens", 512)
LLM_TEMPERATURE = _config["llm"].get("temperature", 0.7)
LLM_TOP_P = _config["llm"].get("top_p", 0.95)

# Azure API key from environment variable
AZURE_API_KEY = os.environ.get("AZURE_API_KEY")
if not AZURE_API_KEY:
    logger.warning("AZURE_API_KEY is not set — LLM and embedding calls will fail.")

# Prompt template
_PROMPT_TEMPLATE_PATH = Path(__file__).parent / "prompts" / "rag_prompt.txt"
RAG_PROMPT_TEMPLATE = _PROMPT_TEMPLATE_PATH.read_text(encoding="utf-8")

logger.info(f"Embedding model configured: {EMBEDDING_MODEL_NAME} via Azure OpenAI")


# ---------------------------------------------------------------------------
# Initialize Vector Store
# ---------------------------------------------------------------------------

logger.info(f"Initializing ChromaDB at: {CHROMA_PERSIST_DIR}")

chroma_client = chromadb.PersistentClient(
    path=CHROMA_PERSIST_DIR,
    settings=Settings(anonymized_telemetry=False),
)

collection = chroma_client.get_or_create_collection(
    name=COLLECTION_NAME,
    metadata={"hnsw:space": "cosine"},
)

logger.info(f"ChromaDB collection '{COLLECTION_NAME}' ready. Documents: {collection.count()}")
logger.info(f"LLM configured: {LLM_MODEL_NAME} via {LLM_ENDPOINT_URL}")

# Cache global du BM25.
# Il est reconstruit seulement quand le nombre de chunks change.
BM25_INDEX_CACHE = {
    "count": -1,
    "ids": [],
    "documents": [],
    "metadatas": [],
    "doc_tokens": [],
    "doc_lengths": [],
    "avgdl": 0.0,
    "idf": {},
    "inverted_index": {},
}


# ---------------------------------------------------------------------------
# Helper Functions
# ---------------------------------------------------------------------------

def extract_text_from_pdf(pdf_path: Path) -> str:
    """
    Extract text content from a PDF file using pypdf.
    """
    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()}")

    full_text = "\n\n".join(pages_text)

    logger.info(
        f"Extracted {len(reader.pages)} pages from PDF: "
        f"{pdf_path.name} ({len(full_text)} chars)"
    )

    return full_text


def chunk_text(text: str, source: str = "unknown") -> list[dict]:
    """
    Chunking fixe de base.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
    )

    chunks = splitter.split_text(text)

    return [
        {
            "text": chunk,
            "source": source,
            "chunk_index": i,
        }
        for i, chunk in enumerate(chunks)
        if chunk.strip()
    ]


def generate_embeddings(texts: list[str]) -> list[list[float]]:
    """
    Generate vector embeddings via the Azure OpenAI /embeddings endpoint.
    """
    headers = {
        "api-key": AZURE_API_KEY,
        "Content-Type": "application/json",
    }

    payload = {
        "input": texts,
        "model": EMBEDDING_MODEL_NAME,
    }

    try:
        resp = http_requests.post(
            EMBEDDING_ENDPOINT_URL,
            headers=headers,
            json=payload,
            timeout=120,
        )
        resp.raise_for_status()
        data = resp.json()
        return [item["embedding"] for item in data["data"]]

    except http_requests.exceptions.HTTPError as e:
        logger.error(f"Embedding API call failed: {e} — {resp.text}")
        raise HTTPException(
            status_code=503,
            detail=f"Embedding service unavailable: {str(e)}",
        )

    except (http_requests.exceptions.JSONDecodeError, ValueError) as e:
        logger.error(
            f"Embedding API returned non-JSON response "
            f"(status {resp.status_code}): {repr(resp.text)}"
        )
        raise HTTPException(
            status_code=502,
            detail="Embedding service returned an invalid response",
        )

    except (KeyError, IndexError) as e:
        logger.error(f"Unexpected embedding response format: {e} — body: {resp.text}")
        raise HTTPException(
            status_code=502,
            detail="Unexpected response from embedding service",
        )


def add_documents_to_vectorstore(documents: list[dict]) -> int:
    """
    Save document embeddings to the ChromaDB vector store.
    """
    if not documents:
        return 0

    texts = [doc["text"] for doc in documents]
    embeddings = generate_embeddings(texts)

    existing_count = collection.count()

    ids = [
        f"doc_{existing_count + i}"
        for i in range(len(documents))
    ]

    metadatas = [
        {
            "source": doc.get("source", "unknown"),
            "chunk_index": doc.get("chunk_index", i),
        }
        for i, doc in enumerate(documents)
    ]

    collection.add(
        ids=ids,
        embeddings=embeddings,
        documents=texts,
        metadatas=metadatas,
    )

    logger.info(f"Added {len(documents)} chunks to vector store. Total: {collection.count()}")

    return len(documents)



# ---------------------------------------------------------------------------
# Lightweight BM25 + Deterministic Reranker
# ---------------------------------------------------------------------------

SEARCH_STOPWORDS = {
    "a", "au", "aux", "avec", "ce", "ces", "cette", "dans", "de", "des",
    "du", "elle", "en", "et", "est", "il", "ils", "je", "la", "le", "les",
    "leur", "leurs", "mais", "ou", "où", "par", "pas", "pour", "que", "qui",
    "sur", "un", "une", "se", "sa", "son", "ses", "the", "of", "and", "to",
    "in", "is", "for", "on", "with", "what", "how", "why",
}


def normalize_for_search(text: str) -> str:
    """
    Normalise un texte pour la recherche lexicale :
    - minuscules ;
    - suppression des accents ;
    - conservation des lettres/chiffres.
    """
    text = unicodedata.normalize("NFKD", text.lower())
    text = "".join(ch for ch in text if not unicodedata.combining(ch))
    return text


def tokenize_for_search(text: str) -> list[str]:
    """
    Tokenisation sobre pour BM25 et le reranking.
    Aucun modèle NLP externe n'est chargé.
    """
    normalized = normalize_for_search(text)
    tokens = re.findall(r"[a-z0-9]+", normalized)

    return [
        token for token in tokens
        if token not in SEARCH_STOPWORDS and (len(token) >= 2 or token.isdigit())
    ]


def build_bm25_index_if_needed() -> dict:
    """
    Construit un index BM25 local à partir des chunks déjà présents dans ChromaDB.
    Sobriété :
    - aucun appel API ;
    - aucun embedding supplémentaire ;
    - reconstruction uniquement si le nombre de chunks change.
    """
    current_count = collection.count()

    if BM25_INDEX_CACHE["count"] == current_count:
        return BM25_INDEX_CACHE

    logger.info(f"Rebuilding BM25 index for {current_count} chunks.")

    if current_count == 0:
        BM25_INDEX_CACHE.update({
            "count": 0,
            "ids": [],
            "documents": [],
            "metadatas": [],
            "doc_tokens": [],
            "doc_lengths": [],
            "avgdl": 0.0,
            "idf": {},
            "inverted_index": {},
        })
        return BM25_INDEX_CACHE

    stored = collection.get(include=["documents", "metadatas"])
    ids = stored.get("ids", [])
    documents = stored.get("documents", []) or []
    metadatas = stored.get("metadatas", []) or []

    doc_tokens = [tokenize_for_search(doc or "") for doc in documents]
    doc_lengths = [len(tokens) for tokens in doc_tokens]
    avgdl = sum(doc_lengths) / max(len(doc_lengths), 1)

    doc_freq = Counter()
    inverted_index = defaultdict(list)

    for doc_idx, tokens in enumerate(doc_tokens):
        counts = Counter(tokens)

        for term, freq in counts.items():
            doc_freq[term] += 1
            inverted_index[term].append((doc_idx, freq))

    total_docs = len(documents)
    idf = {
        term: math.log(1 + (total_docs - freq + 0.5) / (freq + 0.5))
        for term, freq in doc_freq.items()
    }

    BM25_INDEX_CACHE.update({
        "count": current_count,
        "ids": ids,
        "documents": documents,
        "metadatas": metadatas,
        "doc_tokens": doc_tokens,
        "doc_lengths": doc_lengths,
        "avgdl": avgdl,
        "idf": idf,
        "inverted_index": dict(inverted_index),
    })

    return BM25_INDEX_CACHE


def bm25_search(query: str, top_n: int = BM25_CANDIDATES) -> list[dict]:
    """
    Recherche BM25 locale.
    BM25 favorise les correspondances exactes de termes, utile pour :
    BEL, SCR, TVOG, GLM, noms de méthodes, formules, etc.
    """
    index = build_bm25_index_if_needed()

    if index["count"] == 0:
        return []

    query_terms = tokenize_for_search(query)

    if not query_terms:
        return []

    k1 = 1.2
    b = 0.75
    scores = defaultdict(float)
    unique_terms = set(query_terms)

    for term in unique_terms:
        postings = index["inverted_index"].get(term, [])
        term_idf = index["idf"].get(term, 0.0)

        for doc_idx, freq in postings:
            dl = index["doc_lengths"][doc_idx]
            avgdl = index["avgdl"] or 1.0
            denom = freq + k1 * (1 - b + b * dl / avgdl)
            scores[doc_idx] += term_idf * (freq * (k1 + 1)) / max(denom, 1e-9)

    ranked = sorted(scores.items(), key=lambda item: item[1], reverse=True)[:top_n]

    results = []
    for doc_idx, score in ranked:
        metadata = index["metadatas"][doc_idx] or {}
        results.append({
            "id": index["ids"][doc_idx],
            "text": index["documents"][doc_idx],
            "source": metadata.get("source", "unknown"),
            "chunk_index": metadata.get("chunk_index"),
            "distance": None,
            "similarity_score": 0.0,
            "bm25_score": score,
            "retrieval_methods": {"bm25"},
        })

    return results


def minmax_normalize(values: list[float]) -> list[float]:
    """
    Normalisation simple entre 0 et 1 pour fusionner des scores hétérogènes.
    """
    if not values:
        return []

    min_value = min(values)
    max_value = max(values)

    if max_value <= min_value:
        return [1.0 if max_value > 0 else 0.0 for _ in values]

    return [
        (value - min_value) / (max_value - min_value)
        for value in values
    ]


def term_coverage_score(query: str, text: str) -> float:
    """
    Part des termes importants de la requête retrouvés dans le chunk.
    Sert de petit garde-fou lexical dans le reranking.
    """
    query_terms = set(tokenize_for_search(query))

    if not query_terms:
        return 0.0

    text_terms = set(tokenize_for_search(text))
    return len(query_terms & text_terms) / len(query_terms)


def merge_candidates(vector_candidates: list[dict], bm25_candidates: list[dict]) -> list[dict]:
    """
    Fusionne les candidats vectoriels et BM25 sans doublons.
    Clé principale : source + chunk_index.
    Fallback : hash du texte si les métadonnées sont absentes.
    """
    merged = {}

    for candidate in vector_candidates + bm25_candidates:
        key = (
            candidate.get("source"),
            candidate.get("chunk_index"),
        )

        if key[1] is None:
            key = (candidate.get("source"), hash(candidate.get("text", "")))

        if key not in merged:
            merged[key] = {
                **candidate,
                "bm25_score": candidate.get("bm25_score", 0.0),
                "retrieval_methods": set(candidate.get("retrieval_methods", set())),
            }
        else:
            existing = merged[key]

            if candidate.get("distance") is not None:
                existing["distance"] = candidate["distance"]
                existing["similarity_score"] = candidate.get("similarity_score", existing.get("similarity_score", 0.0))

            existing["bm25_score"] = max(
                existing.get("bm25_score", 0.0),
                candidate.get("bm25_score", 0.0),
            )
            existing["retrieval_methods"].update(candidate.get("retrieval_methods", set()))

    return list(merged.values())


def rerank_candidates(candidates: list[dict], query: str, effective_threshold: float) -> list[dict]:
    """
    Reranker déterministe et sobre.
    Il ne charge aucun cross-encoder et ne fait aucun appel LLM.
    Score final = vectoriel + BM25 + couverture lexicale.
    """
    if not candidates:
        return []

    vector_scores = [
        max(0.0, candidate.get("similarity_score", 0.0) or 0.0)
        for candidate in candidates
    ]
    bm25_scores = [
        max(0.0, candidate.get("bm25_score", 0.0) or 0.0)
        for candidate in candidates
    ]
    coverage_scores = [
        term_coverage_score(query, candidate.get("text", ""))
        for candidate in candidates
    ]

    vector_norm = minmax_normalize(vector_scores)
    bm25_norm = minmax_normalize(bm25_scores)

    reranked = []
    for idx, candidate in enumerate(candidates):
        rerank_score = (
            RERANK_VECTOR_WEIGHT * vector_norm[idx]
            + RERANK_BM25_WEIGHT * bm25_norm[idx]
            + RERANK_TERM_COVERAGE_WEIGHT * coverage_scores[idx]
        )

        distance = candidate.get("distance")
        vector_passed = distance is not None and distance <= effective_threshold
        lexical_passed = (
            bm25_norm[idx] > 0
            and coverage_scores[idx] >= MIN_TERM_COVERAGE
            and rerank_score >= MIN_RERANK_SCORE
        )

        candidate["rerank_score"] = rerank_score
        candidate["bm25_normalized_score"] = bm25_norm[idx]
        candidate["term_coverage"] = coverage_scores[idx]
        candidate["passed_threshold"] = vector_passed or lexical_passed
        candidate["retrieval_method"] = "+".join(sorted(candidate.get("retrieval_methods", [])))

        # Valeurs de compatibilité pour l'affichage/API.
        if candidate.get("distance") is None:
            candidate["distance"] = 1.0
        if candidate.get("similarity_score") is None:
            candidate["similarity_score"] = 0.0

        reranked.append(candidate)

    reranked.sort(key=lambda x: x["rerank_score"], reverse=True)

    return reranked[:RERANK_CANDIDATES_LIMIT]


# ---------------------------------------------------------------------------
# Query Rewriting and Query Expansion
# ---------------------------------------------------------------------------

def clean_query_words(query: str) -> list[str]:
    """
    Nettoyage simple d'une question pour détecter mots/acronymes.
    """
    q = query.lower().strip()

    cleaned = (
        q.replace("?", " ")
        .replace(",", " ")
        .replace(".", " ")
        .replace(";", " ")
        .replace(":", " ")
        .replace("'", " ")
        .replace('"', " ")
        .replace("(", " ")
        .replace(")", " ")
    )

    return cleaned.split()


def should_rewrite_query(query: str) -> bool:
    """
    Détermine si la question doit être reformulée par LLM.
    On limite volontairement le rewriting aux questions courtes ou ambiguës
    pour éviter d'ajouter un appel LLM inutile à chaque requête.
    """
    if not QUERY_REWRITE_ENABLED:
        return False

    words = clean_query_words(query)

    if not words:
        return False

    technical_terms = {
        "bel", "scr", "mcr", "glm", "mrh",
        "var", "tvar", "ifrs", "alm"
    }

    # Acronyme seul ou question très courte avec acronyme.
    if len(words) <= QUERY_REWRITE_MAX_WORDS and any(w in technical_terms for w in words):
        return True

    # Question très courte, potentiellement ambiguë.
    if len(words) <= 3:
        return True

    return False


def rewrite_query_with_llm(query: str) -> tuple[str, dict]:
    """
    Reformule la question utilisateur pour améliorer la recherche vectorielle.
    Important :
    - La reformulation sert uniquement au retrieval.
    - La question originale reste utilisée dans le prompt final.
    - La fonction retourne aussi les métriques du rewriting.
    """
    default_info = {
        "query_rewrite_used": False,
        "original_query": query,
        "rewritten_query": query,
        "query_rewrite_prompt_tokens": 0,
        "query_rewrite_completion_tokens": 0,
        "query_rewrite_total_tokens": 0,
        "query_rewrite_co2_grams": None,
        "query_rewrite_energy_kwh": None,
    }

    if not should_rewrite_query(query):
        return query, default_info

    rewrite_prompt = f"""
Tu reformules une question pour améliorer une recherche dans un corpus de mémoires d'actuariat.
Règles :
- Ne réponds pas à la question.
- Ne rajoute pas d'information inventée.
- Explicite seulement les acronymes actuariels évidents si présents :
  BEL = Best Estimate Liability,
  SCR = Solvency Capital Requirement,
  MCR = Minimum Capital Requirement,
  GLM = modèle linéaire généralisé,
  MRH = multirisque habitation,
  VaR = Value at Risk,
  TVaR = Tail Value at Risk,
  ALM = Asset Liability Management.
- Retourne une seule question reformulée, en français.
- Maximum 25 mots.
- Ne retourne pas de JSON.
Question originale :
{query}
Question reformulée :
""".strip()

    try:
        rewrite_result = call_llm_with_metrics(
            rewrite_prompt,
            endpoint_url=LLM_ENDPOINT_URL,
            api_key=AZURE_API_KEY,
            model=LLM_MODEL_NAME,
            max_completion_tokens=QUERY_REWRITE_MAX_TOKENS,
            temperature=0,
            top_p=1,
        )

        rewritten_query = rewrite_result["content"].strip()
        rewritten_query = rewritten_query.strip('"').strip("'").strip()

        if not rewritten_query:
            return query, default_info

        tokens = rewrite_result.get("tokens", {})
        info = {
            "query_rewrite_used": True,
            "original_query": query,
            "rewritten_query": rewritten_query,
            "query_rewrite_prompt_tokens": tokens.get("prompt", 0),
            "query_rewrite_completion_tokens": tokens.get("completion", 0),
            "query_rewrite_total_tokens": tokens.get("total", 0),
            "query_rewrite_co2_grams": rewrite_result.get("co2_grams"),
            "query_rewrite_energy_kwh": rewrite_result.get("energy_kwh"),
        }

        logger.info(
            f"Query rewritten: original='{query}' | rewritten='{rewritten_query}'"
        )

        return rewritten_query, info

    except Exception as e:
        logger.error(f"Query rewriting failed: {e}")
        return query, default_info


def expand_query(query: str) -> str:
    """
    Enrichit les acronymes actuariels pour améliorer la recherche vectorielle.
    Important :
    - Cela ne change pas la question envoyée au LLM final.
    - Cela change seulement la requête utilisée pour chercher les chunks.
    - Pas besoin de refaire l'ingestion.
    """
    q = query.lower().strip()

    expansions = {
        "bel": "Best Estimate Liability Best Estimate provision technique assurance vie solvabilité",
        "scr": "Solvency Capital Requirement capital de solvabilité Solvabilité II",
        "mcr": "Minimum Capital Requirement minimum capital requis Solvabilité II",
        "glm": "modèle linéaire généralisé GLM tarification fréquence sévérité sinistres",
        "mrh": "multirisque habitation assurance habitation sinistres habitation",
        "var": "Value at Risk VaR quantile risque capital économique",
        "tvar": "Tail Value at Risk TVaR risque extrême capital économique",
        "ifrs": "IFRS 17 norme comptable assurance contrats d'assurance",
        "alm": "Asset Liability Management gestion actif passif",
    }

    words = set(clean_query_words(q))

    expanded = query

    for term, expansion in expansions.items():
        if term in words or q == term:
            expanded += " " + expansion

    return expanded


def build_retrieval_query(query: str) -> tuple[str, dict]:
    """
    Construit la requête utilisée pour le retrieval :
    1. rewriting LLM conditionnel ;
    2. expansion déterministe des acronymes.
    On retourne aussi les informations de rewriting pour les métriques.
    """
    rewritten_query, rewrite_info = rewrite_query_with_llm(query)
    retrieval_query = expand_query(rewritten_query)

    rewrite_info["retrieval_query"] = retrieval_query

    return retrieval_query, rewrite_info


def get_distance_threshold(query: str) -> float:
    """
    Seuil adaptatif selon le type de question.
    Idée :
    - Requêtes très courtes ou acronymes : seuil plus permissif.
    - Requêtes normales : seuil standard.
    """
    q = query.lower().strip()

    technical_terms = {
        "bel", "scr", "mcr", "glm", "mrh",
        "var", "tvar", "ifrs", "alm"
    }

    words = clean_query_words(query)

    # Acronyme seul : BEL, SCR, GLM...
    if q in technical_terms:
        return 0.70

    # Question courte contenant un terme technique
    if len(words) <= 3 and any(word in technical_terms for word in words):
        return 0.70

    # Question très courte : seuil un peu plus permissif
    if len(words) <= 5:
        return 0.65

    # Cas général
    return DISTANCE_THRESHOLD


# ---------------------------------------------------------------------------
# Retrieval
# ---------------------------------------------------------------------------

def retrieve_relevant_context(
    query: str,
    top_k: int = TOP_K_RESULTS,
) -> tuple[list[dict], dict]:
    """
    Retrieve relevant document chunks with hybrid retrieval and cheap reranking.
    Différence avec le code de base :
    - recherche vectorielle conservée ;
    - recherche BM25 locale ajoutée ;
    - fusion des candidats ;
    - reranking déterministe sans modèle supplémentaire ;
    - filtrage par seuil adaptatif.
    """
    if collection.count() == 0:
        return [], {
            "query_rewrite_used": False,
            "original_query": query,
            "rewritten_query": query,
            "retrieval_query": query,
            "effective_threshold": None,
            "query_rewrite_prompt_tokens": 0,
            "query_rewrite_completion_tokens": 0,
            "query_rewrite_total_tokens": 0,
            "query_rewrite_co2_grams": None,
            "query_rewrite_energy_kwh": None,
        }

    retrieval_query, rewrite_info = build_retrieval_query(query)
    effective_threshold = get_distance_threshold(query)

    query_embedding = generate_embeddings([retrieval_query])[0]

    max_contexts = min(top_k or TOP_K_RESULTS, MAX_CONTEXT_CHUNKS)

    vector_candidate_count = min(
        max(RETRIEVAL_CANDIDATES, max_contexts),
        collection.count(),
    )

    vector_results = collection.query(
        query_embeddings=[query_embedding],
        n_results=vector_candidate_count,
        include=["documents", "metadatas", "distances"],
    )

    vector_candidates = []

    for i in range(len(vector_results["documents"][0])):
        distance = vector_results["distances"][0][i]
        metadata = vector_results["metadatas"][0][i] or {}

        vector_candidates.append({
            "text": vector_results["documents"][0][i],
            "source": metadata.get("source", "unknown"),
            "chunk_index": metadata.get("chunk_index"),
            "distance": distance,
            "similarity_score": 1 - distance,
            "bm25_score": 0.0,
            "retrieval_methods": {"vector"},
        })

    bm25_candidates = bm25_search(retrieval_query, top_n=BM25_CANDIDATES)

    merged_candidates = merge_candidates(vector_candidates, bm25_candidates)
    reranked_candidates = rerank_candidates(
        merged_candidates,
        query=retrieval_query,
        effective_threshold=effective_threshold,
    )

    filtered_contexts = [
        ctx for ctx in reranked_candidates
        if ctx.get("passed_threshold")
    ]

    # Si aucun chunk ne passe les filtres hybrides, on peut soit ne rien envoyer au LLM,
    # soit garder le meilleur chunk selon FALLBACK_TO_BEST_CHUNK.
    if not filtered_contexts and FALLBACK_TO_BEST_CHUNK and reranked_candidates:
        logger.info(
            f"No chunk passed hybrid filters. Falling back to best reranked chunk "
            f"with rerank_score={reranked_candidates[0]['rerank_score']:.4f}."
        )
        filtered_contexts = [reranked_candidates[0]]

    selected_contexts = filtered_contexts[:max_contexts]

    rewrite_info["effective_threshold"] = effective_threshold
    rewrite_info["vector_candidates"] = len(vector_candidates)
    rewrite_info["bm25_candidates"] = len(bm25_candidates)
    rewrite_info["merged_candidates"] = len(merged_candidates)
    rewrite_info["reranked_candidates"] = len(reranked_candidates)
    rewrite_info["passed_threshold"] = len(filtered_contexts)
    rewrite_info["selected_contexts"] = len(selected_contexts)

    logger.info(
        f"Hybrid retrieval with cheap reranking: "
        f"vector={len(vector_candidates)}, "
        f"bm25={len(bm25_candidates)}, "
        f"merged={len(merged_candidates)}, "
        f"passed={len(filtered_contexts)}, "
        f"selected={len(selected_contexts)}, "
        f"threshold={effective_threshold}, "
        f"rewrite_used={rewrite_info.get('query_rewrite_used')}, "
        f"retrieval_query='{retrieval_query[:120]}', "
        f"query='{query[:80]}'"
    )

    return selected_contexts, rewrite_info

def build_rag_prompt(query: str, contexts: list[dict]) -> str:
    """
    Construct the RAG prompt by combining retrieved context with the user question.
    La distance n'est pas ajoutée dans le prompt pour économiser quelques tokens.
    Elle reste disponible dans les sources retournées par l'API.
    """
    context_text = "\n\n".join(
        f"[Source: {ctx['source']}]\n{ctx['text']}"
        for ctx in contexts
    )

    prompt = RAG_PROMPT_TEMPLATE.format(
        context=context_text,
        question=query,
    )

    return prompt


def rag_query(query: str, top_k: int = TOP_K_RESULTS) -> dict:
    """
    End-to-end RAG pipeline.
    """
    start_time = time.perf_counter()

    contexts, retrieval_info = retrieve_relevant_context(query, top_k=top_k)

    if not contexts:
        elapsed_ms = round((time.perf_counter() - start_time) * 1000, 2)

        if collection.count() == 0:
            answer = "No documents have been ingested yet. Please upload documents first."
            explanation = "No documents found in the vector store to retrieve context from."
        else:
            effective_threshold = retrieval_info.get("effective_threshold")
            answer = (
                "I could not find sufficiently relevant context in the ingested documents "
                "to answer this question reliably."
            )
            explanation = (
                f"Documents exist in the vector store, but no retrieved chunk passed "
                f"the adaptive distance threshold of {effective_threshold}. "
                "This avoids sending weak or irrelevant context to the LLM."
            )

        rewrite_tokens = retrieval_info.get("query_rewrite_total_tokens", 0)

        return {
            "answer": answer,
            "sources": [],
            "explanation": explanation,
            "total_token": rewrite_tokens,
            "prompt_tokens": retrieval_info.get("query_rewrite_prompt_tokens", 0),
            "completion_tokens": retrieval_info.get("query_rewrite_completion_tokens", 0),
            "cached_tokens": 0,
            "query_rewrite_used": retrieval_info.get("query_rewrite_used", False),
            "rewritten_query": retrieval_info.get("rewritten_query", query),
            "retrieval_query": retrieval_info.get("retrieval_query", query),
            "query_rewrite_total_tokens": retrieval_info.get("query_rewrite_total_tokens", 0),
            "co2_grams": retrieval_info.get("query_rewrite_co2_grams"),
            "energy_kwh": retrieval_info.get("query_rewrite_energy_kwh"),
            "run_time_in_ms": elapsed_ms,
        }

    prompt = build_rag_prompt(query, contexts)

    llm_result = call_llm_with_metrics(
        prompt,
        endpoint_url=LLM_ENDPOINT_URL,
        api_key=AZURE_API_KEY,
        model=LLM_MODEL_NAME,
        max_completion_tokens=LLM_MAX_TOKENS,
        temperature=LLM_TEMPERATURE,
        top_p=LLM_TOP_P,
    )

    raw_content = llm_result["content"]
    tokens = llm_result["tokens"]

    rewrite_total_tokens = retrieval_info.get("query_rewrite_total_tokens", 0)
    answer_total_tokens = tokens["total"]
    total_token = rewrite_total_tokens + answer_total_tokens

    # Parse structured JSON response from LLM
    json_str = raw_content.strip()

    if json_str.startswith("```"):
        json_str = json_str.split("\n", 1)[-1]
        json_str = json_str.rsplit("```", 1)[0].strip()

    try:
        parsed = json.loads(json_str)
        answer = parsed["answer"]
        explanation = parsed["explanation"]
    except (json.JSONDecodeError, KeyError):
        answer = raw_content
        explanation = "LLM did not return a structured explanation."

    elapsed_ms = round((time.perf_counter() - start_time) * 1000, 2)

    # Addition approximative des métriques CO2/énergie du rewriting et de la réponse finale.
    rewrite_co2 = retrieval_info.get("query_rewrite_co2_grams")
    final_co2 = llm_result.get("co2_grams")

    if isinstance(rewrite_co2, (int, float)) and isinstance(final_co2, (int, float)):
        total_co2 = rewrite_co2 + final_co2
    else:
        total_co2 = final_co2

    rewrite_energy = retrieval_info.get("query_rewrite_energy_kwh")
    final_energy = llm_result.get("energy_kwh")

    if isinstance(rewrite_energy, (int, float)) and isinstance(final_energy, (int, float)):
        total_energy = rewrite_energy + final_energy
    else:
        total_energy = final_energy

    return {
        "answer": answer,
        "sources": [
            {
                "source": ctx["source"],
                "score": ctx["similarity_score"],
                "distance": ctx["distance"],
                "bm25_score": ctx.get("bm25_score", 0.0),
                "bm25_normalized_score": ctx.get("bm25_normalized_score", 0.0),
                "rerank_score": ctx.get("rerank_score", 0.0),
                "term_coverage": ctx.get("term_coverage", 0.0),
                "retrieval_method": ctx.get("retrieval_method", "vector"),
                "passed_threshold": ctx["passed_threshold"],
                "chunk_index": ctx.get("chunk_index"),
                "ref_text": ctx["text"],
            }
            for ctx in contexts
        ],
        "explanation": explanation,

        # Tokens totaux = rewriting éventuel + génération finale.
        "total_token": total_token,
        "prompt_tokens": tokens["prompt"] + retrieval_info.get("query_rewrite_prompt_tokens", 0),
        "completion_tokens": tokens["completion"] + retrieval_info.get("query_rewrite_completion_tokens", 0),
        "cached_tokens": tokens["cached"],

        # Détail du rewriting
        "query_rewrite_used": retrieval_info.get("query_rewrite_used", False),
        "rewritten_query": retrieval_info.get("rewritten_query", query),
        "retrieval_query": retrieval_info.get("retrieval_query", query),
        "query_rewrite_total_tokens": retrieval_info.get("query_rewrite_total_tokens", 0),
        "answer_generation_total_tokens": answer_total_tokens,

        # Métriques CO2/énergie totales approximatives
        "co2_grams": total_co2,
        "energy_kwh": total_energy,

        "run_time_in_ms": elapsed_ms,
    }


# ---------------------------------------------------------------------------
# Ingest Train Documents
# ---------------------------------------------------------------------------

def ingest_train_documents():
    """
    Load and embed training documents into the vector store.
    """
    if collection.count() > 0:
        logger.info("Vector store already has documents, skipping ingestion.")
        return

    if not TRAIN_DOCS_DIR.exists():
        logger.warning(f"No train_data directory found at: {TRAIN_DOCS_DIR}")
        return

    total_chunks = 0

    for file_path in TRAIN_DOCS_DIR.rglob("*.txt"):
        logger.info(f"Ingesting text file: {file_path.name}")
        text = file_path.read_text(encoding="utf-8", errors="ignore")
        chunks = chunk_text(text, source=file_path.name)
        total_chunks += add_documents_to_vectorstore(chunks)

    for file_path in TRAIN_DOCS_DIR.rglob("*.pdf"):
        logger.info(f"Ingesting PDF file: {file_path.name}")
        text = extract_text_from_pdf(file_path)

        if text.strip():
            chunks = chunk_text(text, source=file_path.name)
            total_chunks += add_documents_to_vectorstore(chunks)
        else:
            logger.warning(f"No extractable text found in: {file_path.name}")

    logger.info(
        f"Train document ingestion complete. "
        f"Chunks added: {total_chunks}. Total chunks: {collection.count()}"
    )


# ---------------------------------------------------------------------------
# FastAPI Application
# ---------------------------------------------------------------------------

app = FastAPI(
    title="RAG Chat API - Gustave Eiffel Hackathon 2026",
    description="A RAG system with /query endpoint for evaluation",
    version="1.0.0",
)


class QueryRequest(BaseModel):
    """
    Request schema for the /query endpoint.
    """
    query: str
    top_k: Optional[int] = TOP_K_RESULTS


class IngestRequest(BaseModel):
    """
    Request schema for the /ingest endpoint.
    """
    text: str
    source: str = "user_upload"


@app.post("/query")
async def query_endpoint(request: QueryRequest):
    """
    RAG Query Endpoint.
    """
    logger.info(f"Query received: {request.query}")
    result = rag_query(request.query, top_k=request.top_k)
    return JSONResponse(content=result)


@app.post("/ingest")
async def ingest_endpoint(request: IngestRequest):
    """
    Document Ingestion Endpoint.
    """
    chunks = chunk_text(request.text, source=request.source)
    count = add_documents_to_vectorstore(chunks)

    return JSONResponse(content={
        "status": "success",
        "chunks_added": count,
        "total_chunks": collection.count(),
    })


@app.get("/health")
async def health_check():
    """
    Health check endpoint.
    """
    return {
        "status": "healthy",
        "documents_in_store": collection.count(),
        "embedding_model": EMBEDDING_MODEL_NAME,
        "llm_model": LLM_MODEL_NAME,
        "retrieval_strategy": "hybrid_vector_bm25_plus_deterministic_reranker_plus_adaptive_distance_threshold",
        "query_rewrite_enabled": QUERY_REWRITE_ENABLED,
        "query_rewrite_max_words": QUERY_REWRITE_MAX_WORDS,
        "retrieval_candidates": RETRIEVAL_CANDIDATES,
        "bm25_candidates": BM25_CANDIDATES,
        "rerank_candidates_limit": RERANK_CANDIDATES_LIMIT,
        "rerank_weights": {
            "vector": RERANK_VECTOR_WEIGHT,
            "bm25": RERANK_BM25_WEIGHT,
            "term_coverage": RERANK_TERM_COVERAGE_WEIGHT,
        },
        "default_distance_threshold": DISTANCE_THRESHOLD,
        "max_context_chunks": MAX_CONTEXT_CHUNKS,
        "fallback_to_best_chunk": FALLBACK_TO_BEST_CHUNK,
    }


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

def gradio_query(question: str) -> tuple[str, str, str, str, str]:
    """
    Handle queries from the Gradio chat interface.
    """
    if not question.strip():
        return "Please enter a question.", "", "", "", ""

    result = rag_query(question)

    sources_text = "\n".join(
        (
            f"  - {s['source']} "
            f"(method: {s.get('retrieval_method', 'vector')}, "
            f"distance: {s.get('distance', 0):.4f}, "
            f"vector: {s.get('score', 0):.4f}, "
            f"bm25: {s.get('bm25_score', 0):.4f}, "
            f"rerank: {s.get('rerank_score', 0):.4f})"
        )
        for s in result["sources"]
    )

    rewrite_info = ""
    if result.get("query_rewrite_used"):
        rewrite_info = (
            f"\n\n🔎 Rewritten query used for retrieval:\n"
            f"{result.get('rewritten_query')}"
        )

    answer = (
        f"{result['answer']}\n\n📚 Sources:\n{sources_text}{rewrite_info}"
        if result["sources"]
        else f"{result['answer']}{rewrite_info}"
    )

    explanation = result.get("explanation", "")
    token_info = str(result.get("total_token", 0))

    co2_value = result.get("co2_grams")
    co2_info = f"{co2_value:.4f} g" if isinstance(co2_value, (int, float)) else "N/A"

    run_time = f"{result.get('run_time_in_ms', 0)} ms"

    return answer, explanation, token_info, co2_info, run_time


def gradio_ingest(text: str, source_name: str) -> str:
    """
    Handle document ingestion from the Gradio UI.
    """
    if not text.strip():
        return "Please provide text to ingest."

    chunks = chunk_text(text, source=source_name or "user_upload")
    count = add_documents_to_vectorstore(chunks)

    return (
        f"✅ Ingested {count} chunks. "
        f"Total documents in store: {collection.count()}"
    )


with gr.Blocks(title="RAG Chat API - Gustave Eiffel Hackathon") as demo:
    gr.Markdown("""
    # 🗼 RAG Chat API - Gustave Eiffel Hackathon 2026
    This application demonstrates a complete **Retrieval-Augmented Generation (RAG)** system.
    **Current improvements:**
    - Conditional LLM query rewriting for short or ambiguous questions
    - Deterministic query expansion for actuarial acronyms such as BEL, SCR, GLM, MRH
    - Hybrid vector + BM25 search
    - Cheap deterministic reranking, without extra LLM call
    - Vector search with adaptive distance threshold
    - Fewer context chunks sent to the LLM to reduce tokens, cost and CO2
    **API Endpoint:** Use `POST /query` with `{"query": "your question"}` for programmatic access.
    ---
    """)

    with gr.Tab("💬 Chat"):
        gr.Markdown("Ask questions about the ingested documents.")

        with gr.Row():
            query_input = gr.Textbox(
                label="Your Question",
                placeholder="e.g., BEL, SCR, GLM, ou Comment le SCR est-il modélisé en assurance vie ?",
                lines=2,
            )

        query_button = gr.Button("Ask", variant="primary")
        query_output = gr.Textbox(label="Answer", lines=8, interactive=False)
        query_explanation = gr.Textbox(label="Explanation", lines=3, interactive=False)

        with gr.Row():
            query_tokens = gr.Textbox(label="Total Tokens", interactive=False)
            query_co2 = gr.Textbox(label="CO2 Emission", interactive=False)
            query_runtime = gr.Textbox(label="Run Time", interactive=False)

        query_button.click(
            fn=gradio_query,
            inputs=query_input,
            outputs=[
                query_output,
                query_explanation,
                query_tokens,
                query_co2,
                query_runtime,
            ],
        )

    with gr.Tab("📄 Ingest Documents"):
        gr.Markdown("Add new documents to the knowledge base.")

        doc_text = gr.Textbox(
            label="Document Text",
            placeholder="Paste your document text here...",
            lines=10,
        )

        doc_source = gr.Textbox(
            label="Source Name",
            placeholder="e.g., my_document.txt",
            value="user_upload",
        )

        ingest_button = gr.Button("Ingest Document", variant="primary")
        ingest_output = gr.Textbox(label="Status", interactive=False)

        ingest_button.click(
            fn=gradio_ingest,
            inputs=[doc_text, doc_source],
            outputs=ingest_output,
        )

    with gr.Tab("ℹ️ API Info"):
        gr.Markdown("""
        ## API Endpoints
        ### POST /query
        ```json
        {
            "query": "BEL",
            "top_k": 3
        }
        ```
        **Response:**
        ```json
        {
            "answer": "...",
            "sources": [
                {
                    "source": "document.pdf",
                    "score": 0.82,
                    "distance": 0.18
                }
            ],
            "query_rewrite_used": true,
            "rewritten_query": "Qu'est-ce que le Best Estimate Liability dans les mémoires d'actuariat ?"
        }
        ```
        ### POST /ingest
        ```json
        {
            "text": "Your document text here...",
            "source": "document_name.txt"
        }
        ```
        ### GET /health
        Returns system health, document count and retrieval configuration.
        """)


app = gr.mount_gradio_app(app, demo, path="/")


# ---------------------------------------------------------------------------
# Entry Point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7860)