"""
RAG Chat API - Gustave Eiffel Hackathon 2026
=============================================
Pipeline RAG avance :
1. Recherche HYBRIDE : vectoriel + BM25 fusionnes par RRF                 [v2bm25]
2. Reranking par CROSS-ENCODER local (multilingue, gratuit, CPU)          [v3rerank]
3. ROUTING petit/gros modele (GPT-5-mini <-> GPT-5.1)                      [v3routing]
   - Question courte / QCM ferme  -> mini d'abord, escalade au gros si reponse faible
   - Question longue / ouverte    -> gros modele directement (evite la double facturation)

L'index BM25 et le cross-encoder se construisent au demarrage a partir des
chunks deja presents dans ChromaDB (pas de re-embedding necessaire).
"""

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

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

# [v2bm25] BM25 keyword search (graceful fallback if not installed)
try:
    from rank_bm25 import BM25Okapi
    _BM25_AVAILABLE = True
except ImportError:
    _BM25_AVAILABLE = False

# [v3rerank] Cross-encoder reranker (graceful fallback if not installed)
try:
    from sentence_transformers import CrossEncoder
    _CE_AVAILABLE = True
except ImportError:
    _CE_AVAILABLE = False

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

if not _BM25_AVAILABLE:
    logger.warning("rank-bm25 not installed -- falling back to vector-only retrieval.")
if not _CE_AVAILABLE:
    logger.warning("sentence-transformers not installed -- reranking disabled.")

logging.getLogger("chromadb.telemetry.product.posthog").setLevel(logging.CRITICAL)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
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 = 512
CHUNK_OVERLAP = 50
TOP_K_RESULTS = 2

# [v2bm25] hybrid retrieval candidates before fusion
RRF_CANDIDATES = 20
RRF_K = 60

# [v3rerank] number of candidates fed to the cross-encoder before final cut
RERANK_INPUT_K = 15
CROSS_ENCODER_MODEL = "BAAI/bge-reranker-v2-m3"

# [v3routing] question classification threshold (words)
SHORT_Q_MAX_WORDS = 25

# [v4guardrail] Abstention par score cross-encoder (sigmoide, 0-1) du MEILLEUR
# passage. En dessous du seuil -> question hors corpus -> abstention.
RERANK_ABSTAIN_THRESHOLD = 0.15
ABSTAIN_ANSWER = "Le contexte fourni ne permet pas de repondre a cette question."

# [v4qcm] Directive ajoutee au prompt quand QCM + contexte pertinent.
QCM_DIRECTIVE = (
    "\n\nINSTRUCTION SUPPLEMENTAIRE : Cette question est un QCM et le contexte "
    "fourni est juge pertinent. Tu DOIS choisir la meilleure reponse parmi les "
    "options proposees en t'appuyant sur le contexte, meme si la reponse n'y est "
    "pas formulee explicitement : raisonne a partir des elements disponibles et "
    "tranche. Commence 'answer' par la ou les lettres correctes."
)
# explicit "weak answer" markers that trigger escalation to the large model
_WEAK_MARKERS = [
    "ne permet pas de répondre", "ne permet pas de repondre",
    "don't have enough", "do not have enough", "i don't have enough",
    "cannot answer", "i cannot answer",
    "je ne sais pas",
]

_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)."
    )
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 (large / default) -- 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)

# [v3routing] LLM small (mini) -- optional. If absent, routing is disabled.
_llm_small = _config.get("llm_small")
if _llm_small:
    LLM_SMALL_ENDPOINT_URL = _llm_small["endpoint_url"]
    LLM_SMALL_MODEL_NAME = _llm_small["model"]
    LLM_SMALL_MAX_TOKENS = _llm_small.get("max_completion_tokens", 512)
    LLM_SMALL_TEMPERATURE = _llm_small.get("temperature", 0.7)
    LLM_SMALL_TOP_P = _llm_small.get("top_p", 0.95)
    _ROUTING_ENABLED = True
    logger.info(f"Routing enabled: small={LLM_SMALL_MODEL_NAME}, large={LLM_MODEL_NAME}")
else:
    _ROUTING_ENABLED = False
    logger.info("No llm_small in config -- routing disabled, large model only.")

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_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")

# ---------------------------------------------------------------------------
# Vector Store (ChromaDB)
# ---------------------------------------------------------------------------
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()}")

# ---------------------------------------------------------------------------
# [v2bm25] BM25 keyword index
# ---------------------------------------------------------------------------
_bm25_index = None
_bm25_ids: list[str] = []
_bm25_docs: list[str] = []
_bm25_metas: list[dict] = []


def _tokenize(text: str) -> list[str]:
    return re.findall(r"\w+", text.lower())


def build_bm25_index() -> None:
    global _bm25_index, _bm25_ids, _bm25_docs, _bm25_metas
    if not _BM25_AVAILABLE:
        return
    try:
        # Lecture PAGINEE : un collection.get() global depasse la limite SQLite
        # "too many SQL variables" sur un gros corpus (84k chunks) -> BM25 jamais construit.
        total = collection.count()
        _bm25_ids, _bm25_docs, _bm25_metas = [], [], []
        _off, _step = 0, 2000
        while _off < total:
            _batch = collection.get(include=["documents", "metadatas"], limit=_step, offset=_off)
            _bm25_ids.extend(_batch.get("ids", []) or [])
            _bm25_docs.extend(_batch.get("documents", []) or [])
            _bm25_metas.extend(_batch.get("metadatas", []) or [])
            _off += _step
        if not _bm25_docs:
            _bm25_index = None
            logger.info("BM25 index empty (no documents in store yet).")
            return
        tokenized_corpus = [_tokenize(d) for d in _bm25_docs]
        _bm25_index = BM25Okapi(tokenized_corpus)
        logger.info(f"BM25 index built over {len(_bm25_docs)} chunks.")
    except Exception as e:
        _bm25_index = None
        logger.warning(f"Failed to build BM25 index: {e}")


# ---------------------------------------------------------------------------
# [v3rerank] Cross-encoder reranker
# ---------------------------------------------------------------------------
_cross_encoder = None


def load_cross_encoder() -> None:
    global _cross_encoder
    if not _CE_AVAILABLE:
        return
    try:
        logger.info(f"Loading cross-encoder '{CROSS_ENCODER_MODEL}' (first run downloads the model)...")
        _cross_encoder = CrossEncoder(CROSS_ENCODER_MODEL, max_length=512)
        logger.info("Cross-encoder ready.")
    except Exception as e:
        _cross_encoder = None
        logger.warning(f"Failed to load cross-encoder: {e}")


def rerank_contexts(query: str, contexts: list[dict], top_k: int) -> list[dict]:
    """Rerank candidate chunks with the cross-encoder, keep the top_k best."""
    if _cross_encoder is None or len(contexts) <= 1:
        return contexts[:top_k]
    try:
        pairs = [[query, c["text"]] for c in contexts]
        scores = _cross_encoder.predict(pairs)
        for c, s in zip(contexts, scores):
            c["rerank_score"] = float(s)
        contexts.sort(key=lambda c: c.get("rerank_score", 0.0), reverse=True)
    except Exception as e:
        logger.warning(f"Reranking failed, falling back to fusion order: {e}")
    return contexts[:top_k]


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


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

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()}")
    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]:
    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]]:
    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):
        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:
    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()}")
    build_bm25_index()  # [v2bm25] keep keyword index in sync
    return len(documents)


def retrieve_relevant_context(query: str, top_k: int = TOP_K_RESULTS) -> list[dict]:
    """
    Hybrid retrieval (vector + BM25, RRF) -> optional cross-encoder rerank -> top_k.
    """
    n = collection.count()
    if n == 0:
        return []

    candidate_k = min(max(top_k * 5, RRF_CANDIDATES), n)

    # --- 1) Vector search (semantic) ---
    query_embedding = generate_embeddings([query])[0]
    vres = collection.query(
        query_embeddings=[query_embedding],
        n_results=candidate_k,
        include=["documents", "metadatas", "distances"],
    )
    v_ids = vres["ids"][0]
    v_docs = vres["documents"][0]
    v_metas = vres["metadatas"][0]

    doc_lookup: dict[str, tuple[str, str]] = {}
    vector_rank: dict[str, int] = {}
    for rank, (did, dtext, dmeta) in enumerate(zip(v_ids, v_docs, v_metas)):
        vector_rank[did] = rank
        doc_lookup[did] = (dtext, (dmeta or {}).get("source", "unknown"))

    # --- 2) BM25 search (keywords) ---
    bm25_rank: dict[str, int] = {}
    if _bm25_index is not None:
        scores = _bm25_index.get_scores(_tokenize(query))
        top_positions = sorted(range(len(scores)), key=lambda p: scores[p], reverse=True)[:candidate_k]
        for rank, pos in enumerate(top_positions):
            did = _bm25_ids[pos]
            bm25_rank[did] = rank
            if did not in doc_lookup:
                doc_lookup[did] = (_bm25_docs[pos], (_bm25_metas[pos] or {}).get("source", "unknown"))

    # --- 3) Reciprocal Rank Fusion ---
    fused: dict[str, float] = {}
    for did, rank in vector_rank.items():
        fused[did] = fused.get(did, 0.0) + 1.0 / (RRF_K + rank + 1)
    for did, rank in bm25_rank.items():
        fused[did] = fused.get(did, 0.0) + 1.0 / (RRF_K + rank + 1)

    # keep more candidates if a reranker will refine them
    pre_k = RERANK_INPUT_K if _cross_encoder is not None else top_k
    ranked = sorted(fused.items(), key=lambda x: x[1], reverse=True)[:pre_k]

    contexts = []
    for did, fused_score in ranked:
        text, source = doc_lookup[did]
        contexts.append({"text": text, "source": source, "similarity_score": round(fused_score, 4)})

    # --- 4) [v3rerank] Cross-encoder reranking ---
    if _cross_encoder is not None:
        contexts = rerank_contexts(query, contexts, top_k)
    else:
        contexts = contexts[:top_k]

    return contexts


def build_rag_prompt(query: str, contexts: list[dict]) -> str:
    context_text = "\n\n".join(f"[Source: {ctx['source']}]\n{ctx['text']}" for ctx in contexts)
    return RAG_PROMPT_TEMPLATE.format(context=context_text, question=query)


# ---------------------------------------------------------------------------
# [v3routing] Question classification + model calls
# ---------------------------------------------------------------------------

_QCM_OPTION_RE = re.compile(r"(?:^|\s)[A-Da-d]\s*[\)\.\-]")  # "A)" "B." "c -" ...
_QCM_KEYWORDS = ["parmi les", "laquelle", "lesquelles", "vrai ou faux",
                 "cochez", "réponse correcte", "reponse correcte",
                 "proposition", "qcm"]


def is_qcm(query: str) -> bool:
    q = query.strip()
    return len(_QCM_OPTION_RE.findall(q)) >= 2 or any(k in q.lower() for k in _QCM_KEYWORDS)


def _top_rerank_norm(contexts: list):
    """Normalised (sigmoid) score of the best reranked passage, or None."""
    if not contexts:
        return None
    s = contexts[0].get("rerank_score")
    if s is None:
        return None
    try:
        return 1.0 / (1.0 + math.exp(-float(s)))
    except OverflowError:
        return 0.0 if s < 0 else 1.0


def classify_question(query: str) -> str:
    """
    Return "short" (QCM / short closed question -> mini first)
    or     "long"  (open question -> large model directly).
    """
    q = query.strip()
    low = q.lower()

    # QCM markers => treat as short/closed
    option_hits = len(_QCM_OPTION_RE.findall(q))
    if option_hits >= 2 or any(k in low for k in _QCM_KEYWORDS):
        return "short"

    # otherwise decide by length
    if len(q.split()) <= SHORT_Q_MAX_WORDS:
        return "short"
    return "long"


def _parse_llm_json(raw_content: str):
    """Return (answer, explanation, parsed_ok)."""
    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)
        return parsed["answer"], parsed["explanation"], True
    except (json.JSONDecodeError, KeyError, TypeError):
        return raw_content, "LLM did not return a structured explanation.", False


def _is_weak_answer(answer_text: str, parsed_ok: bool) -> bool:
    if not parsed_ok:
        return True
    if not answer_text or not str(answer_text).strip():
        return True
    low = str(answer_text).lower()
    return any(m in low for m in _WEAK_MARKERS)


def _call_model(prompt: str, which: str) -> dict:
    """which = 'small' or 'large'."""
    if which == "small" and _ROUTING_ENABLED:
        return call_llm_with_metrics(
            prompt, endpoint_url=LLM_SMALL_ENDPOINT_URL, api_key=AZURE_API_KEY,
            model=LLM_SMALL_MODEL_NAME, max_completion_tokens=LLM_SMALL_MAX_TOKENS,
            temperature=LLM_SMALL_TEMPERATURE, top_p=LLM_SMALL_TOP_P,
        )
    return 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,
    )


def _empty_tokens() -> dict:
    return {"prompt": 0, "completion": 0, "cached": 0, "total": 0}


def _add_tokens(a: dict, b: dict) -> dict:
    return {k: (a.get(k, 0) or 0) + (b.get(k, 0) or 0) for k in ("prompt", "completion", "cached", "total")}


def rag_query(query: str, top_k: int = TOP_K_RESULTS) -> dict:
    start_time = time.perf_counter()

    contexts = retrieve_relevant_context(query, top_k=top_k)

    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.",
            "total_token": 0, "prompt_tokens": 0, "completion_tokens": 0, "cached_tokens": 0,
            "co2_grams": None, "energy_kwh": None, "run_time_in_ms": elapsed_ms,
            "model_used": "none", "question_type": "n/a",
        }

    prompt = build_rag_prompt(query, contexts)

    # [v4guardrail] Abstention si contexte trop peu pertinent (memoire absent)
    top_score = _top_rerank_norm(contexts)
    if top_score is not None:
        logger.info(f"top_rerank={top_score:.3f} | qcm={is_qcm(query)} | q={query[:60]!r}")
        if top_score < RERANK_ABSTAIN_THRESHOLD:
            elapsed_ms = round((time.perf_counter() - start_time) * 1000, 2)
            return {
                "answer": ABSTAIN_ANSWER,
                "sources": [{"source": c["source"], "score": c.get("rerank_score", c["similarity_score"]),
                             "ref_text": c["text"]} for c in contexts],
                "explanation": "Aucun passage suffisamment pertinent : question hors corpus.",
                "total_token": 0, "prompt_tokens": 0, "completion_tokens": 0, "cached_tokens": 0,
                "co2_grams": None, "energy_kwh": None, "run_time_in_ms": elapsed_ms,
                "model_used": "abstained", "question_type": "abstain",
                "top_rerank": round(top_score, 3),
            }

    # [v4qcm] QCM + contexte pertinent -> forcer le modele a trancher
    if is_qcm(query):
        prompt = prompt + QCM_DIRECTIVE

    # [v3routing] decide path
    qtype = classify_question(query) if _ROUTING_ENABLED else "long"
    tokens = _empty_tokens()
    co2 = 0.0
    energy = 0.0
    models_used = []

    def _accumulate(res):
        nonlocal tokens, co2, energy
        tokens = _add_tokens(tokens, res.get("tokens", {}))
        if isinstance(res.get("co2_grams"), (int, float)):
            co2 += res["co2_grams"]
        if isinstance(res.get("energy_kwh"), (int, float)):
            energy += res["energy_kwh"]

    if qtype == "short" and _ROUTING_ENABLED:
        # mini first
        small_res = _call_model(prompt, "small")
        _accumulate(small_res)
        models_used.append(LLM_SMALL_MODEL_NAME)
        answer, explanation, parsed_ok = _parse_llm_json(small_res["content"])

        if _is_weak_answer(answer, parsed_ok):
            # escalate to large
            large_res = _call_model(prompt, "large")
            _accumulate(large_res)
            models_used.append(LLM_MODEL_NAME)
            answer, explanation, _ = _parse_llm_json(large_res["content"])
    else:
        # long/open question -> large directly (no double billing)
        large_res = _call_model(prompt, "large")
        _accumulate(large_res)
        models_used.append(LLM_MODEL_NAME)
        answer, explanation, _ = _parse_llm_json(large_res["content"])

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

    return {
        "answer": answer,
        "sources": [{"source": c["source"], "score": c.get("rerank_score", c["similarity_score"]),
                     "ref_text": c["text"]} for c in contexts],
        "explanation": explanation,
        "total_token": tokens["total"],
        "prompt_tokens": tokens["prompt"],
        "completion_tokens": tokens["completion"],
        "cached_tokens": tokens["cached"],
        "co2_grams": co2 if co2 else None,
        "energy_kwh": energy if energy else None,
        "run_time_in_ms": elapsed_ms,
        "model_used": " -> ".join(models_used),   # e.g. "gpt-5-mini -> gpt-5.1" if escalated
        "question_type": qtype,
        "top_rerank": round(top_score, 3) if top_score is not None else None,
    }


# ---------------------------------------------------------------------------
# Ingest Train Documents (on-demand)
# ---------------------------------------------------------------------------

def ingest_train_documents():
    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
    for file_path in TRAIN_DOCS_DIR.rglob("*.txt"):
        text = file_path.read_text(encoding="utf-8")
        add_documents_to_vectorstore(chunk_text(text, source=file_path.name))
    for file_path in TRAIN_DOCS_DIR.rglob("*.pdf"):
        text = extract_text_from_pdf(file_path)
        if text.strip():
            add_documents_to_vectorstore(chunk_text(text, source=file_path.name))
        else:
            logger.warning(f"No extractable text found in: {file_path.name}")
    logger.info(f"Train document ingestion complete. Total chunks: {collection.count()}")


# Build indexes at startup from existing ChromaDB chunks
build_bm25_index()      # [v2bm25]
load_cross_encoder()    # [v3rerank]


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


class QueryRequest(BaseModel):
    query: str
    top_k: Optional[int] = TOP_K_RESULTS


class IngestRequest(BaseModel):
    text: str
    source: str = "user_upload"


@app.post("/query")
async def query_endpoint(request: QueryRequest):
    result = rag_query(request.query, top_k=request.top_k)
    return JSONResponse(content=result)


@app.post("/ingest")
async def ingest_endpoint(request: IngestRequest):
    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():
    return {
        "status": "healthy",
        "documents_in_store": collection.count(),
        "embedding_model": EMBEDDING_MODEL_NAME,
        "llm_model": LLM_MODEL_NAME,
        "bm25_enabled": _bm25_index is not None,
        "reranker_enabled": _cross_encoder is not None,
        "routing_enabled": _ROUTING_ENABLED,
    }


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

def gradio_query(question: str) -> tuple[str, str, str, str, str]:
    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"])
    routing_info = f"\n\n🔀 Modèle: {result.get('model_used','?')} (type: {result.get('question_type','?')})"
    answer = f"{result['answer']}\n\n📚 Sources:\n{sources_text}{routing_info}" 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:
    if not text.strip():
        return "Please provide text to ingest."
    count = add_documents_to_vectorstore(chunk_text(text, source=source_name or "user_upload"))
    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
    Pipeline avancé : recherche hybride (vectoriel + BM25) → reranking cross-encoder → routing mini/gros.
    **API Endpoint:** `POST /query` avec `{"query": "votre question"}`.
    ---
    """)

    with gr.Tab("💬 Chat"):
        gr.Markdown("Posez une question sur les mémoires ingérés.")
        with gr.Row():
            query_input = gr.Textbox(label="Your Question", placeholder="e.g., Qu'est-ce que le risk adjustment ?", 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}
        ```
        ### GET /health
        Returns system health, document count, and which features are active
        (bm25, reranker, routing).
        """)

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

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)