# app.py — KUMPAS RAG Server for Hugging Face Spaces
# Stack: FAISS + HuggingFace all-MiniLM-L6-v2 + Groq llama-3.3-70b

import os
import logging
from contextlib import asynccontextmanager
from typing import List, Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

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

# ── Global chain (loaded once at startup) ───────────────────────
qa_chain = None

def build_chain():
    """
    Loads the FAISS index from disk and builds the QA chain.
    Called once when the server starts.
    """
    global qa_chain

    groq_api_key = os.environ.get("GROQ_API_KEY", "")
    if not groq_api_key:
        raise RuntimeError("GROQ_API_KEY environment variable is not set!")

    faiss_path = os.environ.get("FAISS_PATH", "student_manual_faiss")
    text_path  = os.environ.get("TEXT_PATH",  "harrypotter.txt")

    logger.info("Loading embeddings model...")
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")

    # ── Try loading saved FAISS index first (fast) ──────────────
    if os.path.exists(faiss_path):
        logger.info(f"Loading FAISS index from {faiss_path}...")
        vectorstore = FAISS.load_local(
            faiss_path,
            embeddings,
            allow_dangerous_deserialization=True
        )
        logger.info("FAISS index loaded from disk ✅")

    # ── Otherwise build from text file (slow, first run only) ───
    elif os.path.exists(text_path):
        logger.info(f"Building FAISS index from {text_path} (this takes a few minutes)...")
        from langchain_text_splitters import RecursiveCharacterTextSplitter

        try:
            with open(text_path, "r", encoding="latin-1") as f:
                text = f.read()
        except UnicodeDecodeError:
            with open(text_path, "r", encoding="cp1252") as f:
                text = f.read()

        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
        chunks   = splitter.split_text(text)
        logger.info(f"Created {len(chunks)} chunks")

        vectorstore = FAISS.from_texts(chunks, embedding=embeddings)
        vectorstore.save_local(faiss_path)
        logger.info("FAISS index built and saved ✅")

    else:
        raise RuntimeError(
            f"Neither FAISS index ({faiss_path}) nor text file ({text_path}) found! "
            "Make sure harrypotter.txt is in your Space files."
        )

    # ── Build retriever ──────────────────────────────────────────
    retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 3}
    )

    # ── Groq LLM ────────────────────────────────────────────────
    llm = ChatGroq(
        model="llama-3.3-70b-versatile",
        temperature=0.5,
        groq_api_key=groq_api_key
    )

    # ── Prompt (same as your Colab notebook) ────────────────────
    prompt = ChatPromptTemplate.from_template("""
Answer the question based only on the following context from the Harry Potter series:
{context}

Question: {question}

Answer clearly and concisely. If the answer is not in the context, say so honestly.
""")

    # ── LCEL chain (identical to your Colab) ────────────────────
    qa_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    logger.info("QA chain ready ✅")
    return qa_chain


# ── Lifespan: build chain on startup ────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("🔮 Starting KUMPAS RAG server...")
    build_chain()
    logger.info("🔮 Server ready!")
    yield
    logger.info("Server shutting down.")


# ── FastAPI app ──────────────────────────────────────────────────
app = FastAPI(
    title="KUMPAS Harry Potter RAG API",
    description="Wizarding Oracle — FAISS + Groq LLaMA",
    version="1.0.0",
    lifespan=lifespan
)

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


# ── Request / Response models ────────────────────────────────────
class HistoryItem(BaseModel):
    role: str
    content: str

class AskRequest(BaseModel):
    question: str
    history: Optional[List[HistoryItem]] = []

class AskResponse(BaseModel):
    answer: str


# ── Endpoints ────────────────────────────────────────────────────
@app.get("/")
def root():
    return {"message": "🔮 KUMPAS Wizarding Oracle is alive!", "status": "ok"}

@app.get("/health")
def health():
    return {
        "status": "ok",
        "oracle": "ready" if qa_chain else "loading",
        "model": "llama-3.3-70b-versatile"
    }

@app.post("/ask", response_model=AskResponse)
def ask(req: AskRequest):
    if qa_chain is None:
        raise HTTPException(status_code=503, detail="Oracle is still initializing. Try again in a moment.")

    question = req.question.strip()
    if not question:
        raise HTTPException(status_code=400, detail="Question cannot be empty.")

    # ── Prepend recent history to give the LLM context ──────────
    if req.history:
        recent = req.history[-4:]  # last 2 turns
        history_text = "\n".join(
            f"{'User' if h.role == 'user' else 'Oracle'}: {h.content}"
            for h in recent
        )
        full_question = f"Previous conversation:\n{history_text}\n\nCurrent question: {question}"
    else:
        full_question = question

    try:
        logger.info(f"Question: {question[:80]}...")
        answer = qa_chain.invoke(full_question)
        answer = str(answer).strip()

        if not answer:
            answer = "The wizarding archives hold no record of this. Try rephrasing your question."

        logger.info(f"Answer: {answer[:80]}...")
        return AskResponse(answer=answer)

    except Exception as e:
        logger.error(f"Chain error: {e}")
        raise HTTPException(
            status_code=500,
            detail="The magic is disrupted. The Oracle cannot answer right now."
        )