import os
import time

import torch.multiprocessing as mp

from colbert.infra.run import Run
from colbert.infra.config import ColBERTConfig, RunConfig
from colbert.infra.launcher import Launcher

from colbert.utils.utils import create_directory, print_message

from colbert.indexing.collection_indexer import reuse_prepare, sample, kmeans, encode, finalize

class Indexer:
    def __init__(self, checkpoint, config=None):
        """
           Use Run().context() to choose the run's configuration. They are NOT extracted from `config`.
        """

        self.index_path = None
        self.checkpoint = checkpoint
        self.checkpoint_config = ColBERTConfig.load_from_checkpoint(checkpoint)

        self.config = ColBERTConfig.from_existing(self.checkpoint_config, config, Run().config)
        self.configure(checkpoint=checkpoint)

    def configure(self, **kw_args):
        self.config.configure(**kw_args)

    def get_index(self):
        return self.index_path

    def erase(self):
        assert self.index_path is not None
        directory = self.index_path
        deleted = []

        for filename in sorted(os.listdir(directory)):
            filename = os.path.join(directory, filename)

            delete = filename.endswith(".json")
            delete = delete and ('metadata' in filename or 'doclen' in filename or 'plan' in filename)
            delete = delete or filename.endswith(".pt")
            
            if delete:
                deleted.append(filename)
        
        if len(deleted):
            print_message(f"#> Will delete {len(deleted)} files already at {directory} in 20 seconds...")
            time.sleep(20)

            for filename in deleted:
                os.remove(filename)

        return deleted

    def prepare(self, name, collection, overwrite=False, no_kmeans=False):
        assert overwrite in [True, False, 'reuse', 'resume']

        self.configure(collection=collection, index_name=name, resume=overwrite=='resume')
        self.configure(partitions=None)

        self.index_path = self.config.index_path_
        index_does_not_exist = (not os.path.exists(self.config.index_path_))

        assert (overwrite in [True, 'reuse', 'resume']) or index_does_not_exist, self.config.index_path_
        create_directory(self.config.index_path_)

        if overwrite is True:
            self.erase()

        if index_does_not_exist or overwrite != 'reuse':
            self.__launch(sample, collection, nospawn=self.config.gpus == 1)
            if not no_kmeans:
                self.__launch(kmeans, collection, nospawn=True)

        return self.index_path

    def encode(self, name, collection):
        self.configure(collection=collection, index_name=name, resume=True)
        self.configure(partitions=None)

        self.index_path = self.config.index_path_
            
        if self.config.reuse_centroids_from is not None: 
            plan_path = os.path.join(self.index_path, "plan.json")
            if not os.path.exists(plan_path):
                # do a light prepare step to create the plan.json file along and soft link the centriods 
                self.index_path = self.config.index_path_
                create_directory(self.config.index_path_)
                self.__launch(reuse_prepare, collection, nospawn=True)
            else: 
                print_message(f"############# Ignoring `reuse_centroids_from` since the plan in the index directory already exists -- will only try to resume indexing only")

        assert os.path.exists(self.config.index_path_), "Run first step `prepare` in advance."

        self.__launch(encode, collection, nospawn=self.config.gpus == 1)

        return self.index_path

    def finalize(self, name, collection):
        self.configure(collection=collection, index_name=name, resume=True)
        self.configure(bsize=64, partitions=None)

        self.index_path = self.config.index_path_
            
        

        self.__launch(finalize, collection, nospawn=True)

        return self.index_path


    def __launch(self, callee, collection, nospawn=False):
        manager = mp.Manager()
        shared_lists = [manager.list() for _ in range(self.config.nranks)]
        shared_queues = [manager.Queue(maxsize=1) for _ in range(self.config.nranks)]

        if not nospawn:
            launcher = Launcher(callee)
            launcher.launch(self.config, collection, shared_lists, shared_queues)
        else:
            new_config = type(self.config).from_existing(self.config, RunConfig.from_existing(Run().config))
            callee(new_config, collection, shared_lists, shared_queues)
