import os
import json
import torch
from sentence_transformers import SentenceTransformer
from .utils import get_md5


class ToolRAGModel:
    def __init__(self, rag_model_name):
        self.rag_model_name = rag_model_name
        self.rag_model = None
        self.tool_desc_embedding = None
        self.tool_name = None
        self.tool_embedding_path = None
        self.load_rag_model()

    def load_rag_model(self):
        self.rag_model = SentenceTransformer(self.rag_model_name)
        self.rag_model.max_seq_length = 4096
        self.rag_model.tokenizer.padding_side = "right"

    def load_tool_desc_embedding(self, toolbox):
        self.tool_name, _ = toolbox.refresh_tool_name_desc(enable_full_desc=True)
        all_tools_str = [json.dumps(each) for each in toolbox.prepare_tool_prompts(toolbox.all_tools)]
        md5_value = get_md5(str(all_tools_str))
        print("Computed MD5 for tool embedding:", md5_value)

        self.tool_embedding_path = os.path.join(
            os.path.dirname(__file__),
            self.rag_model_name.split("/")[-1] + f"_tool_embedding_{md5_value}.pt"
        )

        if os.path.exists(self.tool_embedding_path):
            try:
                self.tool_desc_embedding = torch.load(self.tool_embedding_path, map_location="cpu")
                assert len(self.tool_desc_embedding) == len(toolbox.all_tools), \
                    "Tool count mismatch with loaded embeddings."
                print("\033[92mLoaded cached tool_desc_embedding.\033[0m")
                return
            except Exception as e:
                print(f"⚠️ Failed loading cached embeddings: {e}")
                self.tool_desc_embedding = None

        print("\033[93mGenerating new tool_desc_embedding...\033[0m")
        self.tool_desc_embedding = self.rag_model.encode(
            all_tools_str, prompt="", normalize_embeddings=True
        )

        torch.save(self.tool_desc_embedding, self.tool_embedding_path)
        print(f"\033[92mSaved new tool_desc_embedding to {self.tool_embedding_path}\033[0m")

    def rag_infer(self, query, top_k=5):
        torch.cuda.empty_cache()
        queries = [query]
        query_embeddings = self.rag_model.encode(
            queries, prompt="", normalize_embeddings=True
        )
        if self.tool_desc_embedding is None:
            raise RuntimeError("❌ tool_desc_embedding is not initialized. Did you forget to call load_tool_desc_embedding()?")

        scores = self.rag_model.similarity(
            query_embeddings, self.tool_desc_embedding
        )
        top_k = min(top_k, len(self.tool_name))
        top_k_indices = torch.topk(scores, top_k).indices.tolist()[0]
        top_k_tool_names = [self.tool_name[i] for i in top_k_indices]
        return top_k_tool_names