#!/usr/bin/env python3
"""
Merge experiments into a common directory.
"""

import argparse
import logging
import pathlib
import sys
import tempfile

import mlflow
from mlflow.client import MlflowClient

from .common import configure_logging

LOGGER = logging.getLogger(__name__)


def page_iter(func, **kwargs):
    """
    Wrapper function to transform various MLflow paged iterators into a
    continuous generator.

    Args:
        func:
            The function to call. It must accept a page_token keyword parameter.

        **kwargs:
            The keyword arguments to pass to the function.

    Returns:
        A generator over the items return by the function.
    """
    token = kwargs.get("token")
    while True:
        result = func(**kwargs)
        yield from result
        token = result.token
        if token is None:
            return


class Merger:
    """
    Merge experiments and runs from one backend into another. This uses the
    mlflow.client.MlflowClient class which is a wrapper around the MLflow REST
    API.

    * https://mlflow.org/docs/latest/api_reference/python_api/mlflow.client.html
    * https://mlflow.org/docs/latest/api_reference/rest-api.html
    """

    def __init__(self, src_kwargs, dst_kwargs):
        """
        Args:
            src_kwargs:
                Keyword arguments for MlflowClient for the source client.

            dst_kwargs:
                Keyword arguments for MlflowClient for the destination client.
        """
        self.src_client = MlflowClient(**src_kwargs)
        self.dst_client = MlflowClient(**dst_kwargs)
        self.exp_ids = {}

    def copy_experiments(self):
        """
        Copy experiments.
        """
        for exp in page_iter(self.src_client.search_experiments):
            existing_exp = self.dst_client.get_experiment_by_name(exp.name)
            # Upload the experiment if it is missing.
            if existing_exp is None:
                LOGGER.info("Copying experiment: %s", exp.name)
                dst_id = self.dst_client.create_experiment(
                    name=exp.name,
                    artifact_location=exp.artifact_location,
                    tags=exp.tags,
                )
            else:
                LOGGER.info("Found experiment: %s", exp.name)
                dst_id = existing_exp.experiment_id
            self.exp_ids[exp.experiment_id] = dst_id

    def copy_datasets(self):
        """
        Copy datasets.
        """
        raise NotImplementedError("copy_datasets")

        if not hasattr(self.dst_client, "create_dataset"):
            LOGGER.warning(
                "The create_dataset method has been removed: datasets will not be copied"
            )
            return

        exp_ids = list(self.exp_ids)
        for dataset in page_iter(
            self.src_client.search_datasets, experiment_ids=exp_ids
        ):
            # TODO
            # Check if the name is the same as the dataset_id.
            existing_dataset = self.dst_client.get_dataset(dataset.name)
            if existing_dataset is None:
                new_dataset = self.dst_client.create_dataset(
                    name=dataset.name,
                    experiment_id=[self.exp_ids[e] for e in dataset.experiment_ids],
                    tags=dataset.tags,
                )
                # TODO
                # Transfer the data itself.
            # TODO
            # Create a dataset map if the IDs can differ.

    def copy_runs(self):
        """
        Copy runs.
        """
        exp_ids = list(self.exp_ids)
        for run in page_iter(self.src_client.search_runs, experiment_ids=exp_ids):
            try:
                self.dst_client.get_run(run.info.run_id)
                LOGGER.info("Found run: %s", run.info.run_id)
                continue
            # An exception occurs if the run does not exist.
            except mlflow.exceptions.MlflowException as err:
                # Check if the error includes the run ID. If not, something else
                # went wrong so reraise the exception.
                if run.info.run_id not in str(err):
                    raise err from err
            LOGGER.info("Copying run: %s", run.info.run_id)
            new_run = self.dst_client.create_run(
                experiment_id=self.exp_ids[run.info.experiment_id],
                start_time=run.info.start_time,
                tags=run.data.tags,
                run_name=run.info.run_name,
            )
            metrics = []
            for metric_name in run.data.metrics:
                metrics.extend(
                    self.src_client.get_metric_history(run.info.run_id, metric_name)
                )
            self.dst_client.log_batch(
                run_id=new_run.info.run_id,
                metrics=metrics,
                params=(
                    mlflow.entities.Param(k, v) for (k, v) in run.data.params.items()
                ),
            )
            self.copy_run_artifacts(run, new_run)

    def copy_run_artifacts(self, src_run, dst_run):
        """
        Copy run artifacts.
        """
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_dir = pathlib.Path(tmp_dir)
            for artifact in self.src_client.list_artifacts(src_run.info.run_id):
                LOGGER.info(
                    "Copying artifact: %s [run_id: %s]",
                    artifact.path,
                    src_run.info.run_id,
                )
                tmp_path = tmp_dir / artifact.path
                self.src_client.download_artifacts(
                    run_id=src_run.info.run_id, path=artifact.path, dst_path=tmp_path
                )
                self.dst_client.log_artifact(dst_run.info.run_id, tmp_path)

    # https://mlflow.org/docs/latest/ml/model-registry
    def copy_models(self):
        """
        Copy logged models.
        """
        exp_ids = list(self.exp_ids)
        dst_exp_ids = list(self.exp_ids.values())
        for model in page_iter(
            self.src_client.search_logged_models, experiment_ids=exp_ids
        ):
            existing_model = self.dst_client.search_logged_models(
                experiment_ids=dst_exp_ids, filter_string=f"name = {model.name}"
            )
            if existing_model:
                LOGGER.info("Found model: %s", model.name)
                continue
            new_model = self.dst_client.create_logged_model(
                experiment_id=self.exp_ids[model.experiment_id],
                name=model.name,
                source_run_id=model.source_run_id,
                tags=model.tags,
                params=model.params,
                model_type=model.type,
            )
            self.copy_model_artifacts(model, new_model)
            # create_model_version
            # set_registered_model_alias

    def copy_model_artifacts(self, src_model, dst_model):
        """
        Copy model artifacts.
        """
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_dir = pathlib.Path(tmp_dir)
            for artifact in self.src_client.list_logged_model_artifacts(
                src_model.model_id
            ):
                LOGGER.info(
                    "Copying artifact: %s [model name: %s]",
                    artifact.path,
                    src_model.name,
                )
                tmp_path = tmp_dir / artifact.path
                self.src_client.download_artifacts(
                    run_id=src_model.source_run_id,
                    path=artifact.path,
                    dst_path=tmp_path,
                )
                self.dst_client.log_model_artifact(dst_model.model_id, tmp_path)

    def copy_all(self):
        """
        Copy everything that is currently supported.
        """
        self.copy_experiments()
        self.copy_runs()
        self.copy_models()


def script_merge_backend(args=None):
    """
    Copy experiments and runs from one backend to another. This is a work in
    progress and should be used with caution.
    """
    configure_logging()

    parser = argparse.ArgumentParser(description=script_merge_backend.__doc__)
    parser.add_argument(
        "--tua",
        help="Tracking URI for the source backend.",
    )
    parser.add_argument(
        "--rua",
        help="Registry URI for the source backend.",
    )
    parser.add_argument(
        "--wua",
        help="Workspace store URI for the source backend.",
    )
    parser.add_argument(
        "--tub",
        help="Tracking URI for the destination backend.",
    )
    parser.add_argument(
        "--rub",
        help="Registry URI for the destination backend.",
    )
    parser.add_argument(
        "--wub",
        help="Workspace store URI for the destination backend.",
    )
    pargs = parser.parse_args(args=args)

    merger = Merger(
        src_kwargs={
            "tracking_uri": pargs.tua,
            "registry_uri": pargs.rua,
            "workspace_store_uri": pargs.wua,
        },
        dst_kwargs={
            "tracking_uri": pargs.tub,
            "registry_uri": pargs.rub,
            "workspace_store_uri": pargs.wub,
        },
    )
    try:
        merger.copy_all()
    except OSError as err:
        sys.exit(err)


if __name__ == "__main__":
    script_merge_backend()
