import os
import sys
import sky
import time
import inspect
import traceback

from typing import Union as TUnion, get_args, get_origin, List, Dict, Tuple
from metaflow._vendor import click
from metaflow import util
from metaflow import R
from metaflow.parameters import JSONType
from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY
from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
from metaflow.mflog import TASK_LOG_SOURCE

from .skypilot_exceptions import SkypilotKilledException
from .skypilot import Skypilot


class ClickUnion(click.ParamType):
    name = "union"

    def __init__(self, types):
        self.types = types

    def convert(self, value, param, ctx):
        for t in self.types:
            try:
                # Use Click's internal conversion logic for each type
                return t.convert(value, param, ctx)
            except (click.BadParameter, click.UsageError, ValueError, TypeError):
                continue
        self.fail(f"{value} does not match any of the accepted types: {self.types}", param, ctx)


def get_click_type(annotation):
    mapping = {
        int: click.INT,
        float: click.FLOAT,
        bool: click.BOOL,
        str: click.STRING,
    }

    origin = get_origin(annotation)
    
    # Handle complex generics (List, Dict, Tuple)
    if origin in (list, List, dict, Dict, tuple, Tuple):
        return JSONType

    # Handle Unions and Optionals
    if origin is TUnion:
        args = [a for a in get_args(annotation) if a is not type(None)]

        if len(args) == 1:
            return get_click_type(args[0])

        # Recursively resolve types within the Union
        click_types = [get_click_type(a) for a in args]
        return ClickUnion(click_types)

    # Fallback to primitives or string
    return mapping.get(annotation, click.STRING)


def add_sky_options(func):
    try:
        sig = inspect.signature(sky.Resources.__init__)
        for param in sig.parameters.values():
            if param.name in ["self", "cloud"] or param.name.startswith('_'):
                continue

            # Convert underscores to dashes for CLI consistency
            opt_name = param.name.replace('_', '-')
            click_type = get_click_type(param.annotation)

            extra = {}
            if click_type is click.BOOL:
                extra['is_flag'] = False

            func = click.option(
                f"--{opt_name}",
                default=param.default,
                type=click_type,
                help=f"SkyPilot {param.name} resource.",
                **extra,
            )(func)
    except (ImportError, NameError, ModuleNotFoundError):
        pass

    return func


@click.group()
def cli():
    pass


@cli.group(help="Commands related to Skypilot.")
def skypilot():
    pass


@skypilot.command(
    help="Execute a single task using Skypilot. This command is used internally by Metaflow."
)
@click.argument("step-name")
@click.argument("code-package-sha")
@click.argument("code-package-url")
@click.option("--executable", help="Executable requirement for Skypilot.")
# Dynamically inject SkyPilot-specific resource options
@add_sky_options
@click.option("--cluster-name", default=None, help="Name of an existing SkyPilot cluster to reuse.")
# Standard Metaflow options
@click.option("--run-id", help="Passed to the top-level 'step'.")
@click.option("--task-id", help="Passed to the top-level 'step'.")
@click.option("--input-paths", help="Passed to the top-level 'step'.")
@click.option("--split-index", help="Passed to the top-level 'step'.")
@click.option("--clone-path", help="Passed to the top-level 'step'.")
@click.option("--clone-run-id", help="Passed to the top-level 'step'.")
@click.option("--tag", multiple=True, default=None, help="Passed to the top-level 'step'.")
@click.option("--namespace", default=None, help="Passed to the top-level 'step'.")
@click.option("--retry-count", default=0, help="Passed to the top-level 'step'.")
@click.option("--max-user-code-retries", default=0, help="Passed to the top-level 'step'.")
@click.option(
    "--run-time-limit",
    default=5 * 24 * 60 * 60,  # Default is set to 5 days
    help="Run time limit in seconds for Skypilot job.",
)
@click.pass_context
def step(
    ctx,
    step_name,
    code_package_sha,
    code_package_url,
    executable=None,
    run_time_limit=None,
    **kwargs,
):
    def echo(msg, stream="stderr", job_id=None, **kwargs):
        msg = util.to_unicode(msg)
        if job_id:
            msg = "[%s] %s" % (job_id, msg)
        ctx.obj.echo_always(msg, err=(stream == sys.stderr), **kwargs)

    sky_resource_keys = set(inspect.signature(sky.Resources.__init__).parameters.keys()) | {"cluster_name"}
    mf_step_kwargs = {k: v for k, v in kwargs.items() if k not in sky_resource_keys}

    if R.use_r():
        entrypoint = R.entrypoint()
    else:
        executable = ctx.obj.environment.executable(step_name, executable)
        entrypoint = "%s -u %s" % (executable, os.path.basename(sys.argv[0]))

    top_args = " ".join(util.dict_to_cli_options(ctx.parent.parent.params))

    input_paths = kwargs.get("input_paths")
    split_vars = None
    if input_paths:
        max_size = 30 * 1024
        split_vars = {
            "METAFLOW_INPUT_PATHS_%d" % (i // max_size): input_paths[i : i + max_size]
            for i in range(0, len(input_paths), max_size)
        }
        kwargs["input_paths"] = "".join("${%s}" % s for s in split_vars.keys())

    step_args = " ".join(util.dict_to_cli_options(mf_step_kwargs))

    step_cli = "{entrypoint} {top_args} step {step} {step_args}".format(
        entrypoint=entrypoint,
        top_args=top_args,
        step=step_name,
        step_args=step_args,
    )
    node = ctx.obj.graph[step_name]

    # Get retry information
    retry_count = kwargs.get("retry_count", 0)
    retry_deco = [deco for deco in node.decorators if deco.name == "retry"]
    minutes_between_retries = None
    if retry_deco:
        minutes_between_retries = int(
            retry_deco[0].attributes.get("minutes_between_retries", 1)
        )
    if retry_count:
        ctx.obj.echo_always(
            "Sleeping %d minutes before the next retry" % minutes_between_retries
        )
        time.sleep(minutes_between_retries * 60)

    # Set batch attributes
    task_spec = {
        "flow_name": ctx.obj.flow.name,
        "step_name": step_name,
        "run_id": kwargs["run_id"],
        "task_id": kwargs["task_id"],
        "retry_count": str(retry_count),
    }
    attrs = {"metaflow.%s" % k: v for k, v in task_spec.items()}
    attrs["metaflow.user"] = util.get_username()
    attrs["metaflow.version"] = ctx.obj.environment.get_environment_info()[
        "metaflow_version"
    ]

    # Set environment
    env = {}
    env_deco = [deco for deco in node.decorators if deco.name == "environment"]
    if env_deco:
        env = env_deco[0].attributes["vars"]

    # Add the environment variables related to the input-paths argument
    if split_vars:
        env.update(split_vars)

    # Set log tailing.
    ds = ctx.obj.flow_datastore.get_task_datastore(
        mode="w",
        run_id=kwargs["run_id"],
        step_name=step_name,
        task_id=kwargs["task_id"],
        attempt=int(retry_count),
    )
    stdout_location = ds.get_log_location(TASK_LOG_SOURCE, "stdout")
    stderr_location = ds.get_log_location(TASK_LOG_SOURCE, "stderr")

    def _sync_metadata():
        if ctx.obj.metadata.TYPE == "local":
            sync_local_metadata_from_datastore(
                DATASTORE_LOCAL_DIR,
                ctx.obj.flow_datastore.get_task_datastore(
                    kwargs["run_id"], step_name, kwargs["task_id"]
                ),
            )

    try:
        skypilot = Skypilot(
            datastore=ctx.obj.flow_datastore,
            metadata=ctx.obj.metadata,
            environment=ctx.obj.environment,
        )

        with ctx.obj.monitor.measure("metaflow.skypilot.launch_job"):
            skypilot.launch_job(
                step_name=step_name,
                step_cli=step_cli,
                task_spec=task_spec,
                code_package_sha=code_package_sha,
                code_package_url=code_package_url,
                code_package_ds=ctx.obj.flow_datastore.TYPE,
                run_time_limit=run_time_limit,
                env=env,
                attrs=attrs,
                **kwargs,
            )
    except Exception:
        traceback.print_exc(chain=False)
        _sync_metadata()
        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)

    try:
        skypilot.wait(stdout_location, stderr_location, echo=echo)
    except SkypilotKilledException:
        # don't retry killed tasks
        traceback.print_exc()
        sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)
    finally:
        _sync_metadata()
