#!/usr/bin/env python3
"""Functions for managing artifacts."""

import argparse
import functools
import logging
import pathlib
import sys

import mlflow
import yaml

from .common import configure_logging
from .sqlite import MLflowSqlite
from .uri import URITransformer

LOGGER = logging.getLogger(__name__)


def load_map(path):
    """
    Load a YAML file that maps old artifact URIs to new ones.
    """
    path = pathlib.Path(path).resolve()
    with path.open("rb", encoding="utf-8") as handle:
        data = yaml.safe_load(handle)
    uri_map = {}
    for old, new in data.items():
        old = URITransformer(old)
        new = URITransformer(new)
        uri_map[old] = new
    return uri_map


class ArtifactURIMapper:
    """
    Map old local articate URIs to new ones.
    """

    def __init__(self, paths=None, map_path=None):
        """
        Args:
            paths:
                An optional iterable over mlruns-like directories in which to
                search for displaced artifacts.

            map_path:
                A path to a YAML file containing a dict that maps old artifact
                URIs to new ones.
        """
        self.paths = paths
        self.map_path = map_path

    @functools.cached_property
    def user_map(self):
        """
        A dict mapping old URIs to new ones.
        """
        if self.map_path:
            return load_map(self.map_path)
        return {}

    @functools.cached_property
    def user_mapped_paths(self):
        """
        The reverse-sorted list of mapped URIs if user_map is not None, else
        None. The first match when searching for relative subpaths will be the
        most specific due to the sort order.
        """
        return sorted((u.path for u in self.user_map), reverse=True)

    def get_new_artifact_uri(self, artifact_uri, run):
        """
        Get the new URI for the given artifact URI.
        """
        # If there is a direct mapping, use it.
        if artifact_uri in self.user_map:
            return self.user_map[artifact_uri]

        # If the URI is not a local file, skip it.
        if not artifact_uri.is_local_file:
            return None

        artifact_path = artifact_uri.path

        for old in self.user_mapped_paths:
            if old.is_local_file and artifact_path.is_relative_to(old):
                path = self.user_map[old] / artifact_path.relative_to(old)
                return artifact_uri.copy(path=path)

        parts = artifact_path.parts
        expected_parts = (run.info.experiment_id, run.info.run_id, "artifacts")
        if parts[-3:] != expected_parts:
            LOGGER.error(
                "Expected current artifact path to end with %s, not %s",
                expected_parts,
                parts[-3:],
            )
            return None

        subpath = pathlib.Path(*expected_parts)
        for p in self.paths:
            path = p / subpath
            if path.exists():
                return artifact_uri.copy(path=path)

        return None

    def map_artifact_uris(self):
        """
        Get a dict mapping old artifact URIs to new ones.
        """
        artifact_map = {}
        for run in mlflow.search_runs(
            experiment_ids=[e.experiment_id for e in mlflow.search_experiments()],
            output_format="list",
        ):
            artifact_uri = URITransformer(run.info.artifact_uri)
            new_uri = self.get_new_artifact_uri(artifact_uri, run)

            if new_uri is None:
                if artifact_uri.is_local_file and not artifact_uri.path.exists():
                    LOGGER.warning(
                        "Local artifact path could not be located: %s [run_id: %s]",
                        artifact_uri,
                        run.info.run_id,
                    )
            elif artifact_uri != new_uri:
                LOGGER.info(
                    "Artifact URI will be updated from %s to %s [run_id: %s]",
                    artifact_uri,
                    new_uri,
                    run.info.run_id,
                )
                artifact_map[run.info.run_id] = str(new_uri)

        return artifact_map


def script_fix_artifact_uris(args=None):
    """
    Attempt to fix broken artifact URIs in experiments and runs. This only works
    if the artifacts can be found in local mlruns directories or directories
    with equivalent structures, and when the MLflow backend is a local SQLite3
    database.
    """
    configure_logging()

    parser = argparse.ArgumentParser(description=script_fix_artifact_uris.__doc__)
    parser.add_argument(
        "paths",
        nargs="*",
        default=[],
        type=pathlib.Path,
        help="Paths to directories containingartifacts, such as an mlruns directory.",
    )
    parser.add_argument(
        "-m", "--map", help="A path to a YAML file that maps old paths to new paths."
    )
    parser.add_argument(
        "--db",
        "--database",
        type=pathlib.Path,
        default=MLflowSqlite.DEFAULT_PATH,
        help="Path to the local MLflow SQLite database. Default: %(default)s",
    )
    try:
        pargs = parser.parse_args(args=args)
        uri_map = ArtifactURIMapper(
            paths=pargs.paths, map_path=pargs.map
        ).map_artifact_uris()
        if uri_map:
            MLflowSqlite(path=pargs.db).update_artifact_uris(uri_map)
    except OSError as err:
        sys.exit(err)
