from mlflow.utils.logging_utils import eprint

import click

from mlflow.tracking import _get_store
from mlflow.store.artifact_repo import ArtifactRepository
from mlflow.utils.proto_json_utils import message_to_json


@click.group("artifacts")
def commands():
    """
    Upload, list, and download artifacts from an MLflow artifact repository.

    To manage artifacts for a run associated with a tracking server, set the MLFLOW_TRACKING_URI
    environment variable to the URL of the desired server.
    """
    pass


@commands.command("log-artifact")
@click.option("--local-file", "-l", required=True,
              help="Local path to artifact to log")
@click.option("--run-id", "-r", required=True,
              help="Run ID into which we should log the artifact.")
@click.option("--artifact-path", "-a",
              help="If specified, we will log the artifact into this subdirectory of the " +
                   "run's artifact directory.")
def log_artifact(local_file, run_id, artifact_path):
    """
    Logs a local file as an artifact of a run, optionally within a run-specific
    artifact path. Run artifacts can be organized into directories, so you can
    place the artifact in a directory this way.
    """
    store = _get_store()
    artifact_uri = store.get_run(run_id).info.artifact_uri
    artifact_repo = ArtifactRepository.from_artifact_uri(artifact_uri, store)
    artifact_repo.log_artifact(local_file, artifact_path)
    eprint("Logged artifact from local file %s to artifact_path=%s" % (local_file, artifact_path))


@commands.command("log-artifacts")
@click.option("--local-dir", "-l", required=True,
              help="Directory of local artifacts to log")
@click.option("--run-id", "-r", required=True,
              help="Run ID into which we should log the artifact.")
@click.option("--artifact-path", "-a",
              help="If specified, we will log the artifact into this subdirectory of the " +
                   "run's artifact directory.")
def log_artifacts(local_dir, run_id, artifact_path):
    """
    Logs the files within a local directory as an artifact of a run, optionally
    within a run-specific artifact path. Run artifacts can be organized into
    directories, so you can place the artifact in a directory this way.
    """
    store = _get_store()
    artifact_uri = store.get_run(run_id).info.artifact_uri
    artifact_repo = ArtifactRepository.from_artifact_uri(artifact_uri, store)
    artifact_repo.log_artifacts(local_dir, artifact_path)
    eprint("Logged artifact from local dir %s to artifact_path=%s" % (local_dir, artifact_path))


@commands.command("list")
@click.option("--run-id", "-r", required=True,
              help="Run ID to be listed")
@click.option("--artifact-path", "-a",
              help="If specified, a path relative to the run's root directory to list.")
def list_artifacts(run_id, artifact_path):
    """
    Return all the artifacts directly under run's root artifact directory,
    or a sub-directory. The output is a JSON-formatted list.
    """
    artifact_path = artifact_path if artifact_path is not None else ""
    store = _get_store()
    artifact_uri = store.get_run(run_id).info.artifact_uri
    artifact_repo = ArtifactRepository.from_artifact_uri(artifact_uri, store)
    file_infos = artifact_repo.list_artifacts(artifact_path)
    print(_file_infos_to_json(file_infos))


def _file_infos_to_json(file_infos):
    json_list = [message_to_json(file_info.to_proto()) for file_info in file_infos]
    return "[" + ", ".join(json_list) + "]"


@commands.command("download")
@click.option("--run-id", "-r", required=True,
              help="Run ID from which to download")
@click.option("--artifact-path", "-a",
              help="If specified, a path relative to the run's root directory to download")
def download_artifacts(run_id, artifact_path):
    """
    Download an artifact file or directory to a local directory.
    The output is the name of the file or directory on the local disk.
    """
    artifact_path = artifact_path if artifact_path is not None else ""
    store = _get_store()
    artifact_uri = store.get_run(run_id).info.artifact_uri
    artifact_repo = ArtifactRepository.from_artifact_uri(artifact_uri, store)
    artifact_location = artifact_repo.download_artifacts(artifact_path)
    print(artifact_location)
