import os
import json
import logging
from typing import List

from txtai.embeddings import Embeddings

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class EmbeddingsManager:
    def __init__(self, base_path: str = "./indexes", model_path: str = "avsolatorio/GIST-all-MiniLM-L6-v2"):
        """
        Initializes the EmbeddingsManager.

        Args:
            base_path (str): Base directory to store indices.
            model_path (str): Path or identifier for the embeddings model.
        """
        self.base_path = base_path
        os.makedirs(self.base_path, exist_ok=True)
        self.model_path = model_path
        self.embeddings = Embeddings({"path": self.model_path})
        logger.info(f"Embeddings model loaded from '{self.model_path}'. Base path set to '{self.base_path}'.")

    def create_index(self, index_id: str, documents: List[str]) -> None:
        """
        Creates a new embeddings index with the provided documents.

        Args:
            index_id (str): Unique identifier for the index.
            documents (List[str]): List of documents to be indexed.

        Raises:
            ValueError: If the index already exists.
            Exception: For any other errors during indexing or saving.
        """
        index_path = os.path.join(self.base_path, index_id)
        if os.path.exists(index_path):
            logger.error(f"Index with index_id '{index_id}' already exists at '{index_path}'.")
            raise ValueError(f"Index with index_id '{index_id}' already exists.")

        try:
            # Prepare documents for txtai indexing
            document_tuples = [(i, text, None) for i, text in enumerate(documents)]
            self.embeddings.index(document_tuples)
            logger.info(f"Documents indexed for index_id '{index_id}'.")

            # Create index directory
            os.makedirs(index_path, exist_ok=True)

            # Save embeddings
            self.embeddings.save(os.path.join(index_path, "embeddings"))
            logger.info(f"Embeddings saved to '{os.path.join(index_path, 'embeddings')}'.")

            # Save document list
            with open(os.path.join(index_path, "document_list.json"), "w", encoding='utf-8') as f:
                json.dump(documents, f, ensure_ascii=False, indent=4)
            logger.info(f"Document list saved to '{os.path.join(index_path, 'document_list.json')}'.")

            logger.info(f"Index '{index_id}' created and saved successfully.")
        except Exception as e:
            logger.error(f"Failed to create index '{index_id}': {e}")
            raise Exception(f"Failed to create index '{index_id}': {e}")

    def query_index(self, index_id: str, query: str, num_results: int = 5) -> List[str]:
        """
        Queries an existing embeddings index.

        Args:
            index_id (str): Unique identifier for the index to query.
            query (str): The search query.
            num_results (int): Number of top results to return.

        Returns:
            List[str]: List of top matching documents.

        Raises:
            FileNotFoundError: If the index does not exist.
            Exception: For any other errors during querying.
        """
        index_path = os.path.join(self.base_path, index_id)
        if not os.path.exists(index_path):
            logger.error(f"Index '{index_id}' not found at '{index_path}'.")
            raise FileNotFoundError(f"Index '{index_id}' not found.")

        try:
            # Load embeddings from the index
            self.embeddings.load(os.path.join(index_path, "embeddings"))
            logger.info(f"Embeddings loaded from '{os.path.join(index_path, 'embeddings')}' for index '{index_id}'.")

            # Load document list
            document_list_path = os.path.join(index_path, "document_list.json")
            if not os.path.exists(document_list_path):
                logger.error(f"Document list not found at '{document_list_path}'.")
                raise FileNotFoundError(f"Document list not found for index '{index_id}'.")

            with open(document_list_path, "r", encoding='utf-8') as f:
                document_list = json.load(f)
            logger.info(f"Document list loaded from '{document_list_path}'.")

            # Perform the search
            results = self.embeddings.search(query, num_results)
            queried_texts = [document_list[idx[0]] for idx in results]
            logger.info(f"Query executed successfully on index '{index_id}'. Retrieved {len(queried_texts)} results.")

            return queried_texts
        except Exception as e:
            logger.error(f"Failed to query index '{index_id}': {e}")
            raise Exception(f"Failed to query index '{index_id}': {e}")