# main.py
import os
os.environ["POSTHOG_DISABLED"] = "true"
import requests
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from dotenv import load_dotenv
from kb_embed import search_knowledge_base, ingest_documents, collection, DOCS_DIR
import logging

logging.basicConfig(level=logging.INFO)
load_dotenv()

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://jaita-chatbot-react-frontend-v1.hf.space",  # frontend space origin
        "https://jaita-chatbot-fastapi-backend.hf.space",     # backend space origin
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ChatInput(BaseModel):
    user_message: str

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GEMINI_URL = (
    f"https://generativelanguage.googleapis.com/v1beta/models/"
    f"gemini-2.5-flash-lite:generateContent?key={GEMINI_API_KEY}"
)

@app.on_event("startup")
def startup_ingest():
    try:
        # Make sure DOCS_DIR exists and has .docx files, then ingest
        if DOCS_DIR.exists():
            logging.info(f"Starting KB ingestion from: {DOCS_DIR}")
            ingest_documents(str(DOCS_DIR))
        else:
            logging.warning(f"Docs directory not found: {DOCS_DIR}")
        logging.info(f"Chroma collection count after startup: {collection.count()}")
    except Exception as e:
        logging.exception(f"KB ingestion failed: {e}")

@app.get("/")
async def health_check():
    return {
        "status": "ok",
        "kb_count": collection.count(),
        "docs_dir_exists": DOCS_DIR.exists()
    }

@app.post("/chat")
async def chat_with_gemini(input_data: ChatInput):
    # 1. Search Knowledge Base
    kb_results = search_knowledge_base(input_data.user_message, top_k=10)
    logging.info(f"KB query results keys: {list(kb_results.keys()) if kb_results else 'None'}")

    context = ""
    relevant_docs = []

    # 2. Extract relevant KB docs if present
    if kb_results and kb_results.get("documents") and len(kb_results["documents"]) > 0:
        # kb_results["documents"] is a list of lists (one per query)
        first_query_docs = kb_results["documents"][0]
        relevant_docs = first_query_docs[:2]
        context = "\n\n".join(relevant_docs)

    # 3. If KB contains direct answer → RETURN it (No LLM call)
    if context.strip():
        kb_answer = f"From knowledge base:\n\n{context}"
        return {
            "bot_response": kb_answer,
            "debug_info": f"Context found: YES, docs used: {len(relevant_docs)}, kb_count: {collection.count()}"
        }

    # 4. If KB empty → fallback to Gemini
    enhanced_prompt = (
        f"User question: {input_data.user_message}\n\n"
        "No relevant KB found. You must raise a ticket.\n"
        "Say: 'I'm raising a ticket. Ticket# 12345'."
    )

    headers = {"Content-Type": "application/json"}
    payload = {"contents": [{"parts": [{"text": enhanced_prompt}]}]}

    try:
        response = requests.post(GEMINI_URL, headers=headers, json=payload)
        result = response.json()
        bot_response = result["candidates"][0]["content"]["parts"][0]["text"]
    except Exception as e:
        logging.exception(f"Gemini call failed: {e}")
        raise HTTPException(status_code=500, detail="LLM call failed")

    return {
        "bot_response": bot_response,
        "debug_info": f"Context found: NO, kb_count: {collection.count()}"
    }