# app.py  – FastAPI core with public Gradio UI
import os
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
import gradio as gr
from rag import retrieve_info  # <-- function defined in rag.py

# =============== Config ===============
VALID_KEYS = set(os.getenv("RAG_KEYS", "alpha").split(","))  # add real keys in Space » Settings » Secrets

# =============== FastAPI ==============
api = FastAPI(title="Moffitt RAG API")

class QueryIn(BaseModel):
    query: str
    k: int = 5

@api.post("/v1/query")
async def rag_query(body: QueryIn, x_api_key: str = Header(None)):
    if x_api_key not in VALID_KEYS:
        raise HTTPException(status_code=401, detail="Invalid or missing X-API-Key")
    answer = retrieve_info(body.query, body.k)
    return {"answer": answer}

# =============== Public Gradio UI ==============
def ui_fn(q, k):
    return retrieve_info(q, int(k))

demo = gr.Interface(
    fn=ui_fn,
    inputs=["text", gr.Slider(1, 20, value=5, step=1, label="Top-k")],
    outputs=gr.Textbox(lines=25, label="Retrieved chunks"),
    title="🔍 Moffitt Retrieval Demo",
    allow_flagging="never",
)

# Mount UI at "/"
app = gr.mount_gradio_app(api, demo, path="/")

# ---------- Local debugging ----------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)