import re
import spaces
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "SynastriaNetworks/OpenFable-4B"

print("⏳ Carregando modelo...")
_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)
print("✅ Modelo carregado e pronto!")

@spaces.GPU
def gerar_resposta(mensagens_json):
    inputs = _tokenizer.apply_chat_template(
        mensagens_json,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True
    ).to(_model.device)

    saida = _model.generate(
        **inputs,
        max_new_tokens=1024,
        temperature=0.7,
        do_sample=True,
        pad_token_id=_tokenizer.eos_token_id
    )

    resposta_ids = saida[0][inputs["input_ids"].shape[-1]:]
    return _tokenizer.decode(resposta_ids, skip_special_tokens=True).strip()

def formatar_resposta(texto):
    match = re.search(r"<think>(.*?)</think>", texto, re.DOTALL)

    if not match:
        return texto.strip()

    pensamento = match.group(1).strip()
    resposta_final = texto[match.end():].strip()

    if not resposta_final:
        resposta_final = "*(sem resposta final gerada)*"

    return (
        "<details>\n"
        "<summary>💭 Ver raciocínio</summary>\n\n"
        f"{pensamento}\n\n"
        "</details>\n\n"
        f"{resposta_final}"
    )

def responder_chat(mensagem, historico):
    try:
        mensagens_json = [
            {"role": "system", "content": "Você é o OpenFable, um assistente inteligente e prestativo que responde em português."}
        ]
        for msg in historico:
            mensagens_json.append({"role": msg["role"], "content": msg["content"]})

        mensagens_json.append({"role": "user", "content": mensagem})

        resposta_bruta = gerar_resposta(mensagens_json)
        return formatar_resposta(resposta_bruta)

    except Exception as e:
        return f"❌ Erro: {type(e).__name__}: {e}"

with gr.Blocks() as demo:
    gr.ChatInterface(
        fn=responder_chat,
        title="OpenFable Chatbot",
        description="Interface rodando o modelo OpenFable.",
        type="messages",
    )

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