import logging
import threading
from contextlib import asynccontextmanager

from fastapi import BackgroundTasks, FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from . import ingest, store
from .bm25 import bm25
from .config import settings
from .rag import NoReadyDocuments, answer_question
from .registry import get_registry, public

try:
    from groq import RateLimitError
except Exception:  # pragma: no cover
    RateLimitError = ()  # type: ignore

logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
log = logging.getLogger("rag.main")

_warm = threading.Event()


def _warm_up() -> None:
    """Load models in the background so the server starts instantly."""
    try:
        store.get_embedder()
        store.get_qdrant()
        ingest.recover_interrupted()
        store.get_reranker()
        bm25.ensure_fresh()
    except Exception:
        log.exception("Warm-up failed (the app will retry lazily on first use)")
    finally:
        _warm.set()


@asynccontextmanager
async def lifespan(app: FastAPI):
    threading.Thread(target=_warm_up, daemon=True).start()
    yield


app = FastAPI(title="RAG Backend", version="0.3.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[o.strip() for o in settings.cors_origins.split(",") if o.strip()],
    allow_methods=["*"],
    allow_headers=["*"],
)


class Turn(BaseModel):
    role: str = Field(pattern="^(user|assistant)$")
    content: str


class AskRequest(BaseModel):
    question: str
    doc_ids: list[str] | None = None  # None = search every ready document
    history: list[Turn] = []


@app.get("/health")
def health():
    return {"status": "ok", "llm_model": settings.llm_model, "models_loaded": _warm.is_set()}


@app.get("/documents")
def documents():
    return [public(d) for d in get_registry().all()]


@app.post("/documents")
def upload_document(background: BackgroundTasks, file: UploadFile = File(...)):
    data = file.file.read()
    try:
        entry, duplicate = ingest.submit_document(file.filename or "upload", data)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    if not duplicate:
        background.add_task(ingest.process_document, entry["doc_id"], entry["filename"], data)
    return {**public(entry), "duplicate": duplicate}


@app.delete("/documents/{doc_id}")
def delete_document(doc_id: str):
    if not ingest.delete_document(doc_id):
        raise HTTPException(status_code=404, detail="Document not found.")
    return {"deleted": doc_id}


@app.post("/ask")
def ask(req: AskRequest):
    if not req.question.strip():
        raise HTTPException(status_code=400, detail="Question is empty.")
    if not settings.groq_api_key:
        raise HTTPException(status_code=500, detail="GROQ_API_KEY is not set in .env")
    try:
        return answer_question(req.question.strip(), req.doc_ids, [t.model_dump() for t in req.history])
    except NoReadyDocuments:
        raise HTTPException(status_code=409, detail="No document is ready yet. Upload one and wait until it finishes processing.")
    except RateLimitError:
        raise HTTPException(status_code=429, detail="Groq free-tier rate limit reached. Wait a few seconds and try again.")
    except Exception as e:  # bad key, retired model, network...
        log.exception("ask failed")
        raise HTTPException(status_code=502, detail=f"LLM error: {e}")