import os

# DEVE stare prima di importare gradio
os.environ["GRADIO_SSR_MODE"] = "false"

import gradio as gr
import uvicorn
from fastapi import FastAPI, Request, HTTPException
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

MODEL_REPO = "rogiskhan/knee-physio-qwen3-4b-GGUF"
MODEL_FILE = "qwen3-4b-instruct-2507.Q4_K_M.gguf"

print("Scarico il modello...", flush=True)
path = hf_hub_download(MODEL_REPO, MODEL_FILE)

print("Carico il modello...", flush=True)
llm = Llama(
    model_path=path,
    n_ctx=4096,
    n_threads=min(os.cpu_count() or 4, 8),
    n_batch=256,
    verbose=False,
)

API_KEY = os.environ.get("API_KEY")
app = FastAPI()


@app.get("/v1/models")
async def list_models():
    return {"object": "list", "data": [{"id": MODEL_REPO, "object": "model"}]}


@app.post("/v1/chat/completions")
async def chat_completions(req: Request):
    if API_KEY and req.headers.get("authorization") != f"Bearer {API_KEY}":
        raise HTTPException(401, "unauthorized")
    body = await req.json()
    if body.get("stream"):
        raise HTTPException(400, "streaming non supportato")
    return llm.create_chat_completion(
        messages=body["messages"],
        temperature=body.get("temperature", 0.7),
        max_tokens=body.get("max_tokens", 512),
    )


def _to_messages(history):
    msgs = []
    for h in history or []:
        if isinstance(h, dict):
            msgs.append({"role": h["role"], "content": h["content"]})
        else:
            user, bot = h
            if user:
                msgs.append({"role": "user", "content": user})
            if bot:
                msgs.append({"role": "assistant", "content": bot})
    return msgs


def respond(message, history):
    out = llm.create_chat_completion(
        messages=_to_messages(history) + [{"role": "user", "content": message}],
        max_tokens=512,
    )
    return out["choices"][0]["message"]["content"]


demo = gr.ChatInterface(respond, title="Knee Physio Qwen3 4B")

# Le rotte /v1 sono già registrate: il mount su "/" finisce in coda
# e non le intercetta. Niente riordino manuale.
app = gr.mount_gradio_app(app, demo, path="/")

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