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

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 = 512
CHUNK_OVERLAP = 50
TOP_K_RESULTS = 3

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

# ---------------------------------------------------------------------------
# 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 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 retrieve_relevant_context(query: str, top_k: int = TOP_K_RESULTS) -> list[dict]:
    """
    Retrieve the most relevant document chunks for a given query.

    This is the "Retrieval" step in RAG:
    1. Embed the user's query using the same embedding model
    2. Search the vector store for the nearest neighbors
    3. Return the top-k most similar chunks as context
    """
    if collection.count() == 0:
        return []

    query_embedding = generate_embeddings([query])[0]

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=min(top_k, collection.count()),
        include=["documents", "metadatas", "distances"],
    )

    contexts = []
    for i in range(len(results["documents"][0])):
        contexts.append({
            "text": results["documents"][0][i],
            "source": results["metadatas"][0][i]["source"],
            "similarity_score": 1 - results["distances"][0][i],
        })

    return contexts


def call_llm(prompt: str) -> dict:
    """
    Make a call to the LLM via Azure Foundry (GPT-5).

    Calls the Azure OpenAI-compatible chat/completions endpoint.
    Endpoint URL is loaded from config.json; API key from AZURE_API_KEY env var.

    Returns a dict with 'content' (str) and 'total_tokens' (int).
    """
    headers = {
        "api-key": AZURE_API_KEY,
        "Content-Type": "application/json",
    }
    payload = {
        "model": LLM_MODEL_NAME,
        "messages": [{"role": "user", "content": prompt}],
        "max_completion_tokens": LLM_MAX_TOKENS,
        "temperature": LLM_TEMPERATURE,
        "top_p": LLM_TOP_P,
    }
    try:
        resp = http_requests.post(LLM_ENDPOINT_URL, headers=headers, json=payload, timeout=None)
        resp.raise_for_status()
        data = resp.json()
        content = data["choices"][0]["message"]["content"].strip()
        total_tokens = data.get("usage", {}).get("total_tokens", 0)
        return {"content": content, "total_tokens": total_tokens}
    except http_requests.exceptions.HTTPError as e:
        logger.error(f"LLM API call failed: {e} — {resp.text}")
        raise HTTPException(status_code=503, detail=f"LLM service unavailable: {str(e)}")
    except (http_requests.exceptions.JSONDecodeError, ValueError) as e:
        logger.error(f"LLM API returned non-JSON response (status {resp.status_code}): {repr(resp.text)}")
        raise HTTPException(status_code=502, detail="LLM service returned an invalid response")
    except (KeyError, IndexError) as e:
        logger.error(f"Unexpected LLM response format: {e} — body: {resp.text}")
        raise HTTPException(status_code=502, detail="Unexpected response from LLM service")


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: int = TOP_K_RESULTS) -> dict:
    """
    End-to-end RAG pipeline: Query → Retrieve → Generate.

    This demonstrates Step 2d: How to make a RAG system end-to-end.

    Pipeline steps:
    1. Receive user query
    2. Retrieve relevant context from vector store
    3. Build augmented prompt with context
    4. Call LLM to generate answer
    5. Return answer with source metadata, explanation, token count, and timing
    """
    start_time = time.perf_counter()

    # Step 1: Retrieve relevant document chunks
    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 to retrieve context from.",
            "total_token": 0,
            "run_time_in_ms": elapsed_ms,
        }

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

    # Step 3: Generate answer from LLM
    llm_result = call_llm(prompt)
    raw_content = llm_result["content"]
    total_token = llm_result["total_tokens"]

    # 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,
        "total_token": total_token,
        "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_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 - 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
        - 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]:
    """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))
    run_time = f"{result.get('run_time_in_ms', 0)} ms"
    return answer, explanation, token_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_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_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)