import os
import re

import mlflow
from loguru import logger

REQUIRED_TZ = "Asia/Kolkata"

mlflow_tracking_uri = 'https://qa.unifytwin.com/mlflow/'
AZURE_STORAGE_CONNECTION_STRING = 'DefaultEndpointsProtocol=https;AccountName=azrmlilensqa006382180551;AccountKey=tD' \
                                  'GOKfiZ2svfoMvVmS0Fbpf0FTHfTq4wKYuDX7cAxlhve/3991QuzdvJHm9vWc+lo6mtC+x9yPSghWNR' \
                                  '4+gacg==;EndpointSuffix=core.windows.net'
AZURE_STORAGE_ACCESS_KEY = 'tDGOKfiZ2svfoMvVmS0Fbpf0FTHfTq4wKYuDX7cAxlhve/3991QuzdvJHm9vWc+lo6mtC+x9yPSghWNR4+gacg=='

os.environ["MLFLOW_TRACKING_USERNAME"] = 'mlflow'
os.environ["MLFLOW_TRACKING_PASSWORD"] = 'MlFlOwQA#4321'
os.environ["AZURE_STORAGE_CONNECTION_STRING"] = AZURE_STORAGE_CONNECTION_STRING
os.environ["AZURE_STORAGE_ACCESS_KEY"] = AZURE_STORAGE_ACCESS_KEY
mlflow.set_tracking_uri(mlflow_tracking_uri)
mlflow.set_registry_uri(mlflow_tracking_uri)
client = mlflow.tracking.MlflowClient()


class MlFlowUtil:
    @staticmethod
    def log_model(model, model_name):
        try:
            mlflow.sklearn.log_model(model, model_name)
            logger.info("logged the model")
            return True
        except Exception as e:
            logger.exception(str(e))

    @staticmethod
    def log_metrics(metrics):
        try:
            updated_metric = {}
            for key, value in metrics.items():
                key = re.sub(r"[([{})\]]", "", key)
                updated_metric[key] = value
            mlflow.log_metrics(updated_metric)
            return True
        except Exception as e:
            logger.exception(str(e))

    @staticmethod
    def log_hyper_param(hyper_params):
        try:
            mlflow.log_params(hyper_params)
            return True
        except Exception as e:
            logger.exception(str(e))

    @staticmethod
    def set_tag(child_run_id, key, value):
        try:
            client.set_tag(run_id=child_run_id, key=key, value=value)
        except Exception as e:
            logger.exception(f"Exception while setting the tag - {e}")


class ModelLoaderSaver:
    def __init__(self, model, metrics, hyperparams, experiment_name, parent_run_name, model_save_name, model_type):
        self.model = model
        self.metrics = metrics
        self.hyperparams = hyperparams
        self.experiment_name = experiment_name
        self.parent_run_name = parent_run_name
        self.model_save_name = model_save_name
        self.model_type = model_type
        self._mfu_ = MlFlowUtil()

    def check_create_experiment(self):
        """
        check if experiment exists, if not creates a new experiment
        :return: experiment_id of the experiment
        """
        experiment_info = mlflow.get_experiment_by_name(self.experiment_name)
        if experiment_info is None:
            logger.info(f"No experiment found with name {self.experiment_name}, So creating one")
            mlflow.create_experiment(self.experiment_name)
        else:
            logger.info(f"Proceeding with existing Experiment {self.experiment_name}")
        mlflow.set_experiment(experiment_name=self.experiment_name)
        experiment_info = mlflow.get_experiment_by_name(self.experiment_name)
        experiment_id = experiment_info.experiment_id
        return experiment_id

    def check_create_parent_run(self, experiment_id):
        """
        check if a parent run exists in the experiment, if not create it with the mentioned parent run name
        :param experiment_id: Experiment id
        :return: returns the parent run id
        """
        parent_runs_df = mlflow.search_runs(experiment_id)
        run_key = 'tags.mlflow.runName'
        if run_key in parent_runs_df.columns:
            parent_runs_df = parent_runs_df[parent_runs_df[run_key] == self.parent_run_name]
        else:
            parent_runs_df = parent_runs_df.iloc[:0]
        if not parent_runs_df.empty:
            logger.info(f"Proceeding with existing Parent Run {self.parent_run_name}")
            return list(parent_runs_df['run_id'])[0]
        # no parent run found
        logger.info(f"No Parent Run present {self.parent_run_name}")
        with mlflow.start_run(experiment_id=experiment_id, run_name=self.parent_run_name) as run:
            logger.info(f"Creating the parent Run {self.parent_run_name} with Parent Run Id {run.info.run_id}")
            return run.info.run_id

    def check_create_child_run(self, experiment_id, parent_run_id):
        """
        check if a child run exists in the experiment id under the parent run id
        if exists take the child run id which has the model saved and validate when was it lastly trained.
        Based on the lastly trained see if you have to retrain or not. if retrain create a new child run
        else if no child run exists under the parent run id of experiment id, create a new child run
        :param experiment_id: experiment id
        :param parent_run_id: parent run id
        :return: child run id, retrain flag
        """
        child_runs_df = mlflow.search_runs(experiment_id, filter_string=f"tags.mlflow.parentRunId='{parent_run_id}'")
        if not child_runs_df.empty:
            logger.info(f"Already Child runs are present for Parent Run Id {parent_run_id}")
            child_run_id, retrain = self.get_latest_child_run(experiment_id, parent_run_id, child_runs_df)
            return child_run_id, retrain
        else:
            logger.info(f"Child runs are not present for Parent Run Id {parent_run_id}")
            with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
                with mlflow.start_run(experiment_id=experiment_id, nested=True) as child_run:
                    return child_run.info.run_id, True

    def get_latest_child_run(self, experiment_id, parent_run_id, runs_df):
        """
        Check if child runs are present. if not create a new child run. Otherwise, validate the last run time and
        create a new child run if retraining needed or take the last child run id which has model saved
        :param experiment_id: experiment id
        :param parent_run_id: parent run id
        :param runs_df: the child runs of the parent id
        :return: last child run id, retrain flag
        """
        history_key = 'tags.mlflow.log-model.history'
        if history_key in runs_df.columns:
            runs_df = runs_df[runs_df[history_key].notna()]
        else:
            runs_df = runs_df.iloc[:0]
        if runs_df.empty:
            logger.info("Existing Child Runs doesn't contain any model to run. So creating new child run")
            with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
                with mlflow.start_run(experiment_id=experiment_id, nested=True) as child_run:
                    return child_run.info.run_id, True
        latest_child_run_id = list(runs_df['run_id'])[0]
        latest_run_info = runs_df.iloc[:1]
        retrain = False
        # day_check_flag = self.check_existing_model_retrain(latest_child_run_id, latest_run_info, retrain)
        # if day_check_flag:
        if self.model is not None:
            with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
                with mlflow.start_run(experiment_id=experiment_id, nested=True) as child_run:
                    return child_run.info.run_id, True
        return latest_child_run_id, retrain

    @staticmethod
    def load_model_pyfunc(model_path):
        """
        Function to load the model from mlflow artifact path
        :param model_path: model path on mlflow
        :return: loaded model
        """
        try:
            model = mlflow.pyfunc.load_model(model_path)
            logger.info("loading the model")
            return model
        except Exception as e:
            logger.exception(str(e))

    def check_existing_model_retrain(self, latest_child_run_id, child_run_info, retrain):
        """
        If retrain is True, it returns true as retraining is required.
        If retrain is False, it checks the time difference between the last child run and the current time and returns
        true or false depending on the time difference
        :param latest_child_run_id: last child run id
        :param child_run_info: last child run info
        :param retrain: retrain flag
        :return: final retrain flag
        """
        if retrain:
            logger.info("Retraining Needed...")
            return True
        else:
            # uncomment this if auto train should happen
            logger.info(f"Already trained model is present, checking the age of the existing model of run id "
                        f"{latest_child_run_id}")
            # time_diff = self._mfu_.get_last_run_time_diff(child_run_info, self.retrain_param_unit)
            # logger.info(f"Time difference is {time_diff} {self.retrain_param_unit}")
            # if time_diff >= self.retrain_param_value:
            #     logger.info(f"Retraining needed as the last trained model time exceeds the mentioned time difference "
            #                 f"{self.retrain_param_value} {self.retrain_param_unit}")
            #     return True
            # else:
            #     logger.info(f"Retraining not needed as the last trained model time doesnt exceeds the mentioned time "
            #                 f"difference {self.retrain_param_value} {self.retrain_param_unit}")
            #     return False
            return True

    def forming_loading_path(self, latest_run_id):
        """
        Creates the path from the child run id
        :param latest_run_id: latest child run id
        :return: the path to the model
        """
        try:
            model_name = self.model_save_name
            return f"runs:/{latest_run_id}/{model_name}"
        except Exception as e:
            logger.exception(f"Exception while forming loading path  - {e}")

    def model_trainer(self, experiment_id, parent_run_id, child_run_id):
        """
        Using the experiment id, parent run id and child run id, it will train the model
        :param experiment_id: experiment id
        :param parent_run_id: parent run id
        :param child_run_id: child run id
        :return: the final model
        """
        with mlflow.start_run(experiment_id=experiment_id, run_id=parent_run_id, nested=True):
            with mlflow.start_run(experiment_id=experiment_id, run_id=child_run_id, nested=True):
                # metrics, hyperparams, model_name, model = self._pycaret_obj_.perform_task()
                # logger.info(f"Model is {model_name}")
                logger.info(f"Metrics: {self.metrics}")
                logger.info(f"Hyperparams: {self.hyperparams}")
                self._mfu_.log_model(model=self.model, model_name=self.model_save_name)
                model_name = 'ExtraTreeRegressor'
                self._mfu_.log_metrics(metrics=self.metrics)
                self._mfu_.log_hyper_param(hyper_params=self.hyperparams)
                self._mfu_.set_tag(child_run_id=child_run_id, key="algorithm", value=model_name)
        return self.model

    def get_latest_model(self):
        """
        This is the Main function which will return the latest model
        :return:
        """
        experiment_id = self.check_create_experiment()
        parent_run_id = self.check_create_parent_run(experiment_id)
        child_run_id, retrain = self.check_create_child_run(experiment_id, parent_run_id)
        logger.info(f"Retrain flag is {retrain}")
        if self.model is not None:
            logger.info("Retraining needed")
            self.model_trainer(experiment_id, parent_run_id, child_run_id)
            logger.info("New model trained successfully")
        else:
            logger.info(f"No retraining needed. proceeding to load the last child run model {child_run_id}")
        logger.info(f"Loading the model from the child run id {child_run_id}")
        final_model = self.load_model_pyfunc(model_path=self.forming_loading_path(latest_run_id=child_run_id))
        return final_model
