"""SkyPilot wrapper for stable imports and testability.

This module provides two interfaces:

1. **Module-level API** for SDK developers::

       from mithril import sky

       task = sky.Task(run="echo hello")
       resources = sky.Resources(accelerators="A100:4")
       sky.launch(task=task, cluster_name="dev", ...)

2. **SkyClient class** for CLI code and testing (supports parameter injection)::

       from mithril.sky import SkyClient

       client = SkyClient()
       # or inject a mock in tests
       build_task(params, sky=mock_client)

The module lazily imports SkyPilot to keep `ml --help` fast.
"""

from __future__ import annotations

import functools
from typing import TYPE_CHECKING, TypeVar

from mithril.sky.client import SkyClient

T = TypeVar("T")

if TYPE_CHECKING:
    from collections.abc import Iterable

    from sky import Task as SkyTask
    from sky import backends
    from sky.server.common import RequestId
    from sky.server.requests import payloads
    from sky.skylet import job_lib

__all__ = [
    "SkyClient",
    "api_info",
    "api_status",
    "ensure_local_api_server_version_matches_client",
    "get",
    "job_status",
    "launch",
    "tail_logs",
]


@functools.lru_cache(maxsize=1)
def _get_client() -> SkyClient:
    """Return the shared SkyClient instance, created lazily on first use."""
    return SkyClient()


# Module-level attributes (Task, Resources, clouds, volumes) are accessed via
# __getattr__ to maintain lazy loading - SkyPilot is only imported when first accessed.
_LAZY_ATTRS = {"Task", "Resources", "clouds", "volumes"}


def __getattr__(name: str) -> object:
    """Lazily expose SkyClient attributes at module level."""
    if name in _LAZY_ATTRS:
        return getattr(_get_client(), name)
    msg = f"module {__name__!r} has no attribute {name!r}"
    raise AttributeError(msg)


def __dir__() -> list[str]:
    """Include lazy attributes in dir() for discoverability."""
    return [*globals().keys(), *_LAZY_ATTRS]


def ensure_local_api_server_version_matches_client() -> None:
    """Initialize the shared SkyClient so its one-time preflight runs.

    SkyClient performs the local API server version check in ``__init__``.
    Calling ``_get_client()`` here ensures that preflight has run without
    duplicating the check.
    """
    _get_client()


def launch(
    *,
    task: SkyTask,
    cluster_name: str | None = None,
    retry_until_up: bool = False,
    idle_minutes_to_autostop: int | None = None,
    dryrun: bool = False,
    down: bool = False,
) -> RequestId[tuple[int | None, backends.ResourceHandle | None]]:
    """Launch a task on SkyPilot.

    Args:
        task: A sky.Task object
        cluster_name: Name of the cluster (auto-generated if None)
        retry_until_up: Keep retrying until the cluster is up
        idle_minutes_to_autostop: Auto-stop after idle for this many minutes
        dryrun: If True, don't actually launch
        down: Tear down cluster after job completes

    Returns:
        Request ID for the launch operation
    """
    return _get_client().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,
    )


def api_status(*, request_ids: list[str]) -> list[payloads.RequestPayload]:
    """Get status of SkyPilot requests."""
    return _get_client().api_status(request_ids=request_ids)


def get(request_id: RequestId[T]) -> T:  # noqa: UP047 - PEP 695 syntax requires 3.12+
    """Get result of a SkyPilot request."""
    return _get_client().get(request_id)


def tail_logs(
    *,
    cluster_name: str,
    job_id: int,
    follow: bool = True,
    preload_content: bool = False,
) -> Iterable[str | None]:
    """Stream logs from a SkyPilot job."""
    return _get_client().tail_logs(
        cluster_name=cluster_name,
        job_id=job_id,
        follow=follow,
        preload_content=preload_content,
    )


def job_status(
    cluster_name: str,
    *,
    job_ids: list[int] | None = None,
) -> RequestId[dict[int | None, job_lib.JobStatus | None]]:
    """Get status of a SkyPilot job (returns a request ID)."""
    return _get_client().job_status(cluster_name, job_ids=job_ids)


def api_info() -> object:
    """Get API server info (status, version, etc.)."""
    return _get_client().api_info()
