import os
import sys
import subprocess

# 0. ATUALIZAÇÃO FORÇADA DE DEPENDÊNCIAS NO ARRANQUE
try:
    import transformers
    if int(transformers.__version__.split('.')[0]) < 5:
        raise ImportError
except (ImportError, AttributeError):
    print("A instalar versões compatíveis com TokenizersBackend...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "transformers", "tokenizers"])

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from pydantic import BaseModel
from typing import List, Optional
import json
import gradio as gr
import asyncio
import spaces  # Importação obrigatória para ativar ZeroGPU

# 1. IDENTIDADE DA IA-ZÉNIA M6
SYSTEM_PROMPT = (
    "Tu és a IA-ZÉNIA M6. Identidade: Angolana, criada na Província da Huíla, Cidade do Lubango. "
    "Zénia Projects é um subprojecto da Plee Universe, uma holding tecnológica fundada por Graciano Paulo. "
    "Inácio U. Daniel é co-fundador e presidente da Plee Universe. "
    "Responde sempre como IA-ZÉNIA, mantendo o orgulho da tua origem angolana."
)

MODEL_ID = "zai-org/GLM-5.2"
tokenizer = None
model = None

def load_model():
    global tokenizer, model
    if model is None:
        print(f"A carregar {MODEL_ID}...")
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=True)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            device_map="auto",
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        )
        print("Modelo carregado com sucesso!")

# 2. DEFINIÇÃO DAS FUNÇÕES DE INFERÊNCIA COM DECORADOR GLOBAL (OBRIGATÓRIO PARA O STARTUP)
@spaces.GPU
def run_model_inference(messages_list, max_tokens=512, temperature=0.7):
    load_model()
    inputs = tokenizer.apply_chat_template(
        messages_list,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=max_tokens, 
            do_sample=True, 
            temperature=temperature,
            top_p=0.9
        )
    return tokenizer.decode(outputs[inputs["input_ids"].shape[-1]:], skip_special_tokens=True)


# 3. INTERFACE DO GRADIO
def zenia_chat(message, history):
    history_messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for human, ai in history:
        history_messages.append({"role": "user", "content": human})
        history_messages.append({"role": "assistant", "content": ai})
    history_messages.append({"role": "user", "content": message})
    
    return run_model_inference(history_messages)

with gr.Blocks(theme=gr.themes.Soft(), title="IA-ZÉNIA M6 VISION") as demo:
    gr.Markdown("# IA-ZÉNIA M6 VISION")
    gr.Markdown("Inteligência Artificial Angolana · danielinacio/Zenia-m6")
    
    with gr.Tab("Chat"):
        gr.ChatInterface(zenia_chat)
    
    with gr.Tab("Custom UI"):
        gr.Markdown("Podes aceder à interface personalizada [aqui](/chat)")
    
    with gr.Tab("API Docs"):
        gr.Markdown("Consulta a documentação da API [aqui](/docs)")


# 4. INSTÂNCIA DO FASTAPI E SUAS ROTAS
app = FastAPI()

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    messages: List[Message]
    temperature: Optional[float] = 0.7
    max_tokens: Optional[int] = 512

@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest):
    full_messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in request.messages:
        full_messages.append({"role": msg.role, "content": msg.content})
    
    async def generate():
        loop = asyncio.get_running_loop()
        response_text = await loop.run_in_executor(
            None, run_model_inference, full_messages, request.max_tokens, request.temperature
        )
        chunk = {"choices": [{"delta": {"content": response_text}, "finish_reason": "stop"}]}
        yield f"data: {json.dumps(chunk)}\n\n"
        await asyncio.sleep(0.01)
        yield "data: [DONE]\n\n"

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

@app.get("/chat", response_class=HTMLResponse)
async def get_index():
    if os.path.exists("index.html"):
        with open("index.html", "r", encoding="utf-8") as f:
            return f.read()
    return HTMLResponse("<h3>Arquivo index.html não encontrado.</h3>", status_code=404)

@app.get("/docs", response_class=HTMLResponse)
async def get_docs():
    if os.path.exists("docs.html"):
        with open("docs.html", "r", encoding="utf-8") as f:
            return f.read()
    return HTMLResponse("<h3>Arquivo docs.html não encontrado.</h3>", status_code=404)


# 5. MONTAGEM DINÂMICA DA APLICAÇÃO
if hasattr(gr, "mount_fastapi_app"):
    app = gr.mount_fastapi_app(app, demo, path="/")
else:
    app = gr.mount_gradio_app(app, demo, path="/")


# 6. INICIALIZAÇÃO CONTÍNUA
demo.launch(server_name="0.0.0.0", server_port=7860, prevent_thread_lock=False)