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

from __future__ import annotations

import concurrent.futures
import json
from typing import TYPE_CHECKING, Any, Callable, Iterable, Iterator

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,
        volume_size_gb: int = 0,
        gpu_count: int = 1,
        model_names: list[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.volume_size_gb: int = volume_size_gb
        self.gpu_count: int = max(1, gpu_count)
        self.model_names: list[str] | None = model_names
        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
        self._last_logs: list[str] = []

    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)

        # Capture and print stdout from the remote container.
        self._last_logs = response.get("logs") or []
        for line in self._last_logs:
            print(f"[velar] {line}")

        return response.get("result")

    def map(
        self,
        inputs: Iterable[Any],
        *,
        return_exceptions: bool = False,
        max_workers: int = 50,
    ) -> Iterator[Any]:
        """Run the function on each input in parallel, yielding results in order.

        Usage:
            results = list(generate.map(["prompt1", "prompt2", "prompt3"]))

            # Tolerate individual failures:
            for result in generate.map(inputs, return_exceptions=True):
                if isinstance(result, Exception):
                    print(f"failed: {result}")
                else:
                    print(result)
        """
        items = list(inputs)
        if not items:
            return
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=min(len(items), max_workers)
        ) as executor:
            futures = [executor.submit(self.remote, item) for item in items]
            for future in futures:
                try:
                    yield future.result()
                except Exception as exc:
                    if return_exceptions:
                        yield exc
                    else:
                        raise

    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,
        volume_size_gb: int = 0,
        gpu_count: int = 1,
        slug: str | None = None,
        keep_warm: bool = False,
    ):
        self.func = func
        self.name = func.__name__
        self.gpu = gpu_spec
        self.image = image
        self.secrets: dict[str, str] = secrets or {}
        self.volume_size_gb: int = volume_size_gb
        self.gpu_count: int = max(1, gpu_count)
        self.slug: str | None = slug
        self.keep_warm: bool = keep_warm
        self.deployment_type = "endpoint"

        # Set after deploy().
        self.deployment_id: str | None = None
        self.endpoint_url: str | None = None
        self.invoke_url: str | None = None  # slug-based URL if slug is set
        self._client: VelarClient | None = None
        self._last_logs: list[str] = []

    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)

        # Capture and print stdout from the remote container.
        self._last_logs = response.get("logs") or []
        for line in self._last_logs:
            print(f"[velar] {line}")

        return response.get("result")

    def map(
        self,
        inputs: Iterable[Any],
        *,
        return_exceptions: bool = False,
        max_workers: int = 50,
    ) -> Iterator[Any]:
        """Run the endpoint on each input in parallel, yielding results in order."""
        items = list(inputs)
        if not items:
            return
        with concurrent.futures.ThreadPoolExecutor(
            max_workers=min(len(items), max_workers)
        ) as executor:
            futures = [executor.submit(self.remote, item) for item in items]
            for future in futures:
                try:
                    yield future.result()
                except Exception as exc:
                    if return_exceptions:
                        yield exc
                    else:
                        raise

    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,
    volume_size_gb: int = 0,
    gpu_count: int = 1,
    model: str | list[str] | None = None,
) -> Callable[[Callable[..., Any]], FunctionSpec]:
    """Decorator to define a serverless GPU function.

    Args:
        gpu: GPU type, e.g. "L4", "A100", "H100".
        image: Container image with dependencies.
        timeout: Max seconds to wait for the function to run (default 600).
        secrets: Environment variables/secrets injected into the container,
            e.g. ``{"HF_TOKEN": os.environ["HF_TOKEN"]}``.
        volume_size_gb: Extra pod volume in GB mounted at ``/data`` (default 0).
        gpu_count: Number of GPUs to allocate (default 1). Cost scales linearly.
        model: HuggingFace model ID(s) to cache on a persistent volume, e.g.
            ``"meta-llama/Llama-3.1-8B-Instruct"`` or a list of IDs.

    Usage:
        @velar.function(gpu="A100", image=my_image, timeout=300,
                        secrets={"HF_TOKEN": os.environ["HF_TOKEN"]},
                        volume_size_gb=100, gpu_count=2)
        def train(data):
            import os, torch
            token = os.environ["HF_TOKEN"]
            # /data is available for caching models
            ...
    """
    resolved_gpu = gpu_module.resolve(gpu)
    resolved_image = image or Image.debian_slim()

    # normalize model to list[str] or None
    _model_names: list[str] | None = None
    if model is not None:
        _model_names = [model] if isinstance(model, str) else list(model)

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

    return decorator


def endpoint(
    *,
    gpu: str | gpu_module.GPU = "L4",
    image: Image | None = None,
    secrets: dict[str, str] | None = None,
    volume_size_gb: int = 0,
    gpu_count: int = 1,
    slug: str | None = None,
    keep_warm: bool = False,
) -> Callable[[Callable[..., Any]], EndpointSpec]:
    """Decorator to define a persistent GPU endpoint.

    Args:
        gpu: GPU type, e.g. "L4", "A100", "H100".
        image: Container image with dependencies.
        secrets: Environment variables/secrets injected into the container.
        volume_size_gb: Extra pod volume in GB mounted at ``/data`` (default 0).
        gpu_count: Number of GPUs to allocate (default 1). Cost scales linearly.
        slug: Optional stable slug for a persistent URL that survives redeployments.
            Requires a Pro or Business plan. The invoke URL becomes:
            ``invoke.velar.run/{user_id}/{slug}``.
        keep_warm: If True, the pod stays running between requests (billed hourly).
            Requires ``slug`` to be set.

    Usage:
        @velar.endpoint(gpu="L4", image=my_image,
                        secrets={"OPENAI_API_KEY": os.environ["OPENAI_API_KEY"]},
                        volume_size_gb=50,
                        slug="my-model",
                        keep_warm=True)
        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.debian_slim()

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

    return decorator
