import os
import torch
import uvicorn
import faiss
import numpy as np
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from typing import Optional, List
from sentence_transformers import SentenceTransformer
from fastapi.concurrency import run_in_threadpool
from engine import RAGEngine

app = FastAPI(
    title="ENVH.AI Compliance API",
    version="5.1.0",
    description="Kenyan EHS legal RAG — Groq direct + multi-turn conversation + suggested follow-ups",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "https://envhai.co.ke",
        "https://www.envhai.co.ke",
        "http://localhost:3000",
        "http://localhost:5173",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

print("🚀 Loading embedder …")
embedder = SentenceTransformer("BAAI/bge-small-en-v1.5", device="cpu")
rag      = RAGEngine()
print("✅ Server ready.")


# ── Request models ────────────────────────────────────────────────────────────
class HistoryEntry(BaseModel):
    role: str   # "user" or "ai"
    text: str

class QueryRequest(BaseModel):
    text:    str
    history: Optional[List[HistoryEntry]] = []

class DebugRequest(BaseModel):
    text: str
    k:    Optional[int] = 12


# ── Embed helper ──────────────────────────────────────────────────────────────
def _embed(text: str) -> np.ndarray:
    with torch.no_grad():
        vec = embedder.encode([text]).astype("float32")
    faiss.normalize_L2(vec)
    return vec


# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/")
def root():
    return {
        "status":  "online",
        "version": "5.1.0",
        "model":   rag.model_name,
        "engine":  "Groq (direct)",
    }

@app.get("/health")
def health():
    return {"status": "ok", "model": rag.model_name}


@app.post("/ask")
async def ask(request: QueryRequest):
    """
    Always returns HTTP 200 {"answer": "..."}.
    Errors are returned as readable strings inside answer — never a 500.
    """
    text = request.text.strip()
    if not text:
        return JSONResponse({"answer": "⚠️ Please type a question first.", "suggestions": []})

    try:
        history_dicts = [
            {"role": h.role, "text": h.text}
            for h in (request.history or [])
            if h.text.strip()
        ]

        # Step 1 — enrich vague follow-ups with topic context
        retrieval_query = rag.build_retrieval_query(text, history_dicts)

        # Step 2 — expand Swahili / acronyms
        expanded = rag.expand_query(retrieval_query)

        # Step 3 — embed
        query_vector = _embed(expanded)

        # Step 4 — search + answer. Returns {"answer": str, "suggestions": [...]}
        result = await run_in_threadpool(
            rag.search_and_ask,
            text,           # raw query shown to LLM
            query_vector,   # enriched vector for FAISS
            history_dicts,
        )
        return JSONResponse({
            "answer": result.get("answer", ""),
            "suggestions": result.get("suggestions", []),
        })

    except Exception as e:
        # Log real error to Space logs
        print(f"❌ /ask unhandled error: {type(e).__name__}: {e}")
        return JSONResponse({
            "answer": (
                f"⚠️ **Unexpected Error — `{type(e).__name__}: {e}`**\n\n"
                f"Please share this with support@envhai.co.ke"
            ),
            "suggestions": [],
        })


@app.delete("/session")
def clear_session():
    """Stateless ack — history is owned by the frontend."""
    return {"status": "cleared"}


@app.post("/debug")
async def debug(request: DebugRequest):
    """Dev endpoint — returns raw FAISS scores. Not for end users."""
    if not request.text.strip():
        return JSONResponse({"error": "Query cannot be empty."})
    try:
        expanded = rag.expand_query(request.text)
        vec = _embed(expanded).reshape(1, -1)
        k   = min(request.k, int(rag.index.ntotal))
        distances, indices = rag.index.search(vec, k=k)
        results = [
            {
                "rank":              i + 1,
                "index":             int(idx),
                "cosine_similarity": round(float(dist), 4),
                "chunk_preview":     str(rag.chunks[idx])[:200].replace("\n", " "),
            }
            for i, (dist, idx) in enumerate(zip(distances[0], indices[0]))
            if idx != -1
        ]
        return {
            "original_query": request.text,
            "expanded_query": expanded,
            "model":          rag.model_name,
            "thresholds":     {"strong": 0.45, "partial": 0.25},
            "top_score":      results[0]["cosine_similarity"] if results else None,
            "tier":           rag._confidence_tier(distances[0][distances[0] != -1]),
            "results":        results,
        }
    except Exception as e:
        return JSONResponse({"error": str(e)})


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)