"""General Job behavior."""
import json
import os
import warnings
from shutil import SameFileError, copy2, copytree
from typing import Optional

import ML_management
from ML_management import variables
from ML_management.collectors import COLLECTORS
from ML_management.collectors.collector_pattern_to_methods_map import CollectorMethodName
from ML_management.dataset_loader.dataset_loader_pattern_to_methods_map import DatasetLoaderMethodName
from ML_management.executor.executor_pattern_to_methods_map import ExecutorMethodName
from ML_management.executor.upload_model_mode import UploadModelMode
from ML_management.graphql.schema import ExecutionJob
from ML_management.job.params_prepare import method_params_prepare
from ML_management.mlmanagement.load_api import load_dataset, load_executor, load_model
from ML_management.mlmanagement.visibility_options import VisibilityOptions
from ML_management.model.model_type_to_methods_map import ModelMethodName


class Job:
    """Class of job."""

    def __init__(
        self,
        logger,
        preload_models: Optional[dict] = None,
        preload_dataset_loaders: Optional[dict] = None,
        preload_executor=None,
    ):
        self.logger = logger
        self.role_model_map = {}
        self.role_dataset_loader_map = {}
        self.executor = None
        if preload_models:
            self.role_model_map = preload_models

        if preload_dataset_loaders:
            self.role_dataset_loader_map = preload_dataset_loaders
        if preload_executor:
            self.executor = preload_executor

    def unset_dataset(self) -> None:
        """Unset dataset for model and executor instance before logging."""
        self.executor.role_dataset_loader_map = {}
        for role in self.role_model_map:
            self.role_model_map[role].dataset = None

    def set_device(self, device) -> None:
        """Set device for executor instance variable."""
        self.executor.device = device

    def set_model_log_args(
        self, executor_name: str, executor_version: int, new_model_names, source_models_name_version
    ):
        """Set additional arguments for log_model() to model instance."""
        for role, name_version in source_models_name_version.items():
            self.role_model_map[role].source_model_name = name_version["name"]
            self.role_model_map[role].source_model_version = name_version[
                "version"
            ]  # we can't pass int value as job parameters kwargs
            self.role_model_map[role].source_executor_name = executor_name
            self.role_model_map[role].source_executor_version = executor_version
            self.role_model_map[role].source_executor_role = role
            self.role_model_map[role].model_name = (
                name_version["name"] if role not in new_model_names else new_model_names[role]
            )

    @staticmethod
    def get_init_parameters_from_dict(params):
        """Get __init__ params from list of method_params."""
        for param in params:
            if param["method_name"] == ModelMethodName.init.value:
                return param["method_params"]
        return {}

    def model_preparing(self, list_role_model_method_params):
        self.model_params = {
            role_model.role: {
                "model_methods_params": [
                    {
                        "method_name": model_method_params.method_name,
                        "method_params": json.loads(model_method_params.method_params),
                    }
                    for model_method_params in role_model.model_params.list_model_method_params
                ]
            }
            for role_model in list_role_model_method_params
        }

        self.source_models_name_version = {
            role_param.role: {
                "name": role_param.model_params.model_version_choice.name,
                "version": role_param.model_params.model_version_choice.version,
            }
            for role_param in list_role_model_method_params
        }
        if not self.role_model_map:
            for role, model in self.source_models_name_version.items():
                self.logger.info(f"Start to load model name: {model['name']} version: {model['version']}")
                self.role_model_map[role] = load_model(
                    model["name"],
                    model["version"],
                    False,
                    f"model-{role}",
                    self.get_init_parameters_from_dict(self.model_params[role]["model_methods_params"]),
                ).loaded_object
                self.logger.info(f"The model name: {model['name']} version: {model['version']} was loaded successfully")

        self.new_model_names = {
            role_param.role: role_param.model_params.new_model_name for role_param in list_role_model_method_params
        }

    def data_preparing(self, list_role_data_params):
        self.dataset_loader_params = {
            data_param_role.role: {
                "dataset_loader_methods_params": [
                    {
                        "method_name": dataset_loader_method_params.method_name,
                        "method_params": json.loads(dataset_loader_method_params.method_params),
                    }
                    for dataset_loader_method_params in data_param_role.data_params.list_dataset_loader_method_params
                ]
            }
            for data_param_role in list_role_data_params
        }

        self.collector_params = {
            data_param_role.role: {
                "method_name": data_param_role.data_params.collector_method_params.method_name,
                "method_params": data_param_role.data_params.collector_method_params.method_params,
            }
            for data_param_role in list_role_data_params
        }

        self.collectors = {
            data_param_role.role: data_param_role.data_params.collector_name
            for data_param_role in list_role_data_params
        }

        # configure collector
        self.logger.info("Start to load collectors")
        self.role_collectors_map = {
            role: COLLECTORS[collector_name]() for role, collector_name in self.collectors.items()
        }
        self.collector_method_name_map = {
            role: collector_param["method_name"] for role, collector_param in self.collector_params.items()
        }
        self.collector_method_params_map = {
            role: json.loads(collector_param["method_params"])
            for role, collector_param in self.collector_params.items()
        }
        self.logger.info("Collectors was loaded successfully")

        # configure dataset loader
        dataset_loader_name_version = {
            role_param.role: {
                "name": role_param.data_params.dataset_loader_version_choice.name,
                "version": role_param.data_params.dataset_loader_version_choice.version,
            }
            for role_param in list_role_data_params
        }
        if not self.role_dataset_loader_map:
            for role, dataset_loader in dataset_loader_name_version.items():
                self.logger.info(
                    f"Start to load datasetloader name: {dataset_loader['name']} version: {dataset_loader['version']}"
                )
                self.role_dataset_loader_map[role] = load_dataset(
                    dataset_loader["name"],
                    int(dataset_loader["version"]),
                    False,
                    kwargs_for_init=self.get_init_parameters_from_dict(
                        self.dataset_loader_params[role]["dataset_loader_methods_params"]
                    ),
                ).loaded_object
                self.logger.info(
                    f"The datasetloader name: {dataset_loader['name']} version: {dataset_loader['version']} "
                    f"was loaded successfully"
                )

        self.dataset_loader_method_parameters_dict = {}
        for role, dl_params in self.dataset_loader_params.items():
            self.dataset_loader_method_parameters_dict[role] = {
                DatasetLoaderMethodName(p["method_name"]): method_params_prepare(
                    self.role_dataset_loader_map[role], p["method_name"], p["method_params"]
                )
                for p in dl_params["dataset_loader_methods_params"]
            }

        for role, collector in self.role_collectors_map.items():
            try:
                self.role_dataset_loader_map[role].data_path = getattr(
                    collector, CollectorMethodName(self.collector_method_name_map[role]).name
                )(**self.collector_method_params_map[role])
            except Exception as err:
                self.logger.info(f"Error during data collection of {role} collector.")
                raise err

    def run_job(self, job: ExecutionJob):
        """General protocol."""
        experiment_name = job.experiment.name

        params = job.params
        executor_name = params.executor_params.executor_version_choice.name
        executor_version = params.executor_params.executor_version_choice.version
        list_role_model_method_params = params.list_role_model_params
        list_role_data_params = params.list_role_data_params

        self.model_preparing(list_role_model_method_params)

        self.logger.info(f"Start to load executor name: {executor_name} version: {executor_version}")
        if not self.executor:
            self.executor = load_executor(executor_name, executor_version, False).loaded_object
        self.logger.info("The executor was loaded successfully")

        self.set_model_log_args(executor_name, executor_version, self.new_model_names, self.source_models_name_version)
        self.executor.role_model_map = self.role_model_map

        self.data_preparing(list_role_data_params)

        self.executor.dataset_loader_method_parameters_dict = self.dataset_loader_method_parameters_dict

        # configure executor
        executor_method_name = params.executor_params.executor_method_params.method_name
        executor_method_params = json.loads(params.executor_params.executor_method_params.method_params)

        self.executor.model_method_parameters_dict = {}
        for role, m_params in self.model_params.items():
            self.executor.model_method_parameters_dict[role] = {
                ModelMethodName(p["method_name"]): method_params_prepare(
                    self.role_model_map[role], p["method_name"], p["method_params"]
                )
                for p in m_params["model_methods_params"]
            }

        self.executor.role_dataset_loader_map = self.role_dataset_loader_map
        self.executor.role_dataset_map = {}
        for role, dataset_loader in self.role_dataset_loader_map.items():
            try:
                self.executor.role_dataset_map[role] = dataset_loader.get_dataset(
                    **self.dataset_loader_method_parameters_dict[role][DatasetLoaderMethodName.get_dataset]
                )
            except Exception as err:
                self.logger.info(f"Error during data loading by {role} dataset loader.")
                raise err
        self.set_device("cuda" if params.gpu else "cpu")

        self.logger.info("Start executor methods for job ...")

        try:
            executor_method_params = method_params_prepare(
                self.executor, ExecutorMethodName(executor_method_name).name, executor_method_params
            )
            artifacts = getattr(self.executor, ExecutorMethodName(executor_method_name).name)(**executor_method_params)
            self.unset_dataset()
            self.logger.info("Start uploading model")
            upload_model_modes = {
                role: UploadModelMode(mode) for role, mode in self.executor.upload_model_modes.items()
            }

            prepare_new_models_inference = {
                role_param.role: role_param.model_params.prepare_new_model_inference
                for role_param in list_role_model_method_params
            }
            descriptions = {
                role_param.role: role_param.model_params.description for role_param in list_role_model_method_params
            }
            start_build_new_models_image = {
                role_param.role: role_param.model_params.start_build_new_model_image
                for role_param in list_role_model_method_params
            }
            new_models_visibility = {
                role_param.role: role_param.model_params.new_model_visibility
                for role_param in list_role_model_method_params
            }

            for role, model in self.role_model_map.items():
                upload_model_mode = upload_model_modes[role]

                if upload_model_mode == UploadModelMode.none:
                    continue

                registered_model_name = (
                    self.new_model_names[role]
                    if upload_model_mode == UploadModelMode.new_model
                    else self.source_models_name_version[role]["name"]
                )

                if artifacts is not None:
                    src_artifacts = None
                    if len(self.role_model_map) == 1 and isinstance(artifacts, str):
                        src_artifacts = artifacts
                    elif isinstance(artifacts, dict) and role in artifacts:
                        src_artifacts = artifacts.get(role)
                    else:
                        warnings.warn(
                            "Artifacts returned from executor must be str for single model "
                            "or dict with model roles for multiple models. "
                            "Artifacts returned from executor for "
                            "model {registered_model_name} were not logged.".format(
                                registered_model_name=registered_model_name
                            )
                        )
                    if src_artifacts:
                        dst_artifacts = model.artifacts
                        if not os.path.exists(src_artifacts):
                            warnings.warn(
                                "Artifact's path returned from executor for "
                                "model {registered_model_name} does not exists.".format(
                                    registered_model_name=registered_model_name
                                )
                            )

                        elif os.path.isdir(src_artifacts):
                            copytree(src_artifacts, dst_artifacts)
                        else:
                            try:
                                copy2(src_artifacts, dst_artifacts)
                            except SameFileError:
                                pass

                model_path = os.path.dirname(model.artifacts)
                # child model needs to have same additional local dependencies as a parent
                # in current folder structure we know that each folder beside model folder
                # is folder with additional local dependencies
                additional_local_dependencies = os.listdir(os.path.dirname(model_path))
                additional_local_dependencies = [
                    os.path.join(os.path.dirname(model_path), folder) for folder in additional_local_dependencies
                ]
                additional_local_dependencies.remove(model_path)
                ML_management.mlmanagement.log_api._log_object_src(
                    artifact_path="",
                    description=descriptions[role],
                    registered_model_name=registered_model_name,
                    source_model_name=self.source_models_name_version[role]["name"],
                    source_model_version=int(self.source_models_name_version[role]["version"]),
                    source_executor_name=executor_name,
                    source_executor_version=int(executor_version),
                    source_executor_role=role,
                    upload_model_mode=upload_model_mode,
                    model_path=model_path,
                    create_venv_pack=prepare_new_models_inference[role],
                    additional_local_packages=additional_local_dependencies,
                    experiment_name=experiment_name,
                    start_build=start_build_new_models_image[role],
                    visibility=VisibilityOptions[new_models_visibility[role]],
                )
                self.logger.info(f"Model: {registered_model_name}, role: {role} logged successfully.")
        except Exception as err:
            self.logger.error(f"Error during execution of the job with obj_uuid {variables.secret_uuid}.")
            raise err
