import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.routing import Route
import gradio as gr

MODEL_NAME = "teknium/OpenHermes-2.5-Mistral-7B"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.float16, device_map="cpu")
model.eval()

@spaces.GPU(duration=120)
def generate_reply(prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
    model.to("cuda")
    try:
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        with torch.inference_mode():
            output = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                do_sample=True,
                temperature=temperature,
                pad_token_id=tokenizer.eos_token_id,
            )
        return tokenizer.decode(output[0], skip_special_tokens=True)
    finally:
        model.to("cpu")
        torch.cuda.empty_cache()

with gr.Blocks() as demo:
    gr.Markdown("API is running. POST to `/v1/chat/completions`.")

# --- Register the custom route FIRST ---
async def chat_completions(request: Request):
    payload = await request.json()
    messages = payload.get("messages", [])
    temperature = float(payload.get("temperature", 0.7))
    max_tokens = int(payload.get("max_tokens", 512))

    prompt = ""
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        prompt += f"{role}: {content}\n"
    prompt += "assistant:"

    reply = generate_reply(prompt, max_tokens=max_tokens, temperature=temperature)

    return JSONResponse({
        "id": "openhermes-2.5-mistral-7b-response",
        "object": "chat.completion",
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": reply},
            "finish_reason": "stop",
        }],
    })

# Insert at position 0 so it runs before Gradio's catch-all
demo.app.router.routes.insert(0, Route("/v1/chat/completions", chat_completions, methods=["POST"]))

# --- CORS middleware ---
demo.app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)