import json
import os
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
from dotenv import load_dotenv
from sentence_transformers import CrossEncoder
from langchain_core.messages import HumanMessage, AIMessage, BaseMessage

load_dotenv()

catalog_path = os.path.join(os.path.dirname(__file__),'data','catalog.json')

def load_catalog() -> list[dict]:
    with open(catalog_path, "r", encoding="utf-8") as f:
        return json.load(f)
    
def create_document(catalog : list[dict]) -> list[Document]:
    
    docs = []

    for item in catalog:
        text = f"""
        Entity ID: {item.get('entity_id', '')}

        Assessment Name: {item.get('name', '')}

        Description:
        {item.get('description', '')}

        Job Levels:
        {', '.join(item.get('job_levels', []))}

        Languages:
        {', '.join(item.get('languages', []))}

        Duration:
        {item.get('duration', '')}

        Keys:
        {', '.join(item.get('keys', []))}

        Remote Testing:
        {item.get('remote', '')}

        Adaptive:
        {item.get('adaptive', '')}
        """

        metadata = {
            "entity_id": item.get("entity_id"),
            "name": item.get("name"),
            "job_levels": item.get("job_levels", []),
            "languages": item.get("languages", []),
            "keys": item.get("keys", []),
            "url": item.get("link")
        }

        docs.append(Document(page_content=text, metadata=metadata))

    return docs
 
def download_embeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") -> HuggingFaceEmbeddings:
    embedding_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    cache_folder="./cache"
    )

    return embedding_model

def download_reranker(model_name="cross-encoder/ms-marco-MiniLM-L-6-v2") -> CrossEncoder:
    reranker = CrossEncoder(model_name,cache_folder="./cache")
    return reranker

def deduplicate_docs(docs1,docs2):
    return list({doc.metadata['entity_id']: doc for doc in docs1 + docs2}.values())

def rerank_docs(reranker,query, docs,top_k=10):
    pairs = [(query, doc.page_content) for doc in docs]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, docs), reverse=True)
    return [doc for _, doc in ranked[:top_k]]

def convert_messages(messages) -> list[BaseMessage]:

    lc_messages = []

    for msg in messages:

        if msg.role == "user":

            lc_messages.append(
                HumanMessage(content=msg.content)
            )

        elif msg.role == "assistant":

            lc_messages.append(
                AIMessage(content=msg.content)
            )

    return lc_messages

def context_builder(docs : list[Document]) -> HumanMessage:

    context_docs = "\n\n".join([f'{doc.page_content}' for doc in docs])

    context_msg = HumanMessage(content=f""" Retrieved SHL Assessments: {context_docs} 
                           Use ONLY these retrieved assessments for reasoning, clarifications, comparisons and/or recommendation. """ )
    
    return context_msg