import functools
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import click
from rich.console import Console
from tqdm import tqdm

from grid.cli import rich_click
from grid.cli.observables.base import create_table
from grid.cli.utilities import is_experiment
from grid.openapi import V1ClusterType, V1GetClusterResponse
from grid.openapi.api import cluster_service_api
from grid.sdk import env
from grid.sdk.artifacts import download_artifacts, list_artifacts
from grid.sdk.client import create_swagger_client
from grid.sdk.experiments import Experiment
from grid.sdk.runs import Run


def with_regular_certs(f):
    """ helper decorator to download artifacts without any custom cert. We can override with USE_CUSTOM_CERTS_ALWAYS (in ../globals.py) """

    # keep track of custom certs
    ssl_custom_cert_envs = "REQUESTS_CA_BUNDLE", "SSL_CERT_FILE"
    custom_certs = dict()

    if not env.SSL_CA_CERT:
        for k in ssl_custom_cert_envs:
            v = os.environ.get(k, None)
            if v is not None:
                custom_certs[k] = v

    @functools.wraps(f)
    def wrapper(*args, **kwargs):

        try:
            for e in custom_certs.keys():
                os.environ.pop(e, None)
            return f(*args, **kwargs)
        finally:
            for e, v in custom_certs.items():
                os.environ[e] = v

    return wrapper


@rich_click.command()
@click.option(
    "--bucket_path",
    "--bucket_paths",
    "bucket_path",
    is_flag=True,
    default=False,
    help="Do not download. Show the bucket URL for the experiments instead. (BYOC only)",
)
@click.option(
    "--download_dir",
    type=click.Path(exists=False, file_okay=False, dir_okay=True),
    required=False,
    default="./grid_artifacts",
    help="Download directory that will host all artifact files."
)
@click.option(
    "-m",
    "--match_regex",
    type=str,
    default="",
    help="Only show artifacts that match this regex filter. Best if quoted."
)
@rich_click.argument(
    "runs_or_experiments", type=str, required=True, nargs=-1, help="The run or experiment to download artifacts for."
)
def artifacts(
    runs_or_experiments: List[str],
    bucket_path: bool = False,
    download_dir: Optional[str] = None,
    match_regex: str = ""
) -> None:
    """Downloads artifacts for a given run or experiments.

    This will download artifacts generated by the runs / experiments.
    Regex filtering is used to determine which artifacts to download.
    """
    experiments: List[Tuple[Experiment, Path]] = []
    for element in runs_or_experiments:
        if is_experiment(element):
            experiment = Experiment(name=element)
            if not experiment.exists:
                click.echo(f"Experiment {element} does not exist - can not get artifacts")
                continue
            experiments.append((experiment, Path(download_dir) / element))

        else:
            run = Run(name=element)
            if not run.exists:
                click.echo(f"Run {element} does not exist - can not get artifacts")
                continue

            for experiment in run.experiments:
                if not experiment.exists:
                    click.echo(f"Experiment {experiment.name} does not exist - can not get artifacts")
                    continue
                experiments.append((experiment, Path(download_dir) / element / experiment.name))

    if len(experiments) == 0:
        click.echo("Aborting download - there are no valid experiments")
        return

    if bucket_path:
        # only valid for BYOC clusters
        api_client = create_swagger_client()
        cluster_api = cluster_service_api.ClusterServiceApi(api_client)
        clusters: Dict[str, V1GetClusterResponse] = {exp.cluster_id: None for exp, _ in experiments}
        for cluster_id in clusters.keys():
            resp = cluster_api.cluster_service_get_cluster(cluster_id)
            if resp.spec.cluster_type != V1ClusterType.BYOC:
                raise click.ClickException(
                    "--bucket_path is only valid for experiments on BYOC clusters. "
                    f"The experiment(s) requested belong to cluster '{cluster_id}' which is not BYOC."
                )
            clusters[cluster_id] = resp

        # print the s3 uri for the experiment artifacts
        table = create_table(['experiment name', 'bucket path to artifacts'])
        for exp, _ in experiments:
            driver = clusters[exp.cluster_id].spec.driver.kubernetes
            if not driver.aws:
                raise click.ClickException("only aws clusters are supported")
            aws = driver.aws
            table.add_row(exp.name, f"s3://{aws.bucket_name}/experiments/{exp.id}/artifacts/")

        console = Console()
        console.print(table)
        return

    click.echo("Downloading artifacts. This command may take a while")
    # exp.download_artifacts generates two progress bars:
    # 1st: progress of all artifacts in an experiment
    # 2nd: download progress for a single artifact in MB/s
    # This will add a 3rd at the bottom: experiments finished downloading
    exp_iter = tqdm(experiments, unit="experiment", position=2)
    for experiment, target_path in exp_iter:
        exp_iter.set_description(experiment.name)
        artifacts_to_download = list_artifacts(experiment.name, artifact_regex=match_regex)
        download_artifacts(artifacts_to_download, destination=target_path)
    click.echo("Done downloading artifacts!")
