# chatbot_engine.py
# Duarte Grilo - 2201320 - Projeto Informático

from sentence_transformers import SentenceTransformer, util
from models.rag_engine import get_chroma_db
from config import EMBEDDING_MODEL

# 🔧 Carrega o modelo de embeddings definido no ficheiro de configuração
model = SentenceTransformer(EMBEDDING_MODEL)

# 📚 Acede à base local de embeddings via ChromaDB (LangChain)
collection = get_chroma_db()


def busca_local_heuristica(pergunta: str, k: int = 5, retorna=False):
    print(f"🔍 Pergunta: {pergunta}")

    emb_pergunta = model.encode(pergunta, convert_to_tensor=True)
    resultados = collection.similarity_search(query=pergunta, k=k)

    # Depuração: Imprimir os resultados para verificar o conteúdo
    print(f"🔎 Melhores documentos encontrados: {resultados}")

    # Verificar se 'resultados' é uma lista de tuplas e extrair os documentos corretamente
    if isinstance(resultados, list) and isinstance(resultados[0], tuple):
        # Se for uma lista de tuplas, extraímos o texto (assumindo que o texto está na segunda posição)
        melhores_docs = [d[0]['page_content'] for d in resultados if isinstance(d, tuple) and len(d) > 0 and 'page_content' in d[0]]

    else:
        # Caso contrário, tentamos simplesmente acessar a chave 'page_content'
        melhores_docs = [d['page_content'] for d in resultados if 'page_content' in d]

    # Se não houver documentos válidos, retorna uma resposta padrão
    if not melhores_docs:
        return "Desculpe, não encontrei informações relevantes."

    # Encontra o melhor documento baseado na similaridade
    melhor_doc = max(melhores_docs, key=lambda d: util.pytorch_cos_sim(emb_pergunta, model.encode([d], convert_to_tensor=True)).mean().item())

    if retorna:
        return melhor_doc[:800]  # devolve a resposta em vez de imprimir

    print(f"\n✅ Melhor resposta:\n{melhor_doc[:800]}")