"""
Decorators for Compute resources.

This module provides decorators to simplify working with Compute resources,
automatically managing their lifecycle.
"""

import functools
import inspect
from typing import Callable, Dict, List, Optional

try:
    from kubernetes.client import CoreV1Event
except ImportError as e:
    raise ImportError(
        f"kubernetes package is required: {e}\n" "Install with: pip install kubernetes"
    )


def compute(
    name: str,
    *,
    image: Optional[str] = None,
    command: Optional[List[str]] = None,
    args: Optional[List[str]] = None,
    env: Optional[Dict[str, str]] = None,
    resources: Optional[Dict[str, str]] = None,
    labels: Optional[Dict[str, str]] = None,
    annotations: Optional[Dict[str, str]] = None,
    suspend: bool = False,
    keep_alive: bool = False,
    node: Optional[str] = None,
    volumes: Optional[List[Dict[str, str]]] = None,
    on_events: Optional[Callable[[CoreV1Event], None]] = None,
    on_metrics: Optional[Callable[[object], None]] = None,
):
    """
    Decorator that automatically manages a Compute instance lifecycle.

    Creates a Compute instance with the provided parameters, enters its async
    context (waiting for it to be ready), passes it to the decorated function,
    and ensures cleanup on exit (deletes the Compute resource).

    Args:
        name: Name of the Compute resource
        image: Container image for the Compute runtime
        command: Entrypoint command override
        args: Arguments for the entrypoint
        env: Environment variables as dict
        resources: Resource requirements as dict
        labels: Labels to apply to the Compute resources
        annotations: Annotations to apply to the Compute resources
        suspend: Whether to create the Compute in suspended state
        keep_alive: Whether to keep the Compute resource alive after the function exits
        node: Node hostname to schedule the Compute pod on
        volumes: Simplified volume definitions as list of dicts with keys:
                 name (volume name), storage (e.g. "10Gi"), path (mount path),
                 storageClass (optional), accessMode (optional, default "ReadWriteOnce")
        on_events: Optional callback to receive events for this Compute resource
        on_metrics: Optional callback to receive metrics from this Compute resource

    Returns:
        Decorator function that wraps async functions

    Example:
        >>> @compute(
        ...     name="my-compute",
        ...     resources={"nvidia.com/gpu": "1"}
        ... )
        ... async def process_data(compute: Compute):
        ...     tensors = await compute.receive_tensors(count=1)
        ...     # Process tensors...
        ...     result = tensors[0] * 2
        ...     await compute.send_tensors(result)
        >>>
        >>> await process_data()
    """

    def decorator(func):
        # Import here to avoid circular import
        from skytorch.client.compute import Compute

        # 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 or references 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(
                name=name,
                image=image,
                command=command,
                args=args,
                env=env,
                resources=resources,
                labels=labels,
                annotations=annotations,
                suspend=suspend,
                node=node,
                volumes=volumes,
                on_events=on_events,
                on_metrics=on_metrics,
            )

            # Enter context manager (waits for ready and initializes gRPC)
            # When keep_alive is set, skip deletion on exit so the pod
            # persists across invocations (model stays in GPU memory).
            async with compute_instance as c:
                if keep_alive:
                    # Prevent __aexit__ from deleting the Compute resource
                    c._keep_alive = True
                if compute_param_name:
                    func_kwargs[compute_param_name] = c
                    return await func(*func_args, **func_kwargs)
                else:
                    return await func(c, *func_args, **func_kwargs)

        return wrapper

    return decorator
