import streamlit as st
import os
import json
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from prompts import PROMPT_TEMPLATE, CLASSIFIER_PROMPT

# Configuração da página Streamlit
st.set_page_config(page_title="Central da Visão - Assistente", page_icon="👁️")
st.title("👁️ Assistente Virtual - Central da Visão")


load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not OPENAI_API_KEY:
    st.error("A chave da API da OpenAI não foi encontrada. Verifique as Secrets do Space.")
    st.stop()

# Configurações do banco de dados Chroma
CHROMA_PATH = "chroma_db"
DOCS_PATH = "docs"

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

@st.cache_resource
def load_or_create_db():
    """
    Carrega o banco Chroma se existir, ou cria do zero se não encontrar.
    """
    
    model_name = "Qwen/Qwen3-Embedding-0.6B"
    model_kwargs = {"device": "cpu",
                    "trust_remote_code": True}
    encode_kwargs = {"normalize_embeddings": True}

    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )

    # Verifica se o banco já existe
    if os.path.exists(CHROMA_PATH) and os.listdir(CHROMA_PATH):
        print("Banco de dados encontrado. Carregando...")
        return Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
    
    else:
        print("Banco não encontrado. Criando nova base de conhecimento...")
        
        # Garante que a pasta docs existe
        if not os.path.exists(DOCS_PATH):
            os.makedirs(DOCS_PATH)
            st.warning(f"A pasta '{DOCS_PATH}' estava vazia. O assistente não tem documentos para ler.")
            return Chroma(embedding_function=embeddings)

        # Carrega os arquivos .txt
        loader = DirectoryLoader(
            DOCS_PATH,
            glob="**/*.txt",
            loader_cls=TextLoader,
            loader_kwargs={"encoding": "utf-8"}
        )
        documents = loader.load()

        if not documents:
            st.warning("Nenhum arquivo .txt encontrado em 'docs/'.")
            return Chroma(embedding_function=embeddings)

        # Divide em chunks
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        chunks = text_splitter.split_documents(documents)

        # Cria e salva o banco
        db = Chroma.from_documents(
            documents=chunks, 
            embedding=embeddings, 
            persist_directory=CHROMA_PATH
        )
        return db

# Inicializa o banco de dados
db_chroma = load_or_create_db()

# Sessão de chat
if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "Olá! Sou o consultor virtual da Central da Visão. Como posso ajudar você hoje?"}
    ]

def get_response(query):
    historico = st.session_state.messages[-6:]

    chat_history_str = ""
    for msg in historico:
        prefix = "Usuário: " if msg["role"] == "user" else "Consultor: "
        chat_history_str += f"{prefix}{msg['content']}\n"

    model = ChatOpenAI(model_name='gpt-4o-mini', openai_api_key=OPENAI_API_KEY)

    prompt_classifier = ChatPromptTemplate.from_template(CLASSIFIER_PROMPT)
    classifier_input = prompt_classifier.format(
        knowledge_graph=knowledge_graph_data,
        chat_history=chat_history_str,
        question=query
    )

    classification_response = model.invoke(classifier_input).content

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

    # Busca de documentos relevantes
    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
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=full_context, 
                                    chat_history=chat_history_str, 
                                    question=query)
    # Modelo de linguagem
    response = model.invoke(prompt)
    
    return response.content, classification

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

if query := st.chat_input("Digite sua dúvida aqui..."):
    st.session_state.messages.append({"role": "user", "content": query})
    with st.chat_message("user"):
        st.markdown(query)

    with st.chat_message("assistant"):
        with st.spinner("Consultando base de dados..."):
            try:
                response_text, info_vendas = get_response(query)

                with st.expander("📊 Log de Vendas (Inteligência)"):
                    st.write(f"**Estágio:** {info_vendas.get('estagio')}")
                    st.write(f"**Objeção Identificada:** {info_vendas.get('objecao_id')}")
                    st.caption(f"Argumento base: {info_vendas.get('argumento_base')}")
                    
                st.markdown(response_text)
                st.session_state.messages.append({"role": "assistant", "content": response_text})
            except Exception as e:
                st.error(f"Ocorreu um erro ao processar: {e}")