import os
import json
from typing import List, Optional
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
import gradio as gr
import uvicorn
import spaces
import subprocess
import sys

# --- Instalação do llama-cpp-python via wheel (compatível com o ambiente do HF) ---
# Este comando instala a biblioteca a partir de um pacote binário pré-compilado.
# É uma alternativa robusta para evitar erros de compilação no ambiente do Space.
subprocess.check_call([
    sys.executable, "-m", "pip", "install", "--no-cache-dir",
    "https://huggingface.co/Luigi/llama-cpp-python-wheels-hf-spaces-free-cpu/resolve/main/llama_cpp_python-0.3.22-cp310-cp310-linux_x86_64.whl"
])

# Agora que o wheel foi instalado, podemos importar a biblioteca com segurança.
from llama_cpp import Llama

# --- Modelo ---
MODEL_REPO = "bartowski/Dolphin3.0-Llama3.1-8B-GGUF"
MODEL_FILE = "Dolphin3.0-Llama3.1-8B-Q4_K_M.gguf"
N_CTX = 4096
N_THREADS = 4

llm = None

def load_model():
    global llm
    if llm is None:
        print(f"Baixando/carregando modelo: {MODEL_REPO}/{MODEL_FILE}")
        model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
        llm = Llama(
            model_path=model_path,
            n_ctx=N_CTX,
            n_threads=N_THREADS,
            n_gpu_layers=-1,  # Tenta usar a GPU (se disponível)
            verbose=False,
        )
        print("Modelo carregado com sucesso.")
    return llm

# --- Schemas OpenAI ---
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = "dolphin-llama3.1-8b"
    messages: List[ChatMessage]
    temperature: float = 0.7
    max_tokens: Optional[int] = 512
    stream: bool = False

# --- FastAPI ---
api = FastAPI()

@api.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": "dolphin-llama3.1-8b",
                "object": "model",
                "created": 1700000000,
                "owned_by": "local",
            }
        ],
    }

@spaces.GPU  # <-- Esta linha garante que a GPU seja alocada para a inferência
def generate_response(prompt: str, max_tokens: int, temperature: float):
    model = load_model()
    output = model(prompt, max_tokens=max_tokens, temperature=temperature, stream=False)
    return output["choices"][0]["text"]

@api.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    # Formata o prompt para o modelo Dolphin (formato ChatML)
    prompt = ""
    for msg in req.messages:
        if msg.role == "system":
            prompt += f"<|im_start|>system\n{msg.content}<|im_end|>\n"
        elif msg.role == "user":
            prompt += f"<|im_start|>user\n{msg.content}<|im_end|>\n"
        elif msg.role == "assistant":
            prompt += f"<|im_start|>assistant\n{msg.content}<|im_end|>\n"
    prompt += "<|im_start|>assistant\n"

    # --- Modo de streaming ---
    if req.stream:
        async def stream_generator():
            model = load_model()
            output = model(prompt, max_tokens=req.max_tokens, temperature=req.temperature, stream=True)
            for chunk in output:
                token = chunk["choices"][0]["text"]
                data = {
                    "id": "chatcmpl-local",
                    "object": "chat.completion.chunk",
                    "choices": [{"delta": {"content": token}, "index": 0, "finish_reason": None}],
                }
                yield f"data: {json.dumps(data)}\n\n"
            yield "data: [DONE]\n\n"
        return StreamingResponse(stream_generator(), media_type="text/event-stream")

    # --- Modo normal ---
    text = generate_response(prompt, req.max_tokens, req.temperature)
    return {
        "id": "chatcmpl-local",
        "object": "chat.completion",
        "created": 1700000000,
        "model": req.model,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }

# --- Gradio (Interface web + API) ---
def dummy(x):
    return x

demo = gr.Interface(fn=dummy, inputs="text", outputs="text")

# --- Monta o Gradio dentro do FastAPI ---
app = gr.mount_gradio_app(api, demo, path="/ui")

# --- Rodar (o Hugging Face inicia via `python app.py`) ---
if __name__ == "__main__":
    load_model()
    uvicorn.run(app, host="0.0.0.0", port=7860)