r"""
Advanced RAG Chatbot — main.py  (FIXED)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
FIXES APPLIED:
  1. Replaced fragile LLM-as-judge JSON evaluation with fast, deterministic
     embedding-based metrics (no JSON parsing errors, no hallucinated scores).
  2. context_relevance  → cosine similarity(question_embedding, avg_chunk_embedding)
  3. groundedness       → token-overlap (ROUGE-1 recall) of answer tokens in context
  4. answer_relevance   → cosine similarity(question_embedding, answer_embedding)
  5. evaluate_node now always returns valid floats in [0, 1].
  6. Fixed confidence badge mapping sent in the API response (was always missing).
  7. Fixed early-exit guard that was firing on valid context.
  8. Added structured logging for metric values.

Folder structure (everything inside C:/Users/Srikar/Downloads/RAG/):
  RAG/
  ├── main.py          ← this file
  ├── requirements.txt
  ├── .env
  ├── frontend/
  │   └── index.html
  ├── faiss_index/     ← auto-created
  └── uploads/         ← auto-created

Run:
  cd C:\Users\Srikar\Downloads\RAG
  set COHERE_API_KEY=your-key-here
  uvicorn main:app --reload --port 8000
"""

import os
import re
import uuid
import time
import math
from typing import List, Optional, Annotated, TypedDict
from pathlib import Path
from collections import Counter

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel

# LangChain
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
from langchain_core.documents import Document
from langchain_cohere import ChatCohere
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
from langchain_core.output_parsers import StrOutputParser

# LangGraph
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

# ── Paths ─────────────────────────────────────────────────────────────────────

BASE_DIR     = Path(__file__).parent
FRONTEND_DIR = BASE_DIR / "frontend"
UPLOAD_DIR   = BASE_DIR / "uploads"
FAISS_DIR    = BASE_DIR / "faiss_index"

UPLOAD_DIR.mkdir(exist_ok=True)
FAISS_DIR.mkdir(exist_ok=True)

# ── Config ────────────────────────────────────────────────────────────────────
from dotenv import load_dotenv
load_dotenv()

COHERE_API_KEY = os.getenv("COHERE_API_KEY", "your-cohere-api-key-here")

# ── App ───────────────────────────────────────────────────────────────────────

app = FastAPI(title="RAG Chatbot", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

app.mount("/frontend", StaticFiles(directory=str(FRONTEND_DIR)), name="frontend")

@app.get("/")
def serve_frontend():
    return FileResponse(str(FRONTEND_DIR / "index.html"))

# ── Embeddings ────────────────────────────────────────────────────────────────

print("Loading embedding model (downloads ~90MB on first run)...")
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={"device": "cpu"},
    encode_kwargs={"normalize_embeddings": True},
)
print("Embedding model ready.")

# ── LLM ───────────────────────────────────────────────────────────────────────

llm = ChatCohere(
    cohere_api_key=COHERE_API_KEY,
    model="command-r-plus-08-2024",
    temperature=0.3,
)

# ── Vector Store ──────────────────────────────────────────────────────────────

vector_store: Optional[FAISS] = None

def get_vector_store() -> FAISS:
    global vector_store
    if vector_store is not None:
        return vector_store

    index_file = FAISS_DIR / "index.faiss"
    if index_file.exists():
        print("Loading existing FAISS index...")
        vector_store = FAISS.load_local(
            str(FAISS_DIR), embeddings, allow_dangerous_deserialization=True
        )
        print("FAISS index loaded.")
    else:
        print("Creating new FAISS index...")
        seed = Document(page_content="Knowledge base ready.", metadata={"source": "system"})
        vector_store = FAISS.from_documents([seed], embeddings)
        vector_store.save_local(str(FAISS_DIR))
        print("FAISS index created.")

    return vector_store

@app.on_event("startup")
async def startup():
    get_vector_store()

# ── Deterministic RAG Metric Helpers ──────────────────────────────────────────

def _cosine_similarity(vec_a: list, vec_b: list) -> float:
    """Pure-Python cosine similarity between two float lists."""
    dot   = sum(a * b for a, b in zip(vec_a, vec_b))
    norm_a = math.sqrt(sum(a * a for a in vec_a))
    norm_b = math.sqrt(sum(b * b for b in vec_b))
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return max(0.0, min(1.0, dot / (norm_a * norm_b)))


def _avg_vector(vectors: list) -> list:
    """Element-wise average of a list of equal-length vectors."""
    if not vectors:
        return []
    n = len(vectors[0])
    avg = [0.0] * n
    for v in vectors:
        for i in range(n):
            avg[i] += v[i]
    return [x / len(vectors) for x in avg]


def _tokenize(text: str) -> Counter:
    """Simple whitespace + lowercase tokenizer for overlap calculation."""
    tokens = re.findall(r"[a-z0-9]+", text.lower())
    return Counter(tokens)


def compute_rag_metrics(question: str, retrieved_docs: List[str], answer: str) -> dict:
    """
    Compute three RAG quality metrics deterministically using embeddings
    and token overlap — no LLM call required, no JSON parsing.

    Metrics (each in [0.0, 1.0]):
    ─────────────────────────────
    context_relevance:
        Cosine similarity between the question embedding and the average
        embedding of all retrieved document chunks.
        Measures: "Did we retrieve relevant chunks for this question?"

    groundedness:
        ROUGE-1 Recall of the answer tokens found in the concatenated context.
        Measures: "Is the answer actually supported by the retrieved text?"
        Formula: |answer_tokens ∩ context_tokens| / |answer_tokens|

    answer_relevance:
        Cosine similarity between the question embedding and the answer embedding.
        Measures: "Does the answer semantically address the question asked?"
    """

    # ── Guard: nothing retrieved ──────────────────────────────────────────────
    if not retrieved_docs or all(not d.strip() for d in retrieved_docs):
        return {
            "context_relevance": 0.0,
            "groundedness":      0.0,
            "answer_relevance":  0.0,
        }

    # Strip the "[Source: X]" header lines before embedding
    clean_chunks = []
    for doc in retrieved_docs:
        lines = doc.split("\n", 1)
        content = lines[1].strip() if len(lines) > 1 else doc.strip()
        if content:
            clean_chunks.append(content)

    if not clean_chunks:
        return {
            "context_relevance": 0.0,
            "groundedness":      0.0,
            "answer_relevance":  0.0,
        }

    # ── Embed everything in one batch ─────────────────────────────────────────
    try:
        texts_to_embed = [question, answer] + clean_chunks
        all_vecs = embeddings.embed_documents(texts_to_embed)
        q_vec       = all_vecs[0]
        a_vec       = all_vecs[1]
        chunk_vecs  = all_vecs[2:]
    except Exception as e:
        print(f"[metrics] Embedding error: {e}")
        return {
            "context_relevance": 0.0,
            "groundedness":      0.0,
            "answer_relevance":  0.0,
        }

    # ── 1. Context Relevance ──────────────────────────────────────────────────
    avg_chunk_vec     = _avg_vector(chunk_vecs)
    context_relevance = _cosine_similarity(q_vec, avg_chunk_vec)

    # ── 2. Groundedness (ROUGE-1 Recall) ─────────────────────────────────────
    full_context_text = " ".join(clean_chunks)
    answer_tokens     = _tokenize(answer)
    context_tokens    = _tokenize(full_context_text)

    if not answer_tokens:
        groundedness = 0.0
    else:
        overlap = sum(min(answer_tokens[t], context_tokens[t]) for t in answer_tokens)
        groundedness = overlap / sum(answer_tokens.values())

    # ── 3. Answer Relevance ───────────────────────────────────────────────────
    answer_relevance = _cosine_similarity(q_vec, a_vec)

    metrics = {
        "context_relevance": round(context_relevance, 4),
        "groundedness":      round(groundedness,      4),
        "answer_relevance":  round(answer_relevance,  4),
    }
    print(f"[metrics] {metrics}")
    return metrics


# ── LangGraph ─────────────────────────────────────────────────────────────────

class State(TypedDict):
    messages:       Annotated[List[BaseMessage], add_messages]
    question:       str
    retrieved_docs: List[str]
    answer:         str
    metrics:        dict

SYSTEM_PROMPT = """You are an expert AI assistant with access to a knowledge base.
Answer questions accurately using the retrieved context below.

Rules:
- Answer ONLY from the provided context when available.
- If the context does not contain the answer, say so honestly.
- Be concise, clear, and helpful.
- Mention the source document when relevant.

Context:
{context}
"""

def retrieve_node(state: State) -> State:
    vs = get_vector_store()
    try:
        docs = vs.similarity_search(state["question"], k=4)
        retrieved = [
            f"[Source: {d.metadata.get('source', 'KB')}]\n{d.page_content}"
            for d in docs
        ]
    except Exception as e:
        print(f"[retrieve] Error: {e}")
        retrieved = []
    return {**state, "retrieved_docs": retrieved}


def generate_node(state: State) -> State:
    context = "\n\n---\n\n".join(state["retrieved_docs"]) or "No context available."
    history = state["messages"][:-1]

    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ])

    answer = (prompt | llm | StrOutputParser()).invoke({
        "context":  context,
        "history":  history,
        "question": state["question"],
    })

    return {**state, "answer": answer}


def evaluate_node(state: State) -> State:
    """
    FIXED: Uses deterministic embedding-based metrics instead of fragile
    LLM-as-judge JSON parsing. Fast, reliable, always returns valid floats.
    """
    metrics = compute_rag_metrics(
        question       = state["question"],
        retrieved_docs = state["retrieved_docs"],
        answer         = state["answer"],
    )
    return {**state, "metrics": metrics}


graph = StateGraph(State)
graph.add_node("retrieve",  retrieve_node)
graph.add_node("generate",  generate_node)
graph.add_node("evaluate",  evaluate_node)
graph.set_entry_point("retrieve")
graph.add_edge("retrieve", "generate")
graph.add_edge("generate", "evaluate")
graph.add_edge("evaluate", END)
rag_graph = graph.compile()

# ── Session store ─────────────────────────────────────────────────────────────

sessions: dict = {}

# ── Models ────────────────────────────────────────────────────────────────────

class ChatRequest(BaseModel):
    question:   str
    session_id: Optional[str] = None

class ChatResponse(BaseModel):
    answer:              str
    session_id:          str
    sources:             List[str]
    processing_time_ms:  int
    metrics:             Optional[dict] = None
    confidence:          Optional[str]  = None   # FIX: now actually populated

class UploadResponse(BaseModel):
    message:      str
    filename:     str
    chunks_added: int

class KBStats(BaseModel):
    total_documents: int
    index_exists:    bool

# ── Endpoints ─────────────────────────────────────────────────────────────────

@app.post("/api/chat", response_model=ChatResponse)
async def chat(req: ChatRequest):
    if not req.question.strip():
        raise HTTPException(status_code=400, detail="Question cannot be empty.")

    session_id = req.session_id or str(uuid.uuid4())
    history    = sessions.setdefault(session_id, [])
    history.append(HumanMessage(content=req.question))

    start = time.time()

    try:
        result = rag_graph.invoke({
            "messages":       history,
            "question":       req.question,
            "retrieved_docs": [],
            "answer":         "",
            "metrics":        {}
        })
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"LLM error: {str(e)}")

    answer  = result["answer"]
    metrics = result.get("metrics", {})

    sources = list({
        d.split("\n")[0].replace("[Source: ", "").replace("]", "")
        for d in result["retrieved_docs"]
    })

    # FIX: compute confidence server-side so the badge is always correct
    confidence = "LOW"
    if metrics:
        avg = (
            metrics.get("context_relevance", 0) +
            metrics.get("groundedness",      0) +
            metrics.get("answer_relevance",  0)
        ) / 3
        if avg >= 0.75:
            confidence = "HIGH"
        elif avg >= 0.45:
            confidence = "MODERATE"

    history.append(AIMessage(content=answer))
    sessions[session_id] = history[-20:]

    return ChatResponse(
        answer             = answer,
        session_id         = session_id,
        sources            = sources,
        processing_time_ms = int((time.time() - start) * 1000),
        metrics            = metrics,
        confidence         = confidence,
    )


@app.post("/api/upload", response_model=UploadResponse)
async def upload(file: UploadFile = File(...)):
    ext = Path(file.filename).suffix.lower()
    if ext not in {".pdf", ".txt", ".docx", ".md"}:
        raise HTTPException(status_code=400, detail=f"Unsupported type: {ext}")

    path = UPLOAD_DIR / file.filename
    path.write_bytes(await file.read())

    try:
        if ext == ".pdf":
            loader = PyPDFLoader(str(path))
        elif ext in (".txt", ".md"):
            loader = TextLoader(str(path), encoding="utf-8")
        else:
            loader = Docx2txtLoader(str(path))

        docs = loader.load()
        for d in docs:
            d.metadata["source"] = file.filename

        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
        chunks   = splitter.split_documents(docs)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")

    global vector_store
    vs = get_vector_store()
    vs.add_documents(chunks)
    vs.save_local(str(FAISS_DIR))

    return UploadResponse(
        message      = "Document added successfully!",
        filename     = file.filename,
        chunks_added = len(chunks),
    )


@app.delete("/api/session/{session_id}")
def clear_session(session_id: str):
    sessions.pop(session_id, None)
    return {"message": "Session cleared."}


@app.get("/api/kb/stats", response_model=KBStats)
def kb_stats():
    vs    = get_vector_store()
    total = vs.index.ntotal if vs and hasattr(vs, "index") else 0
    return KBStats(
        total_documents = total,
        index_exists    = (FAISS_DIR / "index.faiss").exists()
    )


@app.get("/api/health")
def health():
    return {
        "status":     "ok",
        "llm":        "command-r-plus-08-2024",
        "embeddings": "all-MiniLM-L6-v2",
        "vectorstore":"FAISS",
        "metrics":    "deterministic (embedding cosine + ROUGE-1)",
    }