import os
import json
from dotenv import load_dotenv
from retriever import db_chroma
from prompts import PROMPT_TEMPLATE, CLASSIFIER_PROMPT
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

def main():
    # Bom modelo para classificação e respostas
    model = ChatOpenAI(model_name='gpt-4o-mini', openai_api_key=os.getenv("OPENAI_API_KEY"))

    with open('know_graph.json', 'r', encoding='utf-8') as f:
        knowledge_graph_data = json.dumps(json.load(f), ensure_ascii=False)

    chat_history = ""  # Histórico inicial 
    
    while True:
        query = input("Digite sua pergunta (para sair digite 'sair'): ")
        if query.lower() == 'sair':
            break

        # Classificação inicial
        prompt_classifier = ChatPromptTemplate.from_template(CLASSIFIER_PROMPT)
        classifier_input = prompt_classifier.format(
            knowledge_graph=knowledge_graph_data,
            chat_history=chat_history,
            question=query
        )

        classification_response = model.invoke(classifier_input).content

        cleaned_response = classification_response.replace("```json", "").replace("```", "").strip()
        try:
            classification = json.loads(cleaned_response)
        except:
            classification = {"estagio": "Nao identificado", "objecao_id": "nenhuma", "argumento_base": ""}

        # Busca de contexto relevante
        docs_chroma = db_chroma.similarity_search(query, k=5)
        context_text = "\n\n".join([doc.page_content for doc in docs_chroma])

        diretriz_venda = f"\n[DIRETRIZ ESTRATÉGICA]: Estágio: {classification.get('estagio')}. " \
                         f"Objeção: {classification.get('objecao_id')}. " \
                         f"Argumento sugerido: {classification.get('argumento_base')}"
        
        full_context = context_text + diretriz_venda

        prompt_final = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
        final_input = prompt_final.format(
            context=full_context,
            chat_history=chat_history,
            question=query
        )

        response = model.invoke(final_input)
        
        print(f"\n--- LOG DE VENDAS ---")
        print(f"📍 Estágio: {classification.get('estagio')}")
        print(f"⚠️ Objeção: {classification.get('objecao_id')}")
        print(f"---------------------\n")
        print('***Resposta:*** \n\n' + response.content)

        # Atualizar histórico (simplificado)
        chat_history += f"\nUsuário: {query}\nAssistente: {response.content}"


        print("\n" + "-"*50 + "\n")
if __name__ == "__main__":
    main()