# SPDX-License-Identifier: Apache-2.0
# Copyright (c) Axolotl AI
# Licensed under the Apache License, Version 2.0

"""Minimal Hugging Face Spaces ZeroGPU probe surface.

This module is intentionally small and dependency-light: importing it does not
import Gradio, spaces, or torch. The prototype Space imports those packages
lazily only when the Space app is built or the GPU probe is invoked.
"""

from __future__ import annotations

import os
import time
from collections.abc import Callable
from typing import Any

DEFAULT_PROBE_DURATION_S = 10
MAX_ZEROGPU_DURATION_S = 120


def coerce_duration(value: Any = None) -> int:
    """Clamp a requested ZeroGPU duration to the supported prototype range."""
    if value is None or value == "":
        value = os.environ.get("HATCHERY_ZEROGPU_PROBE_DURATION", DEFAULT_PROBE_DURATION_S)
    try:
        duration = int(float(value))
    except (TypeError, ValueError):
        duration = DEFAULT_PROBE_DURATION_S
    return max(1, min(MAX_ZEROGPU_DURATION_S, duration))


def probe_duration(duration_s: Any = None) -> int:
    """Callable duration hook compatible with ``@spaces.GPU(duration=...)``."""
    return coerce_duration(duration_s)


def _import_spaces() -> Any | None:
    try:
        import spaces
    except ImportError:
        return None
    return spaces


def _decorate_gpu(
    fn: Callable[[Any], dict[str, Any]],
    *,
    spaces_mod: Any | None = None,
) -> Callable[[Any], dict[str, Any]]:
    spaces_mod = _import_spaces() if spaces_mod is None else spaces_mod
    if spaces_mod is None:
        return fn
    return spaces_mod.GPU(duration=probe_duration)(fn)


def make_gpu_probe(*, spaces_mod: Any | None = None) -> Callable[[Any], dict[str, Any]]:
    """Build the ZeroGPU-decorated probe function.

    The returned callable imports torch inside the decorated function body. That
    preserves the ZeroGPU rule that GPU work happens inside ``@spaces.GPU``.
    """

    def _probe(duration_s: Any = None) -> dict[str, Any]:
        requested_duration_s = coerce_duration(duration_s)
        started = time.perf_counter()

        import torch

        cuda_available = bool(torch.cuda.is_available())
        device_count = int(torch.cuda.device_count()) if cuda_available else 0
        device_name = torch.cuda.get_device_name(0) if cuda_available else None
        allocation_ok = False
        if cuda_available:
            torch.zeros(1, device="cuda")
            torch.cuda.synchronize()
            allocation_ok = True

        return {
            "ok": True,
            "requested_duration_s": requested_duration_s,
            "elapsed_s": round(time.perf_counter() - started, 4),
            "torch_version": getattr(torch, "__version__", "unknown"),
            "cuda_available": cuda_available,
            "cuda_device_count": device_count,
            "cuda_device_name": device_name,
            "cuda_allocation_ok": allocation_ok,
        }

    return _decorate_gpu(_probe, spaces_mod=spaces_mod)


def build_fastapi_app(
    gpu_probe: Callable[[Any], dict[str, Any]] | None = None,
    *,
    title: str = "Hatchery ZeroGPU Probe API",
) -> Any:
    """Build the FastAPI probe app mounted by the prototype Gradio Space."""
    from fastapi import FastAPI, Query

    probe = gpu_probe or make_gpu_probe()
    app = FastAPI(title=title)

    @app.get("/health")
    def health() -> dict[str, Any]:
        return {
            "ok": True,
            "service": "hatchery-zerogpu-probe",
            "max_duration_s": MAX_ZEROGPU_DURATION_S,
        }

    @app.get("/zerogpu/probe")
    def zerogpu_probe(
        duration_s: float | None = Query(default=None, ge=1, le=MAX_ZEROGPU_DURATION_S),
    ) -> dict[str, Any]:
        return probe(duration_s)

    return app


def build_demo() -> Any:
    """Build the Gradio UI and mount the FastAPI probe app below ``/v1``."""
    import gradio as gr

    probe = make_gpu_probe()
    with gr.Blocks(title="Hatchery ZeroGPU Probe") as demo:
        gr.Markdown("# Hatchery ZeroGPU Probe")
        gr.Markdown(
            "Run a tiny GPU allocation probe from Gradio, then test the mounted "
            "FastAPI route at `/v1/zerogpu/probe`."
        )
        duration = gr.Slider(
            minimum=1,
            maximum=MAX_ZEROGPU_DURATION_S,
            value=DEFAULT_PROBE_DURATION_S,
            step=1,
            label="Requested duration seconds",
        )
        output = gr.JSON(label="Probe result")
        run = gr.Button("Run GPU probe", variant="primary")
        run.click(fn=probe, inputs=[duration], outputs=[output], api_name="zerogpu_probe")

    api = build_fastapi_app(probe)
    demo.app.mount("/v1", api)
    demo.app.mount("/api/v1", api)
    return demo