import os
import time
import json
import uuid
import torch
import spaces
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Optional, Dict
from transformers import AutoTokenizer, TextIteratorStreamer
from awq import AutoAWQForCausalLM
from threading import Thread
import gradio as gr

MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct-AWQ"

# Cargar Tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)

# Cargar Modelo AWQ
model = AutoAWQForCausalLM.from_quantized(
    MODEL_ID,
    fuse_layers=False,
    trust_remote_code=True,
    safetensors=True
)

# Esquemas para la API OpenAI
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: Optional[str] = MODEL_ID
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.2
    top_p: Optional[float] = 0.9
    max_tokens: Optional[int] = 2048
    stream: Optional[bool] = False

# Generación en GPU con ZeroGPU
@spaces.GPU(duration=120)
def generate_stream_tokens(messages_dict: List[Dict[str, str]], temperature: float, top_p: float, max_tokens: int):
    text = tokenizer.apply_chat_template(
        messages_dict,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=60.0,
        skip_prompt=True,
        skip_special_tokens=True
    )

    generate_kwargs = dict(
        model_inputs,
        streamer=streamer,
        max_new_tokens=max_tokens,
        do_sample=temperature > 0.0,
        temperature=max(temperature, 1e-2) if temperature > 0.0 else None,
        top_p=top_p if temperature > 0.0 else None,
        repetition_penalty=1.05
    )

    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()

    for new_token in streamer:
        yield new_token

# Funciones de soporte para Gradio UI
def gradio_generate(message, history, system_prompt, temperature, top_p, max_tokens):
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    partial_text = ""
    for token in generate_stream_tokens(messages, temperature, top_p, int(max_tokens)):
        partial_text += token
        yield partial_text

# Definición de la interfaz Gradio
with gr.Blocks(title="Qwen2.5-Coder-32B API & UI") as demo:
    gr.Markdown("# 🚀 Qwen2.5-Coder-32B-Instruct (AWQ en ZeroGPU)")
    gr.Markdown("Compatible con OpenAI: `/v1/chat/completions` | Modelo: `Qwen/Qwen2.5-Coder-32B-Instruct-AWQ`")
    
    gr.ChatInterface(
        fn=gradio_generate,
        additional_inputs=[
            gr.Textbox("Eres un asistente de programación experto.", label="System Prompt"),
            gr.Slider(0.0, 1.0, 0.2, step=0.05, label="Temperature"),
            gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top-P"),
            gr.Slider(256, 4096, 2048, step=256, label="Max New Tokens")
        ]
    )

# Definición de FastAPI y vinculación de endpoints a la app subyacente de Gradio
fastapi_app = demo.app

@fastapi_app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_ID,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "huggingface"
            }
        ]
    }

@fastapi_app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created_time = int(time.time())
    messages_dict = [{"role": m.role, "content": m.content} for m in req.messages]

    if req.stream:
        async def event_generator():
            for token in generate_stream_tokens(
                messages_dict,
                temperature=req.temperature or 0.2,
                top_p=req.top_p or 0.9,
                max_tokens=req.max_tokens or 2048
            ):
                chunk = {
                    "id": req_id,
                    "object": "chat.completion.chunk",
                    "created": created_time,
                    "model": req.model,
                    "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}]
                }
                yield f"data: {json.dumps(chunk)}\n\n"
            
            final_chunk = {
                "id": req_id,
                "object": "chat.completion.chunk",
                "created": created_time,
                "model": req.model,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]
            }
            yield f"data: {json.dumps(final_chunk)}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(event_generator(), media_type="text/event-stream")

    full_content = ""
    for token in generate_stream_tokens(
        messages_dict,
        temperature=req.temperature or 0.2,
        top_p=req.top_p or 0.9,
        max_tokens=req.max_tokens or 2048
    ):
        full_content += token

    return {
        "id": req_id,
        "object": "chat.completion",
        "created": created_time,
        "model": req.model,
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": full_content},
                "finish_reason": "stop"
            }
        ],
        "usage": {"prompt_tokens": -1, "completion_tokens": -1, "total_tokens": -1}
    }

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)