import os
import pandas as pd
import gradio as gr
from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
from langchain.chat_models import init_chat_model
from langchain_core.rate_limiters import InMemoryRateLimiter
from dotenv import load_dotenv

load_dotenv()

rate_limiter = InMemoryRateLimiter(
    requests_per_second=0.5,
    check_every_n_seconds=0.1,
    max_bucket_size=10,
)


llm = init_chat_model(
    "llama-3.3-70b-versatile",
    model_provider="groq",
    rate_limiter=rate_limiter,
    api_key=os.getenv("GROQ_API_KEY"),
)

# Carregar datasets
print("Carregando dados...")
df_funcionarios = pd.read_csv("CadastroFuncionarios.csv", sep=";", decimal=",")
df_clientes = pd.read_csv("CadastroClientes.csv", sep=";")
df_servicos = pd.read_excel("BaseServiçosPrestados.xlsx")

df_integrado = df_servicos.merge(df_clientes, on="ID Cliente", how="left")
df_integrado = df_integrado.merge(
    df_funcionarios[["ID Funcionário", "Salario Base", "Cargo", "Area"]],
    on="ID Funcionário",
    how="left",
)
df_integrado["Receita Estimada"] = (
    df_integrado["Valor Contrato Mensal"]
    * df_integrado["Tempo Total de Contrato (Meses)"]
)

print(
    f"✓ {len(df_funcionarios)} funcionários, {len(df_clientes)} clientes, {len(df_servicos)} serviços"
)

# Construir descrição dinâmica dos DataFrames
prefix_prompt = f"""Você é um agente de análise de dados com acesso a 4 DataFrames pandas já carregados:

df1: CadastroFuncionarios ({len(df_funcionarios)} registros)
   Colunas: {', '.join(df_funcionarios.columns.tolist())}

df2: CadastroClientes ({len(df_clientes)} registros)
   Colunas: {', '.join(df_clientes.columns.tolist())}

df3: BaseServiçosPrestados ({len(df_servicos)} registros)
   Colunas: {', '.join(df_servicos.columns.tolist())}

df4: df_integrado ({len(df_integrado)} registros) - Merge completo dos 3 DataFrames
   Colunas: {', '.join(df_integrado.columns.tolist())}

REGRAS OBRIGATÓRIAS:
1. SEMPRE use python_repl_ast para consultar os DataFrames
2. NUNCA responda sem executar código Python
3. NUNCA invente ou crie dados sintéticos
4. Use df1, df2, df3, df4 para acessar os dados
5. Execute código pandas para cada pergunta"""

# Criar agente com TODOS os dataframes
agent = create_pandas_dataframe_agent(
    llm,
    [df_funcionarios, df_clientes, df_servicos, df_integrado],
    verbose=False,
    agent_type="tool-calling",
    allow_dangerous_code=True,
    return_intermediate_steps=True,
    prefix=prefix_prompt,
    suffix="""Lembre-se: SEMPRE execute código Python usando a ferramenta python_repl_ast.
NUNCA responda sem consultar os dados reais nos DataFrames.""",
    include_df_in_prompt=False,
    number_of_head_rows=2,
)


def consultar(pergunta, historico):
    """Processa pergunta e retorna resposta com passos intermediários."""
    try:
        resultado = agent.invoke({"input": pergunta})

        # Construir resposta com tool calls
        resposta_completa = ""

        # Mostrar passos intermediários (tool calls)
        if "intermediate_steps" in resultado:
            for i, (action, observation) in enumerate(
                resultado["intermediate_steps"], 1
            ):
                tool_name = action.tool if hasattr(action, "tool") else "Tool"
                tool_input = action.tool_input if hasattr(action, "tool_input") else ""
                resposta_completa += (
                    f"**🔧 Passo {i}: {tool_name}**\n```\n{tool_input}\n```\n\n"
                )

        # Adicionar resposta final
        output = resultado.get("output", str(resultado))
        resposta_completa += f"**✅ Resposta:**\n{output}"

        return resposta_completa
    except Exception as e:
        return f"❌ Erro: {str(e)}"


demo = gr.ChatInterface(
    fn=consultar,
    title="🤖 Agente de Consultas Analíticas",
    description="""Faça perguntas sobre funcionários, clientes e serviços. O agente mostrará os passos de raciocínio.

**Desenvolvido para:** Estatística Aplicada - Pós-graduação em Inteligência Artificial Aplicada  
**Autor:** João Victor Winderfeld Bussolotto""",
    examples=[
        "Quantos funcionários temos por cargo?",
        "Qual o salário médio por área?",
        "Top 5 clientes com maior contrato mensal",
        "Receita total estimada de todos os serviços",
        "Qual cargo tem o maior salário médio?",
        "Quantos clientes tem contrato acima de R$ 5000?",
        "Liste os 3 funcionários com maior salário",
        "Qual área tem mais funcionários?",
    ],
)
demo.launch(share=False)