import json
import os

import click

import mlflow
from mlflow.entities import ViewType
from mlflow.exceptions import MlflowException
from mlflow.protos import databricks_pb2
from mlflow.tracking import _get_store, fluent
from mlflow.utils.data_utils import is_uri
from mlflow.utils.string_utils import _create_table

EXPERIMENT_ID = click.option("--experiment-id", "-x", type=click.STRING, required=True)


def _validate_max_results(ctx, param, value):
    """Validate that max_results is non-negative."""
    if value is not None and value < 0:
        raise click.BadParameter("max-results must be a non-negative integer")
    return value


@click.group("experiments")
def commands():
    """
    Manage experiments. To manage experiments associated with a tracking server, set the
    MLFLOW_TRACKING_URI environment variable to the URL of the desired server.
    """


@commands.command()
@click.option("--experiment-name", "-n", type=click.STRING, required=True)
@click.option(
    "--artifact-location",
    "-l",
    help="Base location for runs to store artifact results. Artifacts will be stored "
    "at $artifact_location/$run_id/artifacts. See "
    "https://mlflow.org/docs/latest/tracking.html#where-runs-are-recorded for "
    "more info on the properties of artifact location. "
    "If no location is provided, the tracking server will pick a default.",
)
def create(experiment_name, artifact_location):
    """
    Create an experiment.

    All artifacts generated by runs related to this experiment will be stored under artifact
    location, organized under specific run_id sub-directories.

    Implementation of experiment and metadata store is dependent on backend storage. ``FileStore``
    creates a folder for each experiment ID and stores metadata in ``meta.yaml``. Runs are stored
    as subfolders.
    """
    store = _get_store()
    exp_id = store.create_experiment(experiment_name, artifact_location)
    click.echo(f"Created experiment '{experiment_name}' with id {exp_id}")


@commands.command("search")
@click.option(
    "--view",
    "-v",
    default="active_only",
    help="Select view type for experiments. Valid view types are "
    "'active_only' (default), 'deleted_only', and 'all'.",
)
@click.option(
    "--max-results",
    type=click.INT,
    default=None,
    callback=_validate_max_results,
    help="Maximum number of experiments to return. If not provided, returns all experiments.",
)
def search_experiments(view, max_results):
    """
    Search for experiments in the configured tracking server.
    """
    view_type = ViewType.from_string(view) if view else ViewType.ACTIVE_ONLY
    experiments = mlflow.search_experiments(view_type=view_type, max_results=max_results)
    table = [
        [
            exp.experiment_id,
            exp.name,
            exp.artifact_location
            if is_uri(exp.artifact_location)
            else os.path.abspath(exp.artifact_location),
        ]
        for exp in experiments
    ]
    click.echo(_create_table(sorted(table), headers=["Experiment Id", "Name", "Artifact Location"]))


@commands.command("get")
@click.option(
    "--experiment-id",
    "-x",
    type=click.STRING,
    help="ID of the experiment to retrieve.",
)
@click.option(
    "--experiment-name",
    "-n",
    type=click.STRING,
    help="Name of the experiment to retrieve.",
)
@click.option(
    "--output",
    type=click.Choice(["json", "table"]),
    default="table",
    help="Output format: 'table' (default) or 'json'.",
)
def get_experiment(experiment_id, experiment_name, output):
    """
    Get details of an experiment by ID or name.

    Displays experiment information including name, artifact location, lifecycle stage,
    tags, creation time, and last update time.

    \b
    Examples:

    .. code-block:: bash

        # Get experiment by ID in table format (default)
        mlflow experiments get --experiment-id 1

        # Get experiment by name
        mlflow experiments get --experiment-name "My Experiment"

        # Get experiment in JSON format
        mlflow experiments get --experiment-name "My Experiment" --output json

        # Using short options
        mlflow experiments get -x 0
        mlflow experiments get -n "Default"
    """
    # Validate mutual exclusivity
    if (experiment_id is not None and experiment_name is not None) or (
        experiment_id is None and experiment_name is None
    ):
        raise click.UsageError("Must specify exactly one of --experiment-id or --experiment-name.")

    store = _get_store()

    # Retrieve experiment by ID or name
    if experiment_id is not None:
        experiment = store.get_experiment(experiment_id)
    else:
        experiment = store.get_experiment_by_name(experiment_name)
        if experiment is None:
            raise MlflowException(
                f"Experiment with name '{experiment_name}' does not exist.",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )

    if output == "json":
        experiment_dict = dict(experiment)
        click.echo(json.dumps(experiment_dict, indent=2))
    elif output == "table":
        table_data = [
            ["Experiment ID", experiment.experiment_id],
            ["Name", experiment.name],
            ["Artifact Location", experiment.artifact_location],
            ["Lifecycle Stage", experiment.lifecycle_stage],
            ["Creation Time", experiment.creation_time or "N/A"],
            ["Last Update Time", experiment.last_update_time or "N/A"],
        ]

        if experiment.tags:
            tags_str = ", ".join([f"{k}={v}" for k, v in experiment.tags.items()])
            table_data.append(["Tags", tags_str])
        else:
            table_data.append(["Tags", ""])

        max_field_width = max(len(row[0]) for row in table_data)
        for field, value in table_data:
            click.echo(f"{field.ljust(max_field_width + 2)}: {value}")


@commands.command("delete")
@EXPERIMENT_ID
def delete_experiment(experiment_id):
    """
    Mark an active experiment for deletion. This also applies to experiment's metadata, runs and
    associated data, and artifacts if they are store in default location. Use ``list`` command to
    view artifact location. Command will throw an error if experiment is not found or already
    marked for deletion.

    Experiments marked for deletion can be restored using ``restore`` command, unless they are
    permanently deleted.

    Specific implementation of deletion is dependent on backend stores. ``FileStore`` moves
    experiments marked for deletion under a ``.trash`` folder under the main folder used to
    instantiate ``FileStore``. Experiments marked for deletion can be permanently deleted by
    clearing the ``.trash`` folder. It is recommended to use a ``cron`` job or an alternate
    workflow mechanism to clear ``.trash`` folder.
    """
    store = _get_store()
    store.delete_experiment(experiment_id)
    click.echo(f"Experiment with ID {experiment_id} has been deleted.")


@commands.command("restore")
@EXPERIMENT_ID
def restore_experiment(experiment_id):
    """
    Restore a deleted experiment. This also applies to experiment's metadata, runs and associated
    data. The command throws an error if the experiment is already active, cannot be found, or
    permanently deleted.
    """
    store = _get_store()
    store.restore_experiment(experiment_id)
    click.echo(f"Experiment with id {experiment_id} has been restored.")


@commands.command("rename")
@EXPERIMENT_ID
@click.option("--new-name", type=click.STRING, required=True)
def rename_experiment(experiment_id, new_name):
    """
    Renames an active experiment.
    Returns an error if the experiment is inactive.
    """
    store = _get_store()
    store.rename_experiment(experiment_id, new_name)
    click.echo(f"Experiment with id {experiment_id} has been renamed to '{new_name}'.")


@commands.command("csv")
@EXPERIMENT_ID
@click.option("--filename", "-o", type=click.STRING)
def generate_csv_with_runs(experiment_id, filename):
    # type: (str, str) -> None
    """
    Generate CSV with all runs for an experiment
    """
    runs = fluent.search_runs(experiment_ids=experiment_id)
    if filename:
        runs.to_csv(filename, index=False)
        click.echo(
            f"Experiment with ID {experiment_id} has been exported as a CSV to file: {filename}."
        )
    else:
        click.echo(runs.to_csv(index=False))
