import logging
import os
import pickle
import time
from typing import Any
from typing import Dict
from typing import List

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

from rakam_systems.core import VSFile

# Configure logging
logging.basicConfig(level=logging.INFO)


class VectorStores:
    """
    A class for managing multiple vector stores using FAISS and SentenceTransformers.
    """

    def __init__(self, base_index_path: str, embedding_model: str) -> None:
        """
        Initializes the VectorStores with the specified base index path and embedding model.

        :param base_index_path: Base path to store the FAISS indexes.
        :param embedding_model: Pre-trained SentenceTransformer model name.
        """
        self.base_index_path = base_index_path
        if not os.path.exists(self.base_index_path):
            os.makedirs(self.base_index_path)

        self.embedding_model = SentenceTransformer(embedding_model)
        self.stores = {}
        self.load_all_stores()

    def load_all_stores(self) -> None:
        """
        Loads all vector stores from the base directory.
        """
        for store_name in os.listdir(self.base_index_path):
            store_path = os.path.join(self.base_index_path, store_name)
            if os.path.isdir(store_path):
                self.stores[store_name] = self.load_store(store_path)

    def load_store(self, store_path: str) -> Dict[str, Any]:
        """
        Loads a single vector store from the specified directory.

        :param store_path: Path to the store directory.
        :return: Dictionary containing the store's index, nodes, and metadata.
        """
        store = {}
        try:
            store["index"] = faiss.read_index(os.path.join(store_path, "index"))
            with open(
                os.path.join(store_path, "category_index_mapping.pkl"), "rb"
            ) as f:
                store["category_index_mapping"] = pickle.load(f)
            with open(
                os.path.join(store_path, "metadata_index_mapping.pkl"), "rb"
            ) as f:
                store["metadata_index_mapping"] = pickle.load(f)
            with open(os.path.join(store_path, "nodes.pkl"), "rb") as f:
                store["nodes"] = pickle.load(f)
            logging.info(f"Store loaded successfully from {store_path}.")
        except Exception as e:
            logging.warning(f"Error loading store from {store_path}: {e}")
        return store

    def predict_embeddings(self, query: str) -> np.ndarray:
        """
        Predicts embeddings for a given query using the embedding model.

        :param query: Query string to encode.
        :return: Embedding vector for the query.
        """
        logging.info(f"Predicting embeddings for query: {query}")
        query_embedding = self.embedding_model.encode(query)
        query_embedding = np.asarray([query_embedding], dtype="float32")
        return query_embedding

    def search(
        self, store_name: str, query: str, distance_type="cosine", number=5
    ) -> dict:
        """
        Searches the specified store for the closest embeddings to the query.

        :param store_name: Name of the store to search in.
        :param query: Query string to search for.
        :param distance_type: Distance metric to use (default is "cosine").
        :param number: Number of closest embeddings to return.
        :return: Dictionary of search results with suggestion texts and distances.
        """
        logging.info(
            f"Searching in store: {store_name} for query: {query} with distance type: {distance_type} and number of results: {number}"
        )

        store = self.stores.get(store_name)
        if not store:
            raise ValueError(f"No store found with name: {store_name}")

        query_embedding = self.predict_embeddings(query)
        suggested_nodes = []

        if distance_type == "cosine":
            faiss.normalize_L2(query_embedding)
        D, I = store["index"].search(query_embedding, number)

        seen_texts = set()
        valid_suggestions = {}
        count = 0
        for i, id_ in enumerate(I[0]):
            if count >= number:
                break
            if id_ != -1 and id_ in store["category_index_mapping"]:
                suggestion_text = store["category_index_mapping"][id_]
                node_metadata = store["metadata_index_mapping"][id_]
                suggested_nodes.append(store["nodes"][id_])
                if suggestion_text not in seen_texts:
                    seen_texts.add(suggestion_text)
                    valid_suggestions[str(id_)] = (
                        node_metadata,
                        suggestion_text,
                        float(D[0][i]),
                    )
                    count += 1

        while count < number:
            valid_suggestions[f"placeholder_{count}"] = (
                "No suggestion available",
                float("inf"),
            )
            count += 1

        logging.info(f"Search results: {valid_suggestions}")
        return valid_suggestions, suggested_nodes

    def get_embeddings(self, sentences: List[str], parallel: bool = True) -> np.ndarray:
        """
        Generates embeddings for a list of sentences.

        :param sentences: List of sentences to encode.
        :param parallel: Whether to use parallel processing (default is True).
        :return: Embedding vectors for the sentences.
        """
        logging.info(f"Generating embeddings for {len(sentences)} sentences.")
        start = time.time()
        if parallel:
            os.environ["TOKENIZERS_PARALLELISM"] = "false"
            pool = self.embedding_model.start_multi_process_pool(
                target_devices=["cpu"] * 5
            )
            embeddings = self.embedding_model.encode_multi_process(
                sentences, pool, batch_size=16
            )
            self.embedding_model.stop_multi_process_pool(pool)
        else:
            os.environ["TOKENIZERS_PARALLELISM"] = "true"
            embeddings = self.embedding_model.encode(
                sentences,
                batch_size=32,
                show_progress_bar=True,
                convert_to_tensor=True,
            )
        logging.info(
            f"Time taken to encode {len(sentences)} items: {round(time.time() - start, 2)} seconds"
        )
        return embeddings.cpu().detach().numpy()

    def create_from_files(self, stores_files: Dict[str, List[VSFile]]) -> None:
        """
        Creates FAISS indexes from dictionaries of store names and VSFile objects.

        :param stores_files: Dictionary where keys are store names and values are lists of VSFile objects.
        """
        for store_name, files in stores_files.items():
            logging.info(f"Creating FAISS index for store: {store_name}")
            text_chunks = []
            metadata = []
            nodes = []

            for file in files:
                for node in file.nodes:
                    nodes.append(node)
                    text_chunks.append(node.content)
                    formatted_metadata = {
                        "node_id": node.metadata.node_id,
                        "source_file_uuid": node.metadata.source_file_uuid,
                        "position": node.metadata.position,
                        "custom": node.metadata.custom,
                    }
                    metadata.append(formatted_metadata)

            self._create_and_save_index(store_name, nodes, text_chunks, metadata)

    def create_from_nodes(self, store_name: str, nodes: List[Any]) -> None:
        """
        Creates a FAISS index from a list of nodes and stores it under the given store name.

        :param store_name: Name of the store to create.
        :param nodes: List of nodes containing the content and metadata.
        """
        logging.info(f"Creating FAISS index for store: {store_name}")
        text_chunks = []
        metadata = []

        for node in nodes:
            text_chunks.append(node.content)
            formatted_metadata = {
                "node_id": node.metadata.node_id,
                "source_file_uuid": node.metadata.source_file_uuid,
                "position": node.metadata.position,
                "custom": node.metadata.custom,
            }
            metadata.append(formatted_metadata)

        self._create_and_save_index(store_name, nodes, text_chunks, metadata)

    def _create_and_save_index(
        self,
        store_name: str,
        nodes: List[Any],
        text_chunks: List[str],
        metadata: List[Dict[str, Any]],
    ) -> None:
        """
        Helper function to create and save a FAISS index.

        :param store_name: Name of the store to create.
        :param nodes: List of nodes.
        :param text_chunks: List of text chunks to encode and index.
        :param metadata: List of metadata associated with the text chunks.
        """
        store_path = os.path.join(self.base_index_path, store_name)
        if not os.path.exists(store_path):
            os.makedirs(store_path)

        data_embeddings = self.get_embeddings(sentences=text_chunks, parallel=False)
        category_index_mapping = dict(zip(range(len(text_chunks)), text_chunks))

        # Save category index mapping to file
        with open(os.path.join(store_path, "category_index_mapping.pkl"), "wb") as f:
            pickle.dump(category_index_mapping, f)

        # Save nodes to file
        with open(os.path.join(store_path, "nodes.pkl"), "wb") as f:
            pickle.dump(nodes, f)

        index = faiss.IndexIDMap(faiss.IndexFlatIP(data_embeddings.shape[1]))
        faiss.normalize_L2(data_embeddings)
        index.add_with_ids(
            data_embeddings, np.array(list(category_index_mapping.keys()))
        )
        faiss.write_index(index, os.path.join(store_path, "index"))

        metadata_index_mapping = dict(zip(range(len(text_chunks)), metadata))

        # Save metadata nodes to file
        with open(os.path.join(store_path, "metadata_index_mapping.pkl"), "wb") as f:
            pickle.dump(metadata_index_mapping, f)

        self.stores[store_name] = {
            "index": index,
            "nodes": nodes,
            "category_index_mapping": category_index_mapping,
            "metadata_index_mapping": metadata_index_mapping,
        }
        logging.info(
            f"FAISS index for store {store_name} created and saved successfully."
        )
