# https://www.gradio.app/docs/gradio/server

from gradio import Server
import spaces

from fastapi import HTTPException
from pydantic import BaseModel
from typing import List, Optional

from loguru import logger

# C:\mat-dir\playground\huggingface-stuff\gemma-4-e4b-uncensored\gemma-4-e4b-uncensored\app.py
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 128

@spaces.GPU(duration=66)
def run_inference(prompt: str, max_tokens: int, temperature: float) -> str:
    return f"response to {prompt} - {max_tokens} - {temperature}"

app = Server()

@app.api(name="hello")
def hello(name: str) -> str:
    return f"Hello {name}"

@app.get("/")
def root():
    return {"message": "Hello World"}

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatCompletionRequest):
    logger.debug(" {} ", request)
    try:
        # Convert OpenAI message history into a single model prompt string
        formatted_messages = [{"role": m.role, "content": m.content} for m in request.messages]
        logger.debug(" {} ", formatted_messages)

        # prompt = tokenizer.apply_chat_template(formatted_messages, tokenize=False, add_generation_prompt=True)
        prompt = formatted_messages[0].get("content", "default msg")

        logger.debug(" {} ", prompt)

        # Trigger the ZeroGPU execution block
        response_text = run_inference(prompt, request.max_tokens, request.temperature)

        logger.debug(f"{request.model=}")

        # Structure the payload exactly like OpenAI's JSON response
        return {
            "id": "chatcmpl-zerogpu",
            "object": "chat.completion",
            "model": request.model,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": response_text
                },
                "finish_reason": "stop"
            }]
        }
    except Exception as e:
        logger.debug("e: {}", e)
        raise HTTPException(status_code=500, detail=str(e))

app.launch()