import os
import time
import json
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import SQLDatabase
from huggingface_hub import InferenceClient
import gradio as gr
from dotenv import load_dotenv

load_dotenv()

CSV_FILE_PATH = "tabela_tiago_formated.csv"
SQL_DB_PATH = "data.db"
HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
LLAMA_MODEL = "meta-llama/Llama-3.3-70B-Instruct"
LLAMA_MODEL_FINAL = "Qwen/QwQ-32B"

hf_client = InferenceClient(
    provider="sambanova",
    api_key=HUGGINGFACE_API_KEY,
)

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

query_cache = {}

# Criar banco de dados SQL a partir do JSON
def create_or_load_sql_database(csv_path, sql_db_path):
    if os.path.exists(sql_db_path):
        print("Banco de dados SQL já existe. Carregando...")
        return create_engine(f"sqlite:///{sql_db_path}")
    else:
        print("Banco de dados SQL não encontrado. Criando..")
        engine = create_engine(f"sqlite:///{sql_db_path}")
        df = pd.read_csv(csv_path, sep=";", on_bad_lines="skip")
        print(f"CSV carregado: {len(df)} linhas, {len(df.columns)} colunas")
        df.to_sql("tabela_tiago_formated", engine, index=False, if_exists="replace")
        print("Banco de dados SQL criado com sucesso!")
        return engine

engine = create_or_load_sql_database(CSV_FILE_PATH, SQL_DB_PATH)
db = SQLDatabase(engine=engine)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
sql_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True, max_iterations=20, return_intermediate_steps=True)

def generate_initial_context(db_sample):
    return (
        f"Você é um assistente que gera queries SQL objetivas e eficientes. Sempre inclua LIMIT 15 nas queries. Aqui está o banco de dados:\n\n"
        ###f"Colunas: {', '.join(db_sample.columns)}\n"
        f"Exemplos do banco de dados:\n{db_sample.head().to_string(index=False)}\n\n"
        "\n***IMPORTANTE***: Detecte automaticamente o idioma da pergunta do usuário e responda sempre no mesmo idioma.\n"
        "\nRetorne apenas a pergunta e a query SQL mais eficiente para entregar ao agent SQL do LangChain para gerar uma resposta para a pergunta. O formato deve ser:\n"
        "\nPergunta: <pergunta do usuário>\n"
        "\nOpção de Query SQL:\n<query SQL>"
        "\nIdioma: <idioma>"
    )

def is_greeting(user_query):
    greetings = ["olá", "oi", "bom dia", "boa tarde", "boa noite", "oi, tudo bem?"]
    return user_query.lower().strip() in greetings

def query_with_llama(user_query, db_sample, chat_history):
    initial_context = generate_initial_context(db_sample)

    #formatted_history = ""  # Não inclui histórico temporariamente

    formatted_history = "\n".join(
        [f"{msg['role'].capitalize()}: {msg['content']}" for msg in chat_history[-1:]]
    )

    full_prompt = f"{initial_context}\n\nHistórico recente:\n{formatted_history}\n\nPergunta do usuário:\n{user_query}"

    print("------- Modelo Llama 3.3: Gerando query SQL -------")
    print(f"[DEBUG] Contexto enviado ao Llama:\n{full_prompt}\n")
    start_time = time.time()

    try:
        response = hf_client.chat.completions.create(
            model=LLAMA_MODEL,
            messages=[{"role": "system", "content": full_prompt}],
            max_tokens=800,
            stream=False
        )
        llama_response = response["choices"][0]["message"]["content"]
        end_time = time.time()
        print(f"[DEBUG] Resposta do Llama: {llama_response.strip()}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
        chat_history.append({"role": "assistant", "content": llama_response})
        return llama_response.strip(), chat_history
    except Exception as e:
        print(f"[ERRO] Falha ao interagir com o Llama: {e}")
        return None, chat_history

def refine_response_with_llama(user_query, sql_response):
    prompt = f"""
    Você é um assistente especializado em consolidar e explicar as informações. Abaixo estão as entradas:\n
    Pergunta: 
    {user_query}
    Resposta do Agente:
    {sql_response}\n
    """

    print("------- Modelo Final: Refinando resposta -------")
    print(f"[DEBUG] Contexto enviado ao Llama Final:\n{prompt}\n")
    start_time = time.time()

    try:
        response = hf_client.chat.completions.create(
            model=LLAMA_MODEL_FINAL,
            messages=[{"role": "system", "content": prompt}],
            max_tokens=8000
        )
        
        final_response = response["choices"][0]["message"]["content"]
        end_time = time.time()
        print(f"[DEBUG] Resposta do Llama Final: {final_response.strip()}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
        return final_response.strip()
        
    except Exception as e:
        print(f"[ERRO] Falha ao interagir com o Llama Final: {e}")
        return "Erro ao consolidar a resposta. Por favor, tente novamente."

def query_sql_agent(user_query, chat_history):
    try:
        if user_query in query_cache:
            print("[CACHE] Retornando resposta do cache para a consulta: {user_query}")
            return query_cache[user_query], chat_history

        if is_greeting(user_query):
            greeting_response = "Olá! Estou aqui para ajudar com suas consultas. Pergunte algo relacionado aos dados carregados no agente!"
            return greeting_response, chat_history

        column_data = pd.read_sql_query("SELECT * FROM tabela_tiago_formated LIMIT 10", engine)

        llama_instruction, chat_history = query_with_llama(user_query, column_data, chat_history)
        if not llama_instruction:
            return "Erro: O modelo Llama não conseguiu gerar uma instrução válida.", chat_history

        print("------- Agent SQL: Executando query -------")
        print(f"[DEBUG] Instrução gerada pelo Llama:\n{llama_instruction}\n")
        start_time = time.time()
        response = sql_agent.invoke({"input": llama_instruction})
        sql_response = response.get("output", "Erro ao obter a resposta do agente...")
        end_time = time.time()
        print(f"[DEBUG] Resposta do Agent SQL: {sql_response}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")

        final_response = refine_response_with_llama(user_query, sql_response)
        query_cache[user_query] = final_response  # Armazenar no cache
        return final_response, chat_history

    except Exception as e:
        return f"Erro ao consultar o agente SQL: {e}", chat_history

def gradio_interface():
    chat_history = []

    def respond(user_input):
        nonlocal chat_history
        response, chat_history = query_sql_agent(user_input, chat_history)
        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": response})
        return [{"role": "user", "content": user_input}, {"role": "assistant", "content": response}]

    with gr.Blocks() as demo:
        gr.Markdown("# Agent")
        chatbot = gr.Chatbot(label="Result Agent", type="messages")
        user_input = gr.Textbox(label="Digite sua pergunta")
        submit_button = gr.Button("Enviar")

        submit_button.click(respond, inputs=user_input, outputs=chatbot)

    return demo

if __name__ == "__main__":
    demo = gradio_interface()
    demo.launch(share=False)