import os
from typing import List, Dict
from luntaiDs.CommonTools.accessor import loadJSON, toJSON

class _BaseModelRegistry:
    """Base class for Simple Model Registry
    model versioning is achieved by tracking model_id and its configuration

    configuration format:
    {
        "prod" : "MODEL_ID",
        "archive" : {
            "MODEL_ID_A" : {CONFIG_A},
            "MODEL_ID_B" : {CONFIG_B},
        }
    }
    """

    def load_config(self) -> dict:
        """load configuration
        
        :return dict: configuration in dictionary format
        """
        raise NotImplementedError("")

    def save_config(self, config:dict):
        """save configuration dictionary

        :param dict config: configuration in dictionary format
        """
        raise NotImplementedError("")

    def get_prod_model_id(self) -> str:
        """get deployed (prod) model id

        :return str: model id that is deployed
        """
        return self.load_config().get("prod")

    def get_model_list(self) -> List[str]:
        """get a list of model ids

        :return List[str]: list of model ids registered
        """
        return list(self.load_config().get("archive").keys())

    def get_model_config(self, model_id: str) -> dict:
        """get specific configuration part for a given model id

        :param str model_id: model id of the model
        :return dict: the specific configuration part
        """
        assert model_id in self.get_model_list(), "Model Id not found in registry, please register one first"
        return self.load_config().get("archive").get(model_id)

    def load_model(self, model_id: str):
        """loading the model into memory

        :param str model_id: the model id for the model
        :return _type_: model object into memory
        """
        config = self.get_model_config(model_id)
        return self.load_model_by_config(config)

    def load_prod(self):
        """load deployed model into memory

        :return _type_: deployed model object into memory
        """
        return self.load_model(self.get_prod_model_id())

    def register(self, model_id: str, *args, **kws):
        """register your model into the system

        :param str model_id: your model id
        """
        assert model_id not in self.get_model_list(), "Model Id already registered, please try another one"
        model_config = self.save_model_and_generate_config(model_id, *args, **kws)
        config = self.load_config()
        config['archive'][model_id] = model_config
        self.save_config(config)

    def remove(self, model_id: str):
        """delete model from registry

        :param str model_id: model id to be removed
        """
        assert model_id in self.get_model_list(), "Model Id not found"

        config = self.load_config()
        self.delete_model_files(model_id)
        config['archive'].pop(model_id)
        if config['prod'] == model_id:
            config['prod'] = None
        self.save_config(config)

    def deploy(self, model_id: str):
        """deploy selected model to prod, simply achieved by setting it to prod field

        :param str model_id: model id to be deployed
        """
        # check if model id exists
        assert model_id in self.get_model_list(), "Model Id not found in registry, please register one first"
        config = self.load_config()
        config["prod"] = model_id
        self.save_config(config)

    def delete_model_files(self, model_id: str):
        """delete model related files/data

        :param str model_id: model id to be deleted
        """
        raise NotImplementedError("")

    def load_model_by_config(self, config: dict):
        """load the model using configuration file

        :param dict config: configuration for the model
        """
        raise NotImplementedError("")

    def save_model_and_generate_config(self, model_id:str, *args, **kws) -> dict:
        """save model and generate configuration for this model

        :param str model_id: model id to be generated
        :return dict: the configuration generated by creating model into the system
        """
        raise NotImplementedError("")
    
    
class ModelRegistryLocalFS(_BaseModelRegistry):
    def __init__(self, config_js_path: str):
        """

        :param config_js_path: maybe not used if not local file system

        js format:
        {
            "prod" : "MODEL_ID",
            "archive" : {
                "MODEL_ID_A" : {CONFIG_A},
                "MODEL_ID_B" : {CONFIG_B},
            }
        }
        """
        self.config_js_path = config_js_path

    def load_config(self) -> dict:
        """load configuration
        
        :return dict: configuration in dictionary format
        """
        if os.path.exists(self.config_js_path):
            config = loadJSON(self.config_js_path)
        else:
            config = {
                "prod" : None,
                "archive": {}
            }
        return config

    def save_config(self, config:dict):
        """save configuration dictionary

        :param dict config: configuration in dictionary format
        """
        toJSON(config, self.config_js_path)