"""
RAG Chat API - Gustave Eiffel Hackathon 2026
=============================================

This application demonstrates a complete Retrieval-Augmented Generation (RAG) system
deployed as a Hugging Face Space. It includes:

1. /query endpoint - Called by the RAG evaluation system
2. LLM API integration via Hugging Face Inference API
3. Document embedding using sentence-transformers
4. ChromaDB as the vector store (runs locally within HF Spaces)
5. End-to-end RAG pipeline: ingest → embed → retrieve → generate

Architecture:
    User Query → Embedding → Vector Search → Context Retrieval → LLM Generation → Response
"""

import os
import re
import json
import logging
import time
from pathlib import Path
from typing import Optional

# Must be set before chromadb is imported so the module never registers its
# posthog telemetry hook (workaround for posthog v3 API incompatibility).
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 sentence_transformers import CrossEncoder
from rank_bm25 import BM25Okapi

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 (posthog v3 API incompatibility)
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"
CHUNK_SIZE = 1000   # était 512
CHUNK_OVERLAP = 150  # était 50

# ---------------------------------------------------------------------------
# Reranking configuration
# ---------------------------------------------------------------------------
# Mettre à True pour activer le reranking, False pour le désactiver
ENABLE_RERANKING = True

# Nombre de chunks récupérés lors de la première étape (recherche vectorielle)
# Plus ce nombre est élevé, meilleur est le rappel — mais plus le reranker est lent
TOP_K_RETRIEVAL = 20

# Nombre de chunks conservés après le reranking (envoyés au LLM)
# Doit être inférieur ou égal à TOP_K_RETRIEVAL
TOP_K_RERANKED = 2

# Modèle cross-encoder multilingue (supporte le français)
RERANKER_MODEL_NAME = "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1"

# ---------------------------------------------------------------------------
# BM25 configuration (Hybrid Search)
# ---------------------------------------------------------------------------
# Mettre à True pour activer la recherche hybride BM25 + vectorielle (Reciprocal Rank Fusion)
ENABLE_BM25 = True

# Constante RRF : valeur standard = 60 (plus k est grand, moins les rangs élevés sont pénalisés)
RRF_K = 60

# ---------------------------------------------------------------------------
# Configuration par type de question — valeurs de fallback si le routeur échoue
# ---------------------------------------------------------------------------
QCM_TOP_K_RETRIEVAL  = 12
QCM_TOP_K_RERANKED   = 2
QCM_LLM_MODEL        = "gpt-5.mini"

OPEN_TOP_K_RETRIEVAL = 20
OPEN_TOP_K_RERANKED  = 4
OPEN_LLM_MODEL       = "gpt-5.1"

# ---------------------------------------------------------------------------
# Router configuration (routage intelligent avant retrieval)
# ---------------------------------------------------------------------------
# Mettre à False pour désactiver le routeur et revenir aux heuristiques simples
ENABLE_ROUTER = True
ROUTER_MODEL  = "gpt-5.mini"  # modèle rapide/peu coûteux pour le seul routage

# Settings loaded from data/config.json (runtime config).
# Falls back to the root config.json which serves as the example/template.
_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 (example file). "
        "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 (Azure OpenAI)
EMBEDDING_ENDPOINT_URL = _config["embedding"]["endpoint_url"]
EMBEDDING_MODEL_NAME = _config["embedding"]["model"]

# LLM (Azure OpenAI)
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 (shared by both LLM and embedding endpoints)
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 loaded from file so it can be edited without touching application code
_PROMPT_TEMPLATE_PATH = Path(__file__).parent / "prompts" / "rag_prompt.txt"
RAG_PROMPT_TEMPLATE = _PROMPT_TEMPLATE_PATH.read_text(encoding="utf-8")

_ROUTER_PROMPT_PATH = Path(__file__).parent / "prompts" / "router_prompt.txt"
ROUTER_PROMPT = _ROUTER_PROMPT_PATH.read_text(encoding="utf-8")

# ---------------------------------------------------------------------------
# Step 1: Embedding via Azure OpenAI
# ---------------------------------------------------------------------------
# Embeddings are generated by calling the Azure OpenAI /embeddings endpoint.
# No local model is loaded — the API handles all inference.

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

# ---------------------------------------------------------------------------
# Step 2: Initialize the Vector Store (ChromaDB)
# ---------------------------------------------------------------------------
# ChromaDB runs in-process with persistent storage. This is ideal for HF Spaces
# because it requires no external database service.

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()}")

# ---------------------------------------------------------------------------
# Step 2b: Initialize Cross-Encoder Reranker
# ---------------------------------------------------------------------------
reranker = None
if ENABLE_RERANKING:
    logger.info(f"Loading reranker model: {RERANKER_MODEL_NAME}")
    reranker = CrossEncoder(RERANKER_MODEL_NAME)
    logger.info("Reranker ready.")
else:
    logger.info("Reranking disabled (ENABLE_RERANKING = False).")

# ---------------------------------------------------------------------------
# Step 2c: Build BM25 Index (Hybrid Search)
# ---------------------------------------------------------------------------
_bm25_index: Optional[BM25Okapi] = None
_bm25_ids:     list[str] = []
_bm25_texts:   list[str] = []
_bm25_sources: list[str] = []


def _build_bm25_index() -> None:
    """Construit l'index BM25 à partir de tous les documents de ChromaDB."""
    global _bm25_index, _bm25_ids, _bm25_texts, _bm25_sources
    n = collection.count()
    if n == 0:
        logger.warning("BM25: collection vide, index non construit.")
        return
    logger.info(f"Construction de l'index BM25 ({n} documents)...")
    t0 = time.perf_counter()
    all_docs = collection.get(limit=n, include=["documents", "metadatas"])
    _bm25_ids     = all_docs["ids"]
    _bm25_texts   = all_docs["documents"]
    _bm25_sources = [m["source"] for m in all_docs["metadatas"]]
    tokenized     = [re.sub(r"[^\w\s]", " ", txt.lower()).split() for txt in _bm25_texts]
    _bm25_index   = BM25Okapi(tokenized)
    logger.info(f"Index BM25 prêt en {time.perf_counter() - t0:.1f}s.")


if ENABLE_BM25:
    _build_bm25_index()
else:
    logger.info("BM25 désactivé (ENABLE_BM25 = False).")

# ---------------------------------------------------------------------------
# Step 3: LLM Configuration (Azure Foundry GPT-5)
# ---------------------------------------------------------------------------
# We call the Azure OpenAI-compatible endpoint directly via requests.
# Endpoint URL and model are configured in config.json.

logger.info(f"LLM configured: {LLM_MODEL_NAME} via {LLM_ENDPOINT_URL}")


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

def extract_text_from_pdf(pdf_path: Path) -> str:
    """
    Extract text content from a PDF file using pypdf.

    This demonstrates how to convert PDF documents into plain text
    for downstream embedding. Each page is extracted sequentially
    and concatenated with page separators for traceability.
    """
    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: {pdf_path.name} ({len(full_text)} chars)")
    return full_text


def chunk_text(text: str, source: str = "unknown") -> list[dict]:
    """
    Split a document into smaller chunks for embedding.

    We use LangChain's RecursiveCharacterTextSplitter which intelligently
    splits on paragraph/sentence boundaries to preserve semantic meaning.
    """
    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)]


def generate_embeddings(texts: list[str]) -> list[list[float]]:
    """
    Generate vector embeddings via the Azure OpenAI /embeddings endpoint.

    The endpoint, model name, and API key are loaded from config.json
    and the AZURE_API_KEY environment variable.
    """
    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 (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.

    This demonstrates Step 2c: How to persist embeddings for retrieval.
    Each document gets a unique ID, its embedding vector, the raw text,
    and metadata (source file, chunk index).
    """
    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["source"], "chunk_index": doc["chunk_index"]} for doc in 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)


def detect_qcm(query: str) -> bool:
    """Retourne True si la question contient des options de type A) B) C) D)."""
    return bool(re.search(r'\b[A-D][).]\s', query))


def _fallback_routing(query: str) -> dict:
    """Heuristiques de routage si le routeur LLM est désactivé ou en erreur."""
    if detect_qcm(query):
        return {
            "question_type": "qcm",   "answer_type": "factual",
            "breadth": "narrow",      "difficulty": "low",
            "top_k_retrieval": QCM_TOP_K_RETRIEVAL,
            "top_k_reranked":  QCM_TOP_K_RERANKED,
            "model":           QCM_LLM_MODEL or LLM_MODEL_NAME,
        }
    q = query.lower()
    has_comparison = any(w in q for w in ["compar", "différence", "versus", "vs ", "contrast"])
    has_synthesis  = any(w in q for w in ["enjeux", "limites", "défis", "implications", "impact", "challenges"])
    has_method     = any(w in q for w in ["étape", "processus", "comment", "décrire", "expliquer", "démarche", "how to"])
    word_count     = len(query.split())

    if has_comparison:
        return {
            "question_type": "open",   "answer_type": "comparison",
            "breadth": "broad",        "difficulty": "high",
            "top_k_retrieval": 45,     "top_k_reranked": 7,
            "model": OPEN_LLM_MODEL or LLM_MODEL_NAME,
        }
    if has_synthesis:
        return {
            "question_type": "open",   "answer_type": "synthesis",
            "breadth": "broad",        "difficulty": "high",
            "top_k_retrieval": 45,     "top_k_reranked": 7,
            "model": OPEN_LLM_MODEL or LLM_MODEL_NAME,
        }
    if has_method or word_count > 15:
        return {
            "question_type": "open",   "answer_type": "method",
            "breadth": "medium",       "difficulty": "medium",
            "top_k_retrieval": 35,     "top_k_reranked": 5,
            "model": OPEN_LLM_MODEL or LLM_MODEL_NAME,
        }
    return {
        "question_type": "open",   "answer_type": "factual",
        "breadth": "narrow",       "difficulty": "low",
        "top_k_retrieval": 20,     "top_k_reranked": 3,
        "model": QCM_LLM_MODEL or LLM_MODEL_NAME,
    }


def _call_router_llm(query: str) -> str:
    """Appel direct à GPT-5-mini pour classer la question (sans ecologits)."""
    headers = {"api-key": AZURE_API_KEY, "Content-Type": "application/json"}
    payload = {
        "model": ROUTER_MODEL,
        "messages": [
            {"role": "system", "content": ROUTER_PROMPT},
            {"role": "user",   "content": query},
        ],
        "max_completion_tokens": 200,
        "temperature": 0.0,
    }
    resp = http_requests.post(LLM_ENDPOINT_URL, headers=headers, json=payload, timeout=30)
    resp.raise_for_status()
    return resp.json()["choices"][0]["message"]["content"]


def route_question(query: str) -> dict:
    """
    Classe la question et retourne les paramètres RAG adaptés.

    Retourne un dict avec :
        question_type, answer_type, breadth, difficulty,
        top_k_retrieval, top_k_reranked, model, router_used
    """
    if not ENABLE_ROUTER:
        return {**_fallback_routing(query), "router_used": False}

    try:
        raw = _call_router_llm(query).strip()
        if raw.startswith("```"):
            raw = raw.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
        routing = json.loads(raw)
        required = {"question_type", "answer_type", "breadth", "difficulty",
                    "top_k_retrieval", "top_k_reranked", "model"}
        if not required.issubset(routing.keys()):
            raise ValueError(f"Champs manquants : {required - routing.keys()}")
        routing["top_k_retrieval"] = int(routing["top_k_retrieval"])
        routing["top_k_reranked"]  = int(routing["top_k_reranked"])
        routing["router_used"] = True
        logger.info(
            f"Router → type={routing['question_type']} answer={routing['answer_type']} "
            f"breadth={routing['breadth']} diff={routing['difficulty']} "
            f"top_k={routing['top_k_retrieval']} reranked={routing['top_k_reranked']} "
            f"model={routing['model']}"
        )
        return routing
    except Exception as e:
        logger.warning(f"Routeur LLM échoué ({e}) — fallback heuristiques.")
        return {**_fallback_routing(query), "router_used": False}


def retrieve_relevant_context(
    query: str,
    top_k: int = TOP_K_RETRIEVAL,
    enable_reranking: bool = ENABLE_RERANKING,
    top_k_reranked: int = TOP_K_RERANKED,
    enable_bm25: bool = ENABLE_BM25,
) -> list[dict]:
    """
    Retrieve the most relevant document chunks for a given query.

    Pipeline :
    1a. Recherche vectorielle dense (ChromaDB)
    1b. Recherche BM25 sparse (mots-clés exacts) — si enable_bm25
        Fusion des deux via Reciprocal Rank Fusion (RRF)
    2.  Reranking cross-encoder sur les candidats fusionnés — si enable_reranking
    """
    if collection.count() == 0:
        return []

    # ── Étape 1a : recherche vectorielle ─────────────────────────────────────
    query_embedding = generate_embeddings([query])[0]
    n_candidates    = min(top_k, collection.count())

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

    dense_ids     = dense_results["ids"][0]
    dense_id_to_ctx: dict[str, dict] = {}
    for i, doc_id in enumerate(dense_ids):
        dense_id_to_ctx[doc_id] = {
            "text":             dense_results["documents"][0][i],
            "source":           dense_results["metadatas"][0][i]["source"],
            "similarity_score": 1 - dense_results["distances"][0][i],
        }

    if not enable_bm25 or _bm25_index is None:
        contexts = list(dense_id_to_ctx.values())
    else:
        # ── Étape 1b : recherche BM25 ─────────────────────────────────────────
        tokens        = re.sub(r"[^\w\s]", " ", query.lower()).split()
        bm25_scores   = _bm25_index.get_scores(tokens)
        bm25_top_idx  = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:top_k]

        bm25_id_to_ctx: dict[str, dict] = {}
        for idx in bm25_top_idx:
            doc_id = _bm25_ids[idx]
            if doc_id not in dense_id_to_ctx:
                bm25_id_to_ctx[doc_id] = {
                    "text":             _bm25_texts[idx],
                    "source":           _bm25_sources[idx],
                    "similarity_score": float(bm25_scores[idx]),
                }

        # ── Fusion RRF ────────────────────────────────────────────────────────
        rrf_scores: dict[str, float] = {}
        for rank, doc_id in enumerate(dense_ids):
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + 1.0 / (RRF_K + rank + 1)
        for rank, idx in enumerate(bm25_top_idx):
            doc_id = _bm25_ids[idx]
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0.0) + 1.0 / (RRF_K + rank + 1)

        all_ctx    = {**dense_id_to_ctx, **bm25_id_to_ctx}
        sorted_ids = sorted(rrf_scores, key=lambda d: rrf_scores[d], reverse=True)[:top_k]
        contexts   = [all_ctx[d] for d in sorted_ids if d in all_ctx]

    # ── Étape 2 : reranking cross-encoder ────────────────────────────────────
    if not enable_reranking or reranker is None:
        return contexts

    pairs         = [[query, ctx["text"]] for ctx in contexts]
    rerank_scores = reranker.predict(pairs)
    for ctx, score in zip(contexts, rerank_scores):
        ctx["rerank_score"] = float(score)
    contexts.sort(key=lambda x: x["rerank_score"], reverse=True)

    return contexts[:top_k_reranked]


def build_rag_prompt(query: str, contexts: list[dict]) -> str:
    """
    Construct the RAG prompt by combining retrieved context with the user question.

    The prompt template instructs the LLM to:
    - Answer based ONLY on the provided context
    - Acknowledge when information is insufficient
    - Cite sources when possible
    """
    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: Optional[int] = None,
    *,
    temperature: float = LLM_TEMPERATURE,
    max_tokens: int = LLM_MAX_TOKENS,
    enable_reranking: bool = ENABLE_RERANKING,
    top_k_reranked: Optional[int] = None,
    model: Optional[str] = None,
    enable_bm25: bool = ENABLE_BM25,
    enable_router: bool = ENABLE_ROUTER,
) -> dict:
    """
    End-to-end RAG pipeline: Query → Route → Retrieve → Generate.

    Le routeur LLM détermine automatiquement top_k, top_k_reranked et model
    selon le type, la nature et la complexité de la question.
    Les valeurs passées explicitement priment toujours sur le routeur.
    """
    start_time = time.perf_counter()

    # Routage intelligent (ou fallback heuristique si désactivé / en erreur)
    _router_enabled = enable_router and ENABLE_ROUTER
    routing = route_question(query) if _router_enabled else {**_fallback_routing(query), "router_used": False}

    question_type = routing["question_type"].upper()
    answer_type   = routing.get("answer_type", "factual")
    breadth       = routing.get("breadth", "medium")
    difficulty    = routing.get("difficulty", "medium")
    router_used   = routing.get("router_used", False)

    # Les valeurs explicites priment sur le routeur
    if top_k is None:
        top_k = routing["top_k_retrieval"]
    if top_k_reranked is None:
        top_k_reranked = routing["top_k_reranked"]
    if model is None:
        model = routing["model"]

    logger.info(
        f"Routing → type={question_type} answer={answer_type} breadth={breadth} "
        f"difficulty={difficulty} | top_k={top_k} reranked={top_k_reranked} "
        f"model={model} router={'LLM' if router_used else 'heuristic'}"
    )

    # Step 1: Retrieve relevant document chunks
    contexts = retrieve_relevant_context(
        query,
        top_k=top_k,
        enable_reranking=enable_reranking,
        top_k_reranked=top_k_reranked,
        enable_bm25=enable_bm25,
    )

    if not contexts:
        elapsed_ms = round((time.perf_counter() - start_time) * 1000, 2)
        return {
            "answer": "No documents have been ingested yet. Please upload documents first.",
            "sources": [],
            "explanation": "No documents found in the vector store to retrieve context from.",
            "total_token": 0,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "cached_tokens": 0,
            "co2_grams": None,
            "energy_kwh": None,
            "run_time_in_ms": elapsed_ms,
        }

    # Step 2: Build the augmented prompt
    prompt = build_rag_prompt(query, contexts)

    # Step 3: Generate answer from LLM (with token + CO2 metrics)
    llm_result = call_llm_with_metrics(
        prompt,
        endpoint_url=LLM_ENDPOINT_URL,
        api_key=AZURE_API_KEY,
        model=model,
        max_completion_tokens=max_tokens,
        temperature=temperature,
        top_p=LLM_TOP_P,
    )
    raw_content = llm_result["content"]
    tokens = llm_result["tokens"]
    total_token = tokens["total"]

    # Parse structured JSON response from LLM (handle markdown code fences)
    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)

    # Step 4: Return structured response
    return {
        "answer": answer,
        "sources": [{"source": ctx["source"], "score": ctx["similarity_score"], "ref_text": ctx["text"]} for ctx in contexts],
        "explanation": explanation,
        "question_type":     question_type,
        "answer_type":       answer_type,
        "breadth":           breadth,
        "difficulty":        difficulty,
        "router_used":       router_used,
        "model_used":        model,
        "top_k_reranked_used": top_k_reranked,
        "total_token":       total_token,
        "prompt_tokens":     tokens["prompt"],
        "completion_tokens": tokens["completion"],
        "cached_tokens":     tokens["cached"],
        "co2_grams":         llm_result["co2_grams"],
        "energy_kwh":        llm_result["energy_kwh"],
        "run_time_in_ms":    elapsed_ms,
    }


# ---------------------------------------------------------------------------
# Step 4: Ingest Train Documents (on-demand)
# ---------------------------------------------------------------------------


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

    # Ingest plain text files (recursively in subdirectories)
    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")
        chunks = chunk_text(text, source=file_path.name)
        add_documents_to_vectorstore(chunks)

    # Ingest PDF files (recursively in subdirectories)
    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)
            add_documents_to_vectorstore(chunks)
        else:
            logger.warning(f"No extractable text found in: {file_path.name}")

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


# ---------------------------------------------------------------------------
# Step 5: FastAPI Application with /query Endpoint
# ---------------------------------------------------------------------------

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_RETRIEVAL


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 - Called by the evaluation system.

    Accepts a user query, retrieves relevant context from the vector store,
    and generates an answer using the LLM.

    Request body:
        - query (str): The user's question
        - top_k (int, optional): Number of context chunks to retrieve (default: 3)

    Returns:
        - answer (str): The generated answer
        - sources (list): Source documents used for context
        - explanation (str): Explanation of the retrieval and answer logic
        - total_token (int): Total token count from the LLM call
        - prompt_tokens (int): Prompt token count
        - completion_tokens (int): Completion token count
        - cached_tokens (int): Cached prompt tokens (when reported by provider)
        - co2_grams (float | None): Estimated CO2 emission in grams (via ecologits)
        - energy_kwh (float | None): Estimated energy use in kWh (via ecologits)
        - run_time_in_ms (float): Pipeline execution time in milliseconds
    """
    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.

    Accepts raw text, chunks it, generates embeddings, and stores in the vector store.

    Request body:
        - text (str): The document text to ingest
        - source (str, optional): Source identifier for the document
    """
    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 for monitoring."""
    return {
        "status": "healthy",
        "documents_in_store": collection.count(),
        "embedding_model": EMBEDDING_MODEL_NAME,
        "llm_model": LLM_MODEL_NAME,
    }


# ---------------------------------------------------------------------------
# Step 6: Gradio UI for Interactive Demo
# ---------------------------------------------------------------------------

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']} (relevance: {s['score']:.2f})" for s in result["sources"]
    )
    answer = f"{result['answer']}\n\n📚 Sources:\n{sources_text}" if result["sources"] else result["answer"]
    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. 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.

    **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., What is the Eiffel Tower made of?",
                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": "What is the Eiffel Tower?",
            "top_k": 3
        }
        ```

        **Response:**
        ```json
        {
            "answer": "The Eiffel Tower is...",
            "sources": [{"source": "eiffel_tower.txt", "score": 0.85}],
            "query": "What is the Eiffel Tower?"
        }
        ```

        ### POST /ingest
        ```json
        {
            "text": "Your document text here...",
            "source": "document_name.txt"
        }
        ```

        ### GET /health
        Returns system health and document count.
        """)

# Mount Gradio onto FastAPI so both the UI and API endpoints are served
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)