from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import os
import torch
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline

app = FastAPI(title="Flykite Model Serving Engine")

class QueryRequest(BaseModel):
    question: str
    chunk_size: int = 1000
    chunk_overlap: int = 200
    k_depth: int = 3
    temperature: float = 0.20
    top_p: float = 0.90

# Pre-caching models on container start
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16
)

model_id = "mistralai/Mistral-7B-Instruct-v0.3"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, quantization_config=bnb_config, device_map="auto"
)

pdf_path = "Flykite Airlines_ HRP.pdf"
raw_documents = PyMuPDFLoader(pdf_path).load() if os.path.exists(pdf_path) else None

@app.get("/")
def health():
    return {"status": "ready", "document_indexed": raw_documents is not None}

@app.post("/api/generate")
@app.post("/api/generate")
def generate(payload: QueryRequest):
    # Checkpoint 1: Confirm payload extraction
    print(f"📬 [API RECEIVED] Incoming query captured: '{payload.question}'")
    
    if not raw_documents:
        raise HTTPException(status_code=500, detail="Handbook missing on host filesystem.")
    
    # Checkpoint 2: Monitor the Data Engineering processing phase
    print(f"🔄 [DATA PROCESSING] Slicing text chunks... Size: {payload.chunk_size} | Overlap: {payload.chunk_overlap}")
    splitter = RecursiveCharacterTextSplitter(chunk_size=payload.chunk_size, chunk_overlap=payload.chunk_overlap)
    chunks = splitter.split_documents(raw_documents)
    
    # Checkpoint 3: Monitor Vector Database indexing
    print(f"📦 [INDEXING] Rebuilding dynamic FAISS vector matrix with {len(chunks)} nodes...")
    db = FAISS.from_documents(chunks, embeddings)
    retriever = db.as_retriever(search_kwargs={"k": payload.k_depth})
    
    # Checkpoint 4: Confirm Retrieval Success
    docs = retriever.invoke(payload.question)
    print(f"🔍 [RETRIEVED] Pulling context... Extracted top {len(docs)} segments for prompt insertion.")
    context = "\n\n".join([c.page_content for c in docs])
    audit = [{"page": d.metadata.get("page", "unknown"), "content": d.page_content} for d in docs]
    
    prompt = (
        f"You are the official Flykite Assistant. Base answers ONLY on context.\n\n"
        f"CONTEXT:\n{context}\n\nQUERY: {payload.question}\n\nRESPONSE:"
    )
    
    # Checkpoint 5: Model Inference Start Execution
    print("🧠 [LLM RUNNING] Tokenizing prompt and spinning compute core via Mistral-7B...")
    gen = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512, 
                   temperature=payload.temperature, top_p=payload.top_p, do_sample=True, pad_token_id=tokenizer.eos_token_id)
    
    response_payload = gen(prompt)[0]["generated_text"].split("RESPONSE:")[-1].strip()
    
    # Checkpoint 6: Success confirmation
    print("🚀 [SUCCESS] Response compiled. Shipping payload back to frontend.")
    return {"answer": response_payload, "audit_trail": audit}