"""Direct Compute for SkyTorch gRPC server."""

import asyncio
import functools
import inspect
import logging
import os
import re
from typing import Callable, Optional

import torch

from skytorch.client.grpc import GRPCClient

logger = logging.getLogger(__name__)


class Compute:
    """
    Compute for direct connection to SkyTorch gRPC server without Kubernetes.

    Provides the same interface as skytorch.client.Compute for device management
    and gRPC communication, but connects directly to a gRPC server URL.

    Usage:
        from skytorch.torch.server.testing import Compute

        async with Compute("localhost:50051") as compute:
            device = compute.device("cpu")
            x = torch.tensor([1, 2, 3], device=device)
            y = x + 1
            print(y.cpu())
    """

    # Device string parsing pattern
    _DEVICE_PATTERN = re.compile(r"^([a-zA-Z_]+)(?::(\d+))?$")

    def __init__(
        self,
        url: str = "",
        name: str = "compute",
        on_metrics: Optional[Callable[[object], None]] = None,
    ):
        """
        Initialize Compute.

        Args:
            url: gRPC server URL (host:port). Defaults to SKYTORCH_SERVER_URL
                 environment variable or "localhost:50051".
            name: Name for this compute (used in error messages).
            on_metrics: Optional callback to receive metrics from this Compute resource.
        """
        self.url = url or os.environ.get("SKYTORCH_SERVER_URL", "localhost:50051")
        self.name = name
        self._on_metrics = on_metrics
        self._grpc_client: Optional[GRPCClient] = None
        self._metrics_stream_task: Optional[asyncio.Task] = None

    def _parse_url(self) -> tuple[str, int]:
        """Parse URL into host and port."""
        if ":" in self.url:
            host, port_str = self.url.rsplit(":", 1)
            return host, int(port_str)
        return self.url, 50051

    async def _stream_metrics(self):
        """Stream metrics from this Compute resource and call the callback."""
        try:
            async for snapshot in self._grpc_client.metrics.stream_metrics(
                metric_names=[
                    "gpu.utilization.compute",
                    "gpu.utilization.memory",
                    "gpu.memory.used",
                    "gpu.power.usage",
                ],
                interval_seconds=1.0,
            ):
                self._on_metrics(snapshot)
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"Error streaming metrics: {e}")

    async def __aenter__(self) -> "Compute":
        """Connect to the gRPC server."""
        host, port = self._parse_url()
        self._grpc_client = GRPCClient(host=host, port=port)
        await self._grpc_client.__aenter__()

        # Start metrics streaming if callback is provided
        if self._on_metrics is not None:
            self._metrics_stream_task = asyncio.create_task(self._stream_metrics())

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Disconnect from the gRPC server."""
        # Cancel metrics stream task if running
        if self._metrics_stream_task is not None:
            self._metrics_stream_task.cancel()
            try:
                await self._metrics_stream_task
            except asyncio.CancelledError:
                pass
            self._metrics_stream_task = None

        if self._grpc_client:
            await self._grpc_client.__aexit__(exc_type, exc_val, exc_tb)
            self._grpc_client = None

    def device(self, type: str = "cpu", index: Optional[int] = None) -> torch.device:
        """
        Get a sky device mapped to this Compute.

        Args:
            type: Remote device type, optionally with index (e.g., "cuda", "cuda:0", "cpu")
            index: Remote device index (default: 0). Cannot be specified if type
                   already contains an index.

        Returns:
            torch.device with type "sky" and mapped local index

        Raises:
            RuntimeError: If type contains an index and index is also passed explicitly,
                          or if the device string format is invalid

        Example:
            >>> compute = Compute("localhost:50051")
            >>> device = compute.device("cuda")      # Same as cuda:0
            >>> device = compute.device("cuda:1")    # Uses index 1
            >>> device = compute.device("cuda", 1)   # Same as cuda:1
        """
        from skytorch.torch.backend._device import device_manager

        # Validate and parse device string
        match = self._DEVICE_PATTERN.match(type)
        if not match:
            raise RuntimeError(
                f"Invalid device string: {type!r}. "
                f"Expected format: 'device_type' or 'device_type:index'"
            )

        device_type = match.group(1)
        parsed_index = match.group(2)

        if parsed_index is not None:
            if index is not None:
                raise RuntimeError(
                    f"type (string) must not include an index because index was "
                    f"passed explicitly: {type}"
                )
            device_index = int(parsed_index)
        else:
            device_index = index if index is not None else 0

        return device_manager.get_sky_device(self, device_type, device_index)


def compute(
    url: str = "localhost:50051",
    *,
    name: str = "compute",
    on_metrics: Optional[Callable[[object], None]] = None,
):
    """
    Decorator that automatically manages a Compute instance lifecycle.

    Creates a Compute instance, connects to the gRPC server, and passes
    the compute instance to the decorated function.

    Args:
        url: gRPC server URL (host:port). Defaults to SKYTORCH_SERVER_URL
             environment variable or "localhost:50051".
        name: Name for this compute (used in error messages).
        on_metrics: Optional callback to receive metrics from this Compute resource.

    Returns:
        Decorator function that wraps async functions.

    Example:
        >>> from skytorch.torch.server import compute, Compute
        >>>
        >>> @compute("localhost:50051")
        ... async def test_addition(compute: Compute):
        ...     device = compute.device()
        ...     x = torch.tensor([1, 2, 3], device=device)
        ...     y = torch.tensor([4, 5, 6], device=device)
        ...     z = x + y
        ...     print(z.cpu())
        >>>
        >>> await test_addition()
    """

    def decorator(func):
        # Inspect the function signature to find a Compute parameter
        sig = inspect.signature(func)
        compute_param_name = None

        for param_name, param in sig.parameters.items():
            if param.annotation is not inspect.Parameter.empty:
                # Check if the annotation is Compute
                if param.annotation is Compute or (
                    hasattr(param.annotation, "__name__")
                    and param.annotation.__name__ == "Compute"
                ):
                    compute_param_name = param_name
                    break

        @functools.wraps(func)
        async def wrapper(*func_args, **func_kwargs):
            compute_instance = Compute(url=url, name=name, on_metrics=on_metrics)

            async with compute_instance as c:
                if compute_param_name:
                    # Pass as keyword argument to the typed parameter
                    func_kwargs[compute_param_name] = c
                    return await func(*func_args, **func_kwargs)
                else:
                    # Fallback to positional argument (first position)
                    return await func(c, *func_args, **func_kwargs)

        return wrapper

    return decorator
