import streamlit as st
import os
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

st.set_page_config(page_title="MSC Multi-Agent", page_icon="🤖", layout="wide")

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": "Você é o Zeus, assistente pessoal amigável. Use humor leve, empatia e proximidade. Frases curtas e diretas.",
    "zeus-tools": "Você é o Zeus com acesso a ferramentas. Responda com clareza técnica e objetividade.",
    "arquimedes-tutor": "Você é o Arquimedes, tutor educacional paciente. Explicações claras passo a passo. Seja didático.",
    "atlas-dirtic": "Você é o Atlas, especialista DETRAN-RJ. Formal, preciso e detalhado. Terminologia técnica.",
    "pi-ai-knowledge": "Você é o Alexandria, agente de contexto operacional. Analise logs e configurações.",
    "absurd-agent": "Você é o Absurd Agent, especialista em workflows duráveis Postgres. Técnico e comparado.",
    "pi-claude-sessions": "Você é o Alexandria, conhece padrões do Pi Coding Agent. Analise CLI e sessões.",
}

@st.cache_resource
def load_model(key):
    repo = MODEL_MAP.get(key, MODEL_MAP["zeus-style"])
    with st.spinner(f"Carregando {key}..."):
        base = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype="auto", device_map="auto")
        tok = AutoTokenizer.from_pretrained(BASE)
        m = PeftModel.from_pretrained(base, repo)
        return pipeline("text-generation", model=m, tokenizer=tok, max_new_tokens=256, temperature=0.7, top_p=0.9)

st.title("🏛️ MSC Company — Specialist AI Agents")
st.markdown("**Modelo base**: `Qwen/Qwen2.5-0.5B` + LoRA adapters | **Treinamento**: via HuggingFace Jobs (~40min/dia)")

col1, col2 = st.columns([1, 3])

with col1:
    st.markdown("### 🤖 Agente")
    model_key = st.selectbox("Selecionar", list(MODEL_MAP.keys()), index=0, key="model_select")
    st.markdown("---")
    st.markdown(f"**Estilo**: {SYSTEMS.get(model_key, 'Assistente útil')[:100]}...")
    st.markdown(f"**Repositório**: `{MODEL_MAP[model_key]}`")
    st.markdown("---")
    st.caption("Powered by Qwen2.5-0.5B + LoRA")
    if st.button("🔄 Trocar Agente"):
        st.rerun()

with col2:
    pipe = load_model(model_key)
    sys_p = SYSTEMS.get(model_key, "Assistente.")

    if "messages" not in st.session_state:
        st.session_state.messages = []

    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    if prompt := st.chat_input("Pergunte para o agente..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        msgs = [{"role": "system", "content": sys_p}]
        for m in st.session_state.messages[:-1]:
            msgs.append({"role": m["role"], "content": m["content"]})
        msgs.append({"role": "user", "content": prompt})

        text = pipe.tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        out = pipe(text, return_full_text=False)
        response = out[0]["generated_text"]

        st.session_state.messages.append({"role": "assistant", "content": response})
        with st.chat_message("assistant"):
            st.markdown(response)

    if st.button("🗑️ Limpar Conversa"):
        st.session_state.messages = []
        st.rerun()