"""Modal-style decorators for defining GPU functions and endpoints."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any, Callable

from velar import gpu as gpu_module
from velar.image import Image
from velar.serialization import serialize

if TYPE_CHECKING:
    from velar.client import VelarClient


class FunctionSpec:
    """Spec for a serverless function decorated with @velar.function."""

    def __init__(
        self,
        func: Callable[..., Any],
        *,
        gpu_spec: gpu_module.GPU,
        image: Image,
        timeout: int,
        secrets: dict[str, str] | None = None,
    ):
        self.func = func
        self.name = func.__name__
        self.gpu = gpu_spec
        self.image = image
        self.timeout = timeout
        self.secrets: dict[str, str] = secrets or {}
        self.deployment_type = "serverless"

        # Set after deploy() so .remote() knows where to call.
        self.deployment_id: str | None = None
        self.endpoint_url: str | None = None
        self._client: VelarClient | None = None

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Call locally for testing."""
        return self.func(*args, **kwargs)

    def remote(self, *args: Any, **kwargs: Any) -> Any:
        """Run the function on a remote GPU.

        The function must have been deployed via ``app.deploy()`` before
        calling ``.remote()``.  The method will:
        1. Poll the deployment until it is running.
        2. POST the serialized arguments to the container endpoint.
        3. Return the deserialized result.
        """
        if self.deployment_id is None or self._client is None:
            raise RuntimeError(
                f"Function '{self.name}' is not deployed yet. "
                "Use `app.deploy()` first, then call `.remote()`."
            )

        # Wait for the container to be ready.
        self._client.poll_deployment(self.deployment_id, timeout=self.timeout)

        # Invoke the function via invoke.velar.run proxy.
        payload = serialize(list(args), kwargs)
        raw_response = self._client.invoke_function(
            self.deployment_id, payload, timeout=self.timeout
        )

        # Parse the response envelope produced by the generated handler.
        response = json.loads(raw_response)

        # Print any stdout captured on the remote container.
        for line in response.get("logs") or []:
            print(f"[velar] {line}")

        return response.get("result")

    def __repr__(self) -> str:
        return f"<VelarFunction {self.name} gpu={self.gpu.name}>"


class EndpointSpec:
    """Spec for a persistent endpoint decorated with @velar.endpoint."""

    def __init__(
        self,
        func: Callable[..., Any],
        *,
        gpu_spec: gpu_module.GPU,
        image: Image,
        secrets: dict[str, str] | None = None,
    ):
        self.func = func
        self.name = func.__name__
        self.gpu = gpu_spec
        self.image = image
        self.secrets: dict[str, str] = secrets or {}
        self.deployment_type = "endpoint"

        # Set after deploy().
        self.deployment_id: str | None = None
        self.endpoint_url: str | None = None
        self._client: VelarClient | None = None

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        """Call locally for testing."""
        return self.func(*args, **kwargs)

    def remote(self, *args: Any, **kwargs: Any) -> Any:
        """Invoke the endpoint on a remote GPU.

        Semantics are the same as :meth:`FunctionSpec.remote` but the
        deployment is expected to stay alive as a persistent service.
        """
        if self.deployment_id is None or self._client is None:
            raise RuntimeError(
                f"Endpoint '{self.name}' is not deployed yet. "
                "Use `app.deploy()` first, then call `.remote()`."
            )

        self._client.poll_deployment(self.deployment_id)

        payload = serialize(list(args), kwargs)
        raw_response = self._client.invoke_function(self.deployment_id, payload)

        response = json.loads(raw_response)

        # Print any stdout captured on the remote container.
        for line in response.get("logs") or []:
            print(f"[velar] {line}")

        return response.get("result")

    def __repr__(self) -> str:
        return f"<VelarEndpoint {self.name} gpu={self.gpu.name}>"


def function(
    *,
    gpu: str | gpu_module.GPU = "L4",
    image: Image | None = None,
    timeout: int = 600,
    secrets: dict[str, str] | None = None,
) -> Callable[[Callable[..., Any]], FunctionSpec]:
    """Decorator to define a serverless GPU function.

    Usage:
        @velar.function(gpu="A100", image=my_image, timeout=300,
                        secrets={"HF_TOKEN": os.environ["HF_TOKEN"]})
        def train(data):
            import os
            token = os.environ["HF_TOKEN"]  # available in the container
            ...
    """
    resolved_gpu = gpu_module.resolve(gpu)
    resolved_image = image or Image()

    def decorator(func: Callable[..., Any]) -> FunctionSpec:
        return FunctionSpec(
            func,
            gpu_spec=resolved_gpu,
            image=resolved_image,
            timeout=timeout,
            secrets=secrets,
        )

    return decorator


def endpoint(
    *,
    gpu: str | gpu_module.GPU = "L4",
    image: Image | None = None,
    secrets: dict[str, str] | None = None,
) -> Callable[[Callable[..., Any]], EndpointSpec]:
    """Decorator to define a persistent GPU endpoint.

    Usage:
        @velar.endpoint(gpu="L4", image=my_image,
                        secrets={"OPENAI_API_KEY": os.environ["OPENAI_API_KEY"]})
        def predict(request):
            import os
            client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])
            ...
    """
    resolved_gpu = gpu_module.resolve(gpu)
    resolved_image = image or Image()

    def decorator(func: Callable[..., Any]) -> EndpointSpec:
        return EndpointSpec(
            func,
            gpu_spec=resolved_gpu,
            image=resolved_image,
            secrets=secrets,
        )

    return decorator
