import os
import msgpack
import numpy as np
from scipy.sparse import csr_matrix, vstack, save_npz, load_npz
from some_wikipedia_loader_module import WikipediaLoader
from some_chunking_module import chunk_documents, SparseEmbedding

class GameScribesCacheEmbeddings:
    def __init__(self, namespace, queries, dense_model, sparse_model, hf_home=None):
        self.namespace = namespace
        self.queries = queries
        self.dense_model = dense_model
        self.sparse_model = sparse_model
        self.hf_home = hf_home or os.getenv('HF_HOME')
        self.embeddings_path = os.path.join(self.hf_home, 'embeddings')
        self.chunks_path = os.path.join(self.embeddings_path, f'{self.namespace}_chunks.msgpack')
        self.dense_path = os.path.join(self.embeddings_path, f'{self.namespace}_dense.npz')
        self.sparse_path = os.path.join(self.embeddings_path, f'{self.namespace}_sparse.npz')
        
        if not os.path.exists(self.embeddings_path):
            os.mkdir(self.embeddings_path)

    def load_or_create_chunks(self):
        if not os.path.exists(self.chunks_path):
            docs = self._load_documents()
            chunks, _, _ = chunk_documents(docs, self.dense_model, self.sparse_model)
            self._save_chunks(chunks)
        else:
            chunks = self._load_chunks()
        return chunks

    def load_or_create_dense_embeddings(self):
        if not os.path.exists(self.dense_path):
            docs = self._load_documents()
            _, dense_embeddings, _ = chunk_documents(docs, self.dense_model, self.sparse_model)
            self._save_dense_embeddings(dense_embeddings)
        else:
            dense_embeddings = self._load_dense_embeddings()
        return dense_embeddings

    def load_or_create_sparse_embeddings(self):
        if not os.path.exists(self.sparse_path):
            docs = self._load_documents()
            _, _, sparse_embeddings = chunk_documents(docs, self.dense_model, self.sparse_model)
            self._save_sparse_embeddings(sparse_embeddings)
        else:
            sparse_embeddings = self._load_sparse_embeddings()
        return sparse_embeddings

    def _load_documents(self):
        docs = []
        for query in self.queries:
            docs.extend(WikipediaLoader(query=query).load())
        return docs

    def _save_chunks(self, chunks):
        with open(self.chunks_path, "wb") as outfile:
            packed = msgpack.packb(chunks, use_bin_type=True)
            outfile.write(packed)

    def _save_dense_embeddings(self, dense_embeddings):
        np.savez_compressed(self.dense_path, *dense_embeddings)

    def _save_sparse_embeddings(self, sparse_embeddings):
        max_index = max(np.max(embedding.indices) for embedding in sparse_embeddings)
        sparse_matrices = []
        for embedding in sparse_embeddings:
            data = embedding.values
            indices = embedding.indices
            indptr = np.array([0, len(data)])
            matrix = csr_matrix((data, indices, indptr), shape=(1, max_index + 1))
            sparse_matrices.append(matrix)
        combined_sparse_matrix = vstack(sparse_matrices)
        save_npz(self.sparse_path, combined_sparse_matrix)

    def _load_chunks(self):
        with open(self.chunks_path, "rb") as data_file:
            byte_data = data_file.read()
        return msgpack.unpackb(byte_data, raw=False)

    def _load_dense_embeddings(self):
        return list(np.load(self.dense_path).values())

    def _load_sparse_embeddings(self):
        sparse_embeddings = []
        loaded_sparse_matrix = load_npz(self.sparse_path)
        for i in range(loaded_sparse_matrix.shape[0]):
            row = loaded_sparse_matrix.getrow(i)
            values = row.data
            indices = row.indices
            embedding = SparseEmbedding(values, indices)
            sparse_embeddings.append(embedding)
        return sparse_embeddings