"""SkyPilot client wrapper."""

from __future__ import annotations

import importlib
import warnings
from functools import cached_property
from typing import TYPE_CHECKING, TypeVar

T = TypeVar("T")

if TYPE_CHECKING:
    from collections.abc import Iterable
    from types import ModuleType

    from sky import Resources as SkyResources
    from sky import Task as SkyTask
    from sky import backends
    from sky.schemas.api import responses
    from sky.schemas.api.responses import APIHealthResponse
    from sky.server.common import RequestId
    from sky.server.requests import payloads
    from sky.skylet import job_lib
    from sky.volumes import volume as volume_lib


class SkyClient:
    """SkyPilot client that lazily imports the SkyPilot module.

    Use this class when you need dependency injection (e.g., for testing)::

        from mithril.sky import SkyClient


        def build_task(params, sky: SkyClient):
            return sky.Task(...)


        # In tests, pass a MagicMock as the client
        mock_client = MagicMock()
        build_task(params, sky=mock_client)

    For simple usage, prefer the module-level API::

        from mithril import sky

        task = sky.Task(run="echo hello")
        sky.launch(task=task, cluster_name="dev")
    """

    def __init__(self) -> None:
        """Initialize the client and run one-time API-server compatibility preflight."""
        self._api_server_preflight_done = False
        self.ensure_local_api_server_version_matches_client()

    @cached_property
    def _module(self) -> ModuleType:
        return importlib.import_module("sky")

    @property
    def Task(self) -> type[SkyTask]:  # noqa: N802 - mirrors sky.Task
        """SkyPilot Task class."""
        return self._module.Task

    @property
    def Resources(self) -> type[SkyResources]:  # noqa: N802 - mirrors sky.Resources
        """SkyPilot Resources class."""
        return self._module.Resources

    @property
    def clouds(self) -> ModuleType:
        """SkyPilot cloud provider module."""
        return self._module.clouds

    @cached_property
    def volumes(self) -> VolumesAPI:
        """SkyPilot volumes API."""
        return VolumesAPI(self._module)

    def launch(
        self,
        *,
        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,
        _need_confirmation: bool = False,
    ) -> RequestId[tuple[int | None, backends.ResourceHandle | None]]:
        """Launch a task on SkyPilot."""
        return self._module.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=_need_confirmation,
        )

    def api_status(self, *, request_ids: list[str]) -> list[payloads.RequestPayload]:
        """Return request status objects."""
        return self._module.api_status(request_ids=request_ids)

    def get(self, request_id: RequestId[T]) -> T:
        """Call sky.get."""
        return self._module.get(request_id)

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

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

    def api_info(self) -> APIHealthResponse:
        """Get API server info (status, version, etc.)."""
        return self._module.api_info()

    @cached_property
    def _server_common(self) -> ModuleType:
        return importlib.import_module("sky.server.common")

    @property
    def client_version(self) -> str:
        """SkyPilot client package version."""
        return self._module.__version__

    def api_server_url(self) -> str:
        """Get API server URL."""
        return self._server_common.get_server_url()

    def api_start(self) -> None:
        """Start SkyPilot API server."""
        self._module.api_start()

    def api_stop(self) -> None:
        """Stop SkyPilot API server."""
        self._module.api_stop()

    def api_restart(self) -> None:
        """Restart SkyPilot API server."""
        self.api_stop()
        self.api_start()

    def ensure_local_api_server_version_matches_client(self) -> None:
        """Restart local API server when server/client versions differ."""
        if self._api_server_preflight_done:
            return

        api_server_url = self.api_server_url()
        # Reuse SkyPilot's own "local endpoint" definition instead of keeping
        # a parallel URL parser in mithril.
        if not self._server_common.is_api_server_local(api_server_url):
            self._api_server_preflight_done = True
            return

        # Call the underlying SDK directly here so this check stays non-recursive.
        server_version = self._module.api_info().version
        if server_version and server_version != self.client_version:
            self.api_restart()
            refreshed_server_version = self._module.api_info().version
            if (
                refreshed_server_version
                and refreshed_server_version != self.client_version
            ):
                warnings.warn(
                    "SkyPilot API server version still differs after restart: "
                    f"client={self.client_version}, "
                    f"server={refreshed_server_version}.",
                    stacklevel=2,
                )
        self._api_server_preflight_done = True


class VolumesAPI:
    """SkyPilot volumes API wrapper."""

    def __init__(self, sky_module: ModuleType) -> None:
        """Initialize volumes API with SkyPilot module."""
        self._module = sky_module

    @property
    def Volume(self) -> type[volume_lib.Volume]:  # noqa: N802 - mirrors sky.volumes.Volume
        """SkyPilot Volume class."""
        return self._module.volumes.Volume

    def ls(self) -> RequestId[list[responses.VolumeRecord]]:
        """List volumes registered with SkyPilot.

        Returns:
            Request ID for the list request.
        """
        return self._module.volumes.ls()

    def apply(self, volume: volume_lib.Volume) -> RequestId[None]:
        """Apply a volume specification.

        Args:
            volume: Volume object to apply.

        Returns:
            Request ID of the apply request.
        """
        return self._module.volumes.apply(volume)

    def delete(self, names: list[str], *, purge: bool = False) -> RequestId[None]:
        """Delete volumes.

        Args:
            names: List of volume names to delete.
            purge: If True, delete the volume from the database even if the
                deletion API fails.

        Returns:
            Request ID of the delete request.
        """
        return self._module.volumes.delete(names, purge=purge)
