import os
import sys
import json
import inspect

from metaflow import R
from metaflow import StepMutator
from metaflow.sidecar import Sidecar
from metaflow.decorators import StepDecorator
from metaflow.exception import MetaflowException
from metaflow.metadata_provider import MetaDatum
from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
from metaflow.plugins.timeout_decorator import get_run_time_limit_for_task
from metaflow.metadata_provider.util import sync_local_metadata_to_datastore

from .skypilot_exceptions import SkypilotException


SKYPILOT_DEPS = {"skypilot": "0.11.1"}


def _find_pypi_decospec(step):
    for deco in step.decorator_specs:
        if deco[0] == "pypi":
            return deco
    return None


def _merge_pypi_into_step(step):
    existing = _find_pypi_decospec(step)

    if existing:
        _, _, _, attrs = existing
        merged_attrs = dict(attrs)
        merged_attrs["packages"] = {**SKYPILOT_DEPS, **attrs.get("packages", {})}
        step.remove_decorator("pypi")
        step.add_decorator("pypi", deco_kwargs=merged_attrs)
    else:
        step.add_decorator("pypi", deco_kwargs={"packages": SKYPILOT_DEPS})


class skypilot(StepMutator):
    def init(self, *args, **kwargs):
        if args:
            raise SkypilotException(
                "@skypilot only accepts keyword arguments, e.g. @skypilot(accelerators='A100')"
            )
        self.kwargs = kwargs

    def mutate(self, mutable_step):
        _merge_pypi_into_step(mutable_step)
        mutable_step.add_decorator(
            "skypilot_step",
            deco_kwargs=self.kwargs,
            duplicates=mutable_step.ERROR,
        )


class SkypilotDecorator(StepDecorator):
    """
    Specifies that this step should execute with Skypilot.
    """

    name = "skypilot_step"

    @classmethod
    def _get_sky_defaults(cls):
        try:
            import sky
            sig = inspect.signature(sky.Resources.__init__)
            return {
                param.name: param.default for param in sig.parameters.values()
                if param.name not in ["self", "cloud"] and not param.name.startswith('_')
            }
        except (ImportError, ModuleNotFoundError, NameError):
            return {}

    # Initialize defaults using the safe class method, plus skypilot launch options
    defaults = {**_get_sky_defaults.__func__(None), "cluster_name": None}

    package_url = None
    package_sha = None
    run_time_limit = None

    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
        super(SkypilotDecorator, self).__init__(attributes, statically_defined, inserted_by)

    # Refer https://github.com/Netflix/metaflow/blob/master/docs/lifecycle.png
    # to understand where these functions are invoked in the lifecycle of a
    # Metaflow flow.
    def step_init(self, flow, graph, step, decos, environment, flow_datastore, logger):
        # Set internal state.
        self.logger = logger
        self.environment = environment
        self.step = step
        self.flow_datastore = flow_datastore

        # Set run time limit for the Skypilot Batch job.
        self.run_time_limit = get_run_time_limit_for_task(decos)
        if self.run_time_limit < 60:
            raise SkypilotException(
                "The timeout for step *{step}* should be at "
                "least 60 seconds for execution on Skypilot.".format(step=step)
            )

        if any([deco.name in ["kubernetes", "batch", "slurm"] for deco in decos]):
            raise MetaflowException(
                "Step *{step}* contains @skypilot alongside another compute decorator. "
                "Please use only one.".format(step=step)
            )

        if any([deco.name == "parallel" for deco in decos]):
            raise MetaflowException(
                "Step *{step}* contains a @parallel decorator "
                "with the @skypilot decorator. @parallel is not supported with @skypilot.".format(
                    step=step
                )
            )

    def package_init(self, flow, step_name, environment):
        try:
            # sky is a soft dependency.
            import sky
        except (NameError, ImportError, ModuleNotFoundError):
            raise SkypilotException(
                "Could not import module 'sky'.\n\nInstall skypilot "
                "Python package (https://pypi.org/project/skypilot/) first.\n"
                "You can install the module by executing - "
                "%s -m pip install skypilot\n"
                "or equivalent through your favorite Python package manager."
                % sys.executable
            )

    def runtime_init(self, flow, graph, package, run_id):
        # Set some more internal state.
        self.flow = flow
        self.graph = graph
        self.package = package
        self.run_id = run_id

    def runtime_task_created(
        self, task_datastore, task_id, split_index, input_paths, is_cloned, ubf_context
    ):
        if not is_cloned:
            self._save_package_once(self.flow_datastore, self.package)

    def runtime_step_cli(
        self, cli_args, retry_count, max_user_code_retries, ubf_context
    ):
        if retry_count <= max_user_code_retries:
            cli_args.commands = ["skypilot", "step"]
            cli_args.command_args.append(self.package_sha)
            cli_args.command_args.append(self.package_url)

            # Filter and pass only non-None attributes to keep the CLI clean.
            options = {
                k.replace('_', '-'): (
                    str(v).lower() if isinstance(v, bool)
                    else json.dumps(v) if isinstance(v, (dict, list, tuple))
                    else v
                )
                for k, v in self.attributes.items() if v is not None
            }
            cli_args.command_options.update(options)

            cli_args.command_options["run-time-limit"] = self.run_time_limit
            if not R.use_r():
                cli_args.entrypoint[0] = sys.executable

    def task_pre_step(
        self,
        step_name,
        task_datastore,
        metadata,
        run_id,
        task_id,
        flow,
        graph,
        retry_count,
        max_retries,
        ubf_context,
        inputs,
    ):
        self.metadata = metadata
        self.task_datastore = task_datastore

        meta = {}
        if "METAFLOW_SKYPILOT_WORKLOAD" in os.environ:
            meta["skypilot-task-id"] = os.environ.get("SKYPILOT_TASK_ID")
            skypilot_cluster_info = json.loads(os.environ.get("SKYPILOT_CLUSTER_INFO"))
            for k, v in skypilot_cluster_info.items():
                if v is not None:
                    meta[f"skypilot-{k}"] = v
            meta["skypilot-job-id"] = os.environ.get("SKYPILOT_INTERNAL_JOB_ID")

            self._save_logs_sidecar = Sidecar("save_logs_periodically")
            self._save_logs_sidecar.start()
        
        if len(meta) > 0:
            entries = [
                MetaDatum(
                    field=k,
                    value=v,
                    type=k,
                    tags=["attempt_id:{0}".format(retry_count)],
                )
                for k, v in meta.items()
                if v is not None
            ]
            # Register book-keeping metadata for debugging.
            metadata.register_metadata(run_id, step_name, task_id, entries)

    def task_finished(
        self, step_name, flow, graph, is_task_ok, retry_count, max_retries
    ):
        if "METAFLOW_SKYPILOT_WORKLOAD" in os.environ:
            if hasattr(self, "metadata") and self.metadata.TYPE == "local":
                # Note that the datastore is *always* Amazon S3 (see
                # runtime_task_created function).
                sync_local_metadata_to_datastore(
                    DATASTORE_LOCAL_DIR, self.task_datastore
                )

        try:
            self._save_logs_sidecar.terminate()
        except:
            # Best effort kill
            pass

    @classmethod
    def _save_package_once(cls, flow_datastore, package):
        if cls.package_url is None:
            cls.package_url, cls.package_sha = flow_datastore.save_data(
                [package.blob], len_hint=1
            )[0]
