"""Launch command for the Mithril CLI."""

from __future__ import annotations

from dataclasses import dataclass, field
from importlib.metadata import version as package_version
from typing import TYPE_CHECKING, Annotated, Any, cast

import click
from cyclopts import Parameter
from rich.console import Console

from mithril.api.client import MithrilClient
from mithril.cli.commands.onboarding import interactive_onboarding
from mithril.cli.utils.skypilot_passthrough import SKY_ALIAS_COMMANDS
from mithril.cli.utils.streaming import poll_launch, stream_job_logs
from mithril.cli.utils.volumes import ensure_volumes_registered
from mithril.sky import SkyClient

if TYPE_CHECKING:
    from sky import Resources as SkyResources
    from sky import Task as SkyTask

console = Console()
INDENT_SYMBOL = "├── "
INDENT_LAST_SYMBOL = "└── "


@dataclass(frozen=True, slots=True)
class ResourceParams:
    """Resource-related CLI options."""

    gpus: str | None = None
    cpus: str | None = None
    memory: str | None = None
    cloud: str | None = None
    region: str | None = None
    use_spot: bool | None = None


@dataclass(frozen=True, slots=True)
class TaskParams:
    """Task-related CLI options."""

    entrypoint: str | None
    name: str | None = None
    workdir: str | None = None
    num_nodes: int | None = None
    env: tuple[str, ...] = ()
    resources: ResourceParams = field(default_factory=ResourceParams)


def _exit(code: int) -> None:
    raise SystemExit(code)


def _is_bare_invocation(
    *,
    entrypoint: str | None,
    cluster_name: str | None,
    gpus: str | None,
    cpus: str | None,
    memory: str | None,
    cloud: str | None,
    region: str | None,
    use_spot: bool | None,
    num_nodes: int | None,
    name: str | None,
    workdir: str | None,
    env: tuple[str, ...],
    dryrun: bool,
) -> bool:
    """True when the user typed `ml launch` with no arguments or flags."""
    return (
        entrypoint is None
        and cluster_name is None
        and gpus is None
        and cpus is None
        and memory is None
        and cloud is None
        and region is None
        and use_spot is None
        and num_nodes is None
        and name is None
        and workdir is None
        and not env
        and not dryrun
    )


# `ml launch` is its own command (rather than only exposing `ml sky launch ...`).
# Today it's a thin proxy over SkyPilot: we translate flags into a `sky.Task` /
# `sky.Resources` and delegate the launch.
#
# Keeping `ml launch` first-class gives us control over the launch sequence (e.g. better
# logging/UX), and a stable place to add Mithril features that SkyPilot may choose not
# to support (or that we eventually implement without going through SkyPilot).
def launch(
    entrypoint: str | None = None,
    *,
    cluster_name: Annotated[
        str | None,
        Parameter(name=["-c", "--cluster"], help="Cluster name."),
    ] = None,
    gpus: Annotated[
        str | None,
        Parameter(name="--gpus", help="GPU type and count (e.g., 'A100:4')."),
    ] = None,
    cpus: Annotated[
        str | None,
        Parameter(name="--cpus", help="Number of vCPUs."),
    ] = None,
    memory: Annotated[
        str | None,
        Parameter(name="--memory", help="Memory in GB."),
    ] = None,
    cloud: Annotated[
        str | None,
        Parameter(name="--cloud", help="Cloud provider."),
    ] = None,
    region: Annotated[str | None, Parameter(name="--region", help="Region.")] = None,
    use_spot: Annotated[
        bool | None,
        Parameter(name="use-spot", negative="no-spot", help="Use spot instances."),
    ] = None,
    num_nodes: Annotated[
        int | None,
        Parameter(name="--num-nodes", help="Number of nodes."),
    ] = None,
    idle_minutes_to_autostop: Annotated[
        int | None,
        Parameter(
            name=["-i", "--idle-minutes-to-autostop"],
            help="Auto-stop after idle.",
        ),
    ] = None,
    down: Annotated[
        bool,
        Parameter(name="--down", help="Tear down after jobs finish."),
    ] = False,
    detach_run: Annotated[
        bool,
        Parameter(name=["-d", "--detach-run"], help="Don't stream logs."),
    ] = False,
    retry_until_up: Annotated[
        bool,
        Parameter(name=["-r", "--retry-until-up"], help="Retry until up."),
    ] = False,
    yes: Annotated[
        bool,
        Parameter(name=["-y", "--yes"], help="Skip confirmation prompts."),
    ] = False,
    dryrun: Annotated[
        bool,
        Parameter(name="--dryrun", help="Show what would be launched."),
    ] = False,
    name: Annotated[
        str | None,
        Parameter(name=["-n", "--name"], help="Task name."),
    ] = None,
    workdir: Annotated[
        str | None,
        Parameter(name="--workdir", help="Working directory to sync."),
    ] = None,
    env: Annotated[
        tuple[str, ...],
        Parameter(name=["-e", "--env"], help="Environment variable (KEY=VALUE)."),
    ] = (),
    sky: SkyClient | None = None,
) -> None:
    r"""Launch a cluster and run a task.

    Examples

        ml launch task.yaml -c mycluster
        ml launch 'echo hello' --gpus A100:1
        ml launch --gpus A100:4 -c dev
    """
    if _is_bare_invocation(
        entrypoint=entrypoint,
        cluster_name=cluster_name,
        gpus=gpus,
        cpus=cpus,
        memory=memory,
        cloud=cloud,
        region=region,
        use_spot=use_spot,
        num_nodes=num_nodes,
        name=name,
        workdir=workdir,
        env=env,
        dryrun=dryrun,
    ):
        interactive_onboarding(console)
        return

    if sky is None:
        sky = SkyClient()
    resources = ResourceParams(
        gpus=gpus,
        cpus=cpus,
        memory=memory,
        cloud=cloud,
        region=region,
        use_spot=use_spot,
    )
    task = build_task(
        TaskParams(
            entrypoint=entrypoint,
            name=name,
            workdir=workdir,
            num_nodes=num_nodes,
            env=env,
            resources=resources,
        ),
        sky,
    )

    if dryrun:
        console.print("[cyan]Dry run mode:[/cyan]")
        console.print(f"  Cluster: {cluster_name or '(auto-generated)'}")
        console.print(f"  Task: {task}")
        console.print(f"  Resources: {task.resources}")
        return

    _print_version_banner(console, sky)

    # Sync Mithril volumes to SkyPilot before launch. Volume values that are
    # strings are persistent volume names; dicts are ephemeral configs (no sync
    # needed).
    if task.volumes:
        volume_names = [v for v in task.volumes.values() if isinstance(v, str)]
        if volume_names:
            ensure_volumes_registered(
                volume_names,
                mithril=MithrilClient(),
                sky=sky,
            )
    try:
        request_id = sky.launch(
            task=task,
            cluster_name=cluster_name,
            retry_until_up=retry_until_up,
            idle_minutes_to_autostop=idle_minutes_to_autostop,
            dryrun=dryrun,
            down=down,
            _need_confirmation=not yes,
        )
    except click.exceptions.Abort:
        # User pressed Ctrl+C during confirmation prompt
        return

    result, cluster_name = poll_launch(
        request_id,
        console,
        sky=sky,
        initial_cluster_name=cluster_name,
    )
    job_id, handle = cast(tuple[int | None, object | None], result)
    cluster_name = _get_cluster_name(handle, cluster_name)

    exit_code = 0
    if not detach_run and job_id is not None:
        assert cluster_name is not None, "cluster_name should be set for jobs"
        exit_code = stream_job_logs(cluster_name, job_id, console, sky=sky)
    hints = _format_command_hints(job_id=job_id, cluster_name=cluster_name)
    if hints:
        console.print(hints, markup=False)
    if exit_code != 0:
        _exit(exit_code)


def build_task(params: TaskParams, sky: SkyClient) -> SkyTask:
    """Build a sky.Task from CLI parameters."""
    envs = parse_env_vars(params.env)

    if params.entrypoint and params.entrypoint.endswith((".yaml", ".yml")):
        task = sky.Task.from_yaml(params.entrypoint)
        if params.name:
            task.name = params.name
        if params.workdir:
            task.workdir = params.workdir
        if envs:
            task.update_envs(envs)
        if _has_resource_overrides(params.resources):
            task.set_resources(build_resources(params.resources, sky))
        if params.num_nodes:
            task.num_nodes = params.num_nodes
        return task

    resources = build_resources(params.resources, sky)
    task = sky.Task(
        name=params.name,
        run=params.entrypoint,
        workdir=params.workdir,
        envs=envs if envs else None,
        num_nodes=params.num_nodes or 1,
    )
    task.set_resources(resources)
    return task


def build_resources(params: ResourceParams, sky: SkyClient) -> SkyResources:
    """Build sky.Resources from CLI parameters."""
    kwargs: dict[str, Any] = {}
    clouds = sky.clouds
    if params.gpus:
        kwargs["accelerators"] = params.gpus
    if params.cpus:
        kwargs["cpus"] = params.cpus
    if params.memory:
        kwargs["memory"] = params.memory
    if params.region:
        kwargs["region"] = params.region
    if params.use_spot is not None:
        kwargs["use_spot"] = params.use_spot
    if params.cloud:
        cloud_cls = getattr(clouds, params.cloud.upper(), None)
        if cloud_cls is not None:
            kwargs["cloud"] = cloud_cls()
    return sky.Resources(**kwargs)


def _has_resource_overrides(params: ResourceParams) -> bool:
    """Return True if any resource overrides were provided."""
    return any(
        [
            params.gpus,
            params.cpus,
            params.memory,
            params.cloud,
            params.region,
            params.use_spot is not None,
        ]
    )


def parse_env_vars(env: tuple[str, ...]) -> dict[str, str]:
    """Parse environment variables from CLI args."""
    envs: dict[str, str] = {}
    for item in env:
        if "=" in item:
            key, _, value = item.partition("=")
            envs[key] = value
    return envs


def _print_version_banner(console: Console, sky: SkyClient) -> None:
    client_version = package_version("mithril-client")
    server_version = sky.api_info().version
    console.print(f"[dim]mithril-client {client_version}[/dim]", highlight=False)
    if server_version:
        console.print(
            f"[dim]SkyPilot API server {server_version}[/dim]", highlight=False
        )
    else:
        console.print("[dim]SkyPilot API server (version unknown)[/dim]")


def _get_cluster_name(handle: object | None, fallback: str | None) -> str | None:
    if handle is None:
        return fallback
    getter = getattr(handle, "get_cluster_name", None)
    if callable(getter):
        return cast("str | None", getter())
    return cast("str | None", getattr(handle, "cluster_name", fallback))


def _format_command_hints(job_id: int | None, cluster_name: str | None) -> str:
    if cluster_name is None:
        return ""
    lines = ["📋 Useful Commands"]
    if job_id is not None:
        lines.append(f"Job ID: {job_id}")
        lines.append(
            f"{INDENT_SYMBOL}To cancel the job:\t\t "
            f"{_cli_command('cancel')} {cluster_name} {job_id}"
        )
        lines.append(
            f"{INDENT_SYMBOL}To stream job logs:\t\t "
            f"{_cli_command('logs')} {cluster_name} {job_id}"
        )
        lines.append(
            f"{INDENT_LAST_SYMBOL}To view job queue:\t\t "
            f"{_cli_command('queue')} {cluster_name}"
        )
    lines.append(f"Cluster name: {cluster_name}")
    lines.append(f"{INDENT_SYMBOL}To log into the head VM:\t ssh {cluster_name}")
    lines.append(
        f"{INDENT_SYMBOL}To submit a job:\t\t "
        f"{_cli_command('exec')} {cluster_name} yaml_file"
    )
    lines.append(
        f"{INDENT_SYMBOL}To stop the cluster:\t {_cli_command('stop')} {cluster_name}"
    )
    lines.append(
        f"{INDENT_LAST_SYMBOL}To teardown the cluster:\t {_cli_command('down')} "
        f"{cluster_name}"
    )
    return "\n" + "\n".join(lines)


def _cli_command(command: str) -> str:
    if command in SKY_ALIAS_COMMANDS:
        return f"ml {command}"
    return f"ml sky {command}"
