import gradio as gr
import os
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

BASE = "Qwen/Qwen2.5-0.5B"
MODEL_MAP = {
    "zeus-style": "Finish-him/zeus-style-sft-v1",
    "zeus-tools": "Finish-him/zeus-tools-sft-v1",
    "arquimedes-tutor": "MSC-Company/arquimedes-tutor-sft-v1",
    "atlas-dirtic": "Finish-him/atlas-dirtic-rag-v2",
    "pi-ai-knowledge": "Finish-him/pi-ai-knowledge-v1",
    "absurd-agent": "Finish-him/absurd-agent-sft-v1",
    "pi-claude-sessions": "Finish-him/pi-claude-sessions-rag-v1",
}

SYSTEMS = {
    "zeus-style": "Voce e o Zeus, assistente pessoal amigavel. Use humor leve, empatia e proximidade. Frases curtas e diretas.",
    "zeus-tools": "Voce e o Zeus com acesso a ferramentas. Responda com clareza tecnica e objetividade.",
    "arquimedes-tutor": "Voce e o Arquimedes, tutor educacional paciente. Explicacoes claras passo a passo. Seja didatico.",
    "atlas-dirtic": "Voce e o Atlas, especialista DETRAN-RJ. Formal, preciso e detalhado. Terminologia tecnica.",
    "pi-ai-knowledge": "Voce e o Alexandria, agente de contexto operacional. Analise logs e configuracoes.",
    "absurd-agent": "Voce e o Absurd Agent, especialista em workflows duraveis Postgres. Tecnico e comparado.",
    "pi-claude-sessions": "Voce e o Alexandria, conhece padroes do Pi Coding Agent. Analise CLI e sessoes.",
}

print("[MSC] Carregando base model...")
base = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype="auto", device_map="auto")
tok = AutoTokenizer.from_pretrained(BASE)
print("[MSC] Base OK")

def get_pipe(key):
    key = key or "zeus-style"
    repo = MODEL_MAP.get(key, MODEL_MAP["zeus-style"])
    print(f"[MSC] Carregando: {repo}")
    m = PeftModel.from_pretrained(base, repo)
    return pipeline("text-generation", model=m, tokenizer=tok, max_new_tokens=128, temperature=0.7)

active = os.environ.get("MSC_MODEL", "zeus-style")
pipe = get_pipe(active)
sys_p = SYSTEMS.get(active, "Assistente util.")

def respond(msg, hist, key):
    global pipe, sys_p, active
    if key != active:
        active = key
        pipe = get_pipe(key)
        sys_p = SYSTEMS.get(key, "Assistente.")
    msgs = [{"role": "system", "content": sys_p}]
    for u, a in hist:
        msgs += [{"role":"user","content":u},{"role":"assistant","content":a}]
    msgs.append({"role":"user","content":msg})
    t = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    out = pipe(t, return_full_text=False)
    hist.append((msg, out[0]["generated_text"]))
    return "", hist

with gr.Blocks(title="MSC Multi-Agent") as demo:
    gr.Markdown("# 🤖 MSC Company — Specialist AI Agents")
    gr.Markdown(f"**Base**: Qwen2.5-0.5B + LoRA | **Agente ativo**: {active}")
    gr.Markdown("12 modelos treinados via HuggingFace Jobs (~40min/dia)")
    with gr.Row():
        cb = gr.Chatbot(height=400)
        with gr.Column(scale=1):
            gr.Markdown("### Selecionar Agente")
            dd = gr.Dropdown(list(MODEL_MAP.keys()), value=active, label="Agente")
            def_desc = {
                "zeus-style": "Pessoal, amigável, com humor leve",
                "zeus-tools": "Técnico, preciso, com ferramentas",
                "arquimedes-tutor": "Didático, paciente, educacional",
                "atlas-dirtic": "Formal, DETRAN-RJ, documentos",
                "pi-ai-knowledge": "Contextual, operacional, logs",
                "absurd-agent": "Workflows Postgres duráveis",
                "pi-claude-sessions": "Pi Coding Agent, CLI, sessões",
            }
            for k, v in def_desc.items():
                gr.Markdown(f"- **{k}**: {v}")
    msg = gr.Textbox(placeholder="Pergunte para o agente...", label="Pergunta")
    with gr.Row():
        bt = gr.Button("Enviar", variant="primary")
        cl = gr.Button("Limpar")
    bt.click(respond, [msg, cb, dd], [msg, cb])
    msg.submit(respond, [msg, cb, dd], [msg, cb])
    cl.click(lambda: ("", []), [], [msg, cb])

demo.launch()

# gradio 6.x compatible - fixed HfFolder + audioop issues