# app.py — GraphRAG Portfolio Agent | GPT-4o-mini + Neo4j
import streamlit as st
import os

st.set_page_config(
    page_title="GraphRAG — Portfolio Agent",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# ── SESSION STATE ─────────────────────────────────────────────
for k, v in {
    'messages': [], 'neo4j': None, 'neo4j_ok': False,
    'agent': None, 'schema_ok': False,
    'openai_key': '', 'graph_stats': None,
}.items():
    if k not in st.session_state:
        st.session_state[k] = v

# ── HELPERS ───────────────────────────────────────────────────
def get_neo4j_config():
    cfg = {}
    try:
        s = st.secrets
        if 'NEO4J_URI' in s:
            return {'uri': s['NEO4J_URI'], 'username': s['NEO4J_USERNAME'],
                    'password': s['NEO4J_PASSWORD'],
                    'database': s.get('NEO4J_DATABASE', 'neo4j')}
        if 'neo4j' in s:
            n = s['neo4j']
            return {'uri': n.get('uri',''), 'username': n.get('username',''),
                    'password': n.get('password',''), 'database': n.get('database','neo4j')}
    except Exception:
        pass
    return {'uri': os.getenv('NEO4J_URI',''), 'username': os.getenv('NEO4J_USERNAME',''),
            'password': os.getenv('NEO4J_PASSWORD',''), 'database': os.getenv('NEO4J_DATABASE','neo4j')}

def get_openai_key():
    try:
        if 'OPENAI_API_KEY' in st.secrets:
            return st.secrets['OPENAI_API_KEY']
    except Exception:
        pass
    return os.getenv('OPENAI_API_KEY', st.session_state.openai_key)

def conectar_neo4j():
    if st.session_state.neo4j is not None:
        return st.session_state.neo4j
    try:
        from neo4j import GraphDatabase
        cfg = get_neo4j_config()
        if not all([cfg['uri'], cfg['username'], cfg['password']]):
            return None
        driver = GraphDatabase.driver(cfg['uri'], auth=(cfg['username'], cfg['password']))
        with driver.session(database=cfg['database']) as s:
            s.run('RETURN 1')
        st.session_state.neo4j    = (driver, cfg['database'])
        st.session_state.neo4j_ok = True
        return st.session_state.neo4j
    except Exception as e:
        st.session_state.neo4j_ok = False
        return None

def criar_agente(openai_key):
    try:
        from graph_agent import GraphRAGAgent
        conn = st.session_state.neo4j
        if conn is None:
            return None
        driver, database = conn
        return GraphRAGAgent(openai_key, driver, database)
    except Exception as e:
        st.error(f"Erro ao criar agente: {e}")
        return None

# ── SIDEBAR ───────────────────────────────────────────────────
with st.sidebar:
    st.title("🤖 GraphRAG Agent")
    st.caption("GPT-4o-mini · Neo4j · Cypher")
    st.divider()

    # Neo4j status
    conn = conectar_neo4j()
    if st.session_state.neo4j_ok:
        st.success("🗄️ Neo4j Conectado")
    else:
        st.error("Neo4j offline")
        st.info("Configure NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD nos secrets do HF Space.")

    st.divider()

    # OpenAI Key
    st.markdown("#### 🔑 OpenAI API Key")
    key_input = st.text_input(
        "Chave", type="password",
        value=st.session_state.openai_key,
        placeholder="sk-...",
        help="Ou configure OPENAI_API_KEY nos secrets do HF"
    )
    if key_input:
        st.session_state.openai_key = key_input

    openai_key = get_openai_key()
    if openai_key:
        st.success("✅ OpenAI Key configurada")
    else:
        st.warning("OpenAI Key não configurada")

    st.divider()

    # Popular grafo
    st.markdown("#### 🧬 Base de Conhecimento")
    if st.session_state.neo4j_ok:
        if st.button("📥 Popular Grafo Neo4j", use_container_width=True, type="primary"):
            try:
                from graph_knowledge import popular_neo4j, verificar_schema
                driver, database = st.session_state.neo4j
                with st.spinner("Populando Neo4j com projetos GNN..."):
                    n, erros = popular_neo4j(driver, database)
                    nos, arestas = verificar_schema(driver, database)
                    st.session_state.graph_stats = {'nos': nos, 'arestas': arestas}
                    st.session_state.schema_ok   = True
                st.success(f"✅ {n} statements executados!")
                if erros:
                    st.warning(f"{len(erros)} avisos (normal em re-execuções)")
            except Exception as e:
                st.error(f"Erro: {e}")

        if st.session_state.graph_stats:
            st.markdown("**Nós no grafo:**")
            for item in st.session_state.graph_stats['nos']:
                col1, col2 = st.columns([3,1])
                col1.caption(item['tipo'])
                col2.markdown(f"**{item['total']}**")
            total_ar = sum(a['total'] for a in st.session_state.graph_stats['arestas'])
            st.caption(f"Total de arestas: **{total_ar}**")
    else:
        st.info("Conecte o Neo4j primeiro")

    st.divider()
    if st.button("🗑️ Limpar Chat", use_container_width=True):
        st.session_state.messages = []
        st.rerun()

# ── MAIN ──────────────────────────────────────────────────────
st.title("🤖 Portfolio GraphRAG Agent")
st.markdown("Pergunte qualquer coisa sobre os **5 projetos de GNN** — o agente consulta o Neo4j com Cypher gerado pelo GPT e responde.")
st.divider()

# Alertas de configuração
if not get_openai_key():
    st.warning("⬅️ Configure sua OpenAI API Key na sidebar para usar o agente.")

if not st.session_state.neo4j_ok:
    st.error("Neo4j não conectado. Configure as credenciais na sidebar.")

if st.session_state.neo4j_ok and not st.session_state.schema_ok:
    st.info("⬅️ Clique em **Popular Grafo Neo4j** na sidebar para carregar a base de conhecimento.")

# ── COMO FUNCIONA ─────────────────────────────────────────────
if not st.session_state.messages:
    with st.expander("ℹ️ Como funciona o GraphRAG?", expanded=True):
        c1, c2, c3 = st.columns(3)
        with c1:
            st.markdown("**1. Você pergunta**")
            st.markdown("Em linguagem natural, sobre qualquer projeto, tecnologia ou conceito do portfólio.")
        with c2:
            st.markdown("**2. GPT gera Cypher**")
            st.markdown("O GPT-4o-mini analisa o schema do grafo e gera uma query Cypher para o Neo4j.")
        with c3:
            st.markdown("**3. Neo4j responde**")
            st.markdown("A query é executada, o contexto é passado de volta ao GPT que formula a resposta final.")

    st.markdown("### 💡 Sugestões de perguntas")
    sugestoes = [
        "Quais projetos usam PyTorch Geometric?",
        "Qual projeto tem maior AUC?",
        "Me explique o DOMINANT",
        "Qual a diferença entre HetGNN e GraphSAGE?",
        "Quais papers são referenciados?",
        "Que conceitos o TGN implementa?",
        "Qual projeto usa dado real?",
        "Como funciona Inductive Learning?",
    ]
    cols = st.columns(4)
    for i, sug in enumerate(sugestoes):
        with cols[i % 4]:
            if st.button(sug, key=f"sug_{i}", use_container_width=True):
                st.session_state.pending_question = sug
                st.rerun()

# ── CHAT HISTORY ─────────────────────────────────────────────
for msg in st.session_state.messages:
    with st.chat_message(msg["role"], avatar="🧑" if msg["role"]=="user" else "🤖"):
        if msg["role"] == "assistant" and msg.get("cypher"):
            with st.expander(f"🔍 Cypher gerado ({msg.get('n_resultados', 0)} resultados)", expanded=False):
                st.code(msg["cypher"], language="cypher")
        st.markdown(msg["content"])

# ── INPUT ─────────────────────────────────────────────────────
# Processa pergunta pendente (de botão de sugestão)
pergunta_pendente = st.session_state.pop("pending_question", None)

pergunta = st.chat_input("Pergunte sobre os projetos de GNN...") or pergunta_pendente

if pergunta and get_openai_key() and st.session_state.neo4j_ok:
    # Mostra mensagem do usuário
    with st.chat_message("user", avatar="🧑"):
        st.markdown(pergunta)
    st.session_state.messages.append({"role": "user", "content": pergunta})

    # Cria agente e responde
    agent = criar_agente(get_openai_key())
    if agent:
        with st.chat_message("assistant", avatar="🤖"):
            with st.spinner("Consultando grafo Neo4j..."):
                try:
                    resultado = agent.responder(pergunta)

                    with st.expander(
                        f"🔍 Cypher gerado ({len(resultado['resultados'])} resultados)",
                        expanded=False
                    ):
                        st.code(resultado["cypher"], language="cypher")

                    st.markdown(resultado["resposta"])

                    st.session_state.messages.append({
                        "role":         "assistant",
                        "content":      resultado["resposta"],
                        "cypher":       resultado["cypher"],
                        "n_resultados": len(resultado["resultados"]),
                    })
                except Exception as e:
                    msg_erro = f"Erro ao processar: {e}"
                    st.error(msg_erro)
                    st.session_state.messages.append({
                        "role": "assistant", "content": msg_erro})
    else:
        st.error("Não foi possível criar o agente. Verifique as configurações.")

elif pergunta and not get_openai_key():
    st.warning("Configure a OpenAI API Key na sidebar.")
elif pergunta and not st.session_state.neo4j_ok:
    st.error("Neo4j não conectado.")