# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Demo Gradio app for ZeroGPU with transformers serve.

This app demonstrates how to deploy a Gradio frontend that calls
`transformers serve` over HTTP — **and auto-starts the serve process
when you run ``python app.py``**.

Architecture:

    ┌────────────────────────────────────────────────────────────┐
    │                    Gradio App (app.py)                      │
    │                                                            │
    │  ┌──────────────────────────┐    HTTP /v1/chat/completions  │
    │  │  Gradio UI               │ ────────────────────────────▶ │
    │  └──────────────────────────┘                               │
    │                                    ◀─────────────────────── │
    │  ┌──────────────────────────┐                               │
    │  │ transformers serve       │  (child process, auto-started)│
    │  │ model eager-loaded at    │  on dynamically allocated GPU │
    │  │ startup                  │  (ZeroGPU if enabled)         │
    │  └──────────────────────────┘                               │
    │                                                            │
    │  User just runs:                                            │
    │    python app.py                                            │
    └────────────────────────────────────────────────────────────┘

The Gradio app pre-downloads the model to disk, starts ``transformers serve``
as a subprocess with eager model loading (``force_model``), waits for it to be
ready, and tears it down when the app closes.

ZeroGPU support is automatic when running in a Hugging Face Space
(``SPACE_ID`` env var is set).

To run locally:
    pip install -r requirements.txt
    python app.py

To deploy as a ZeroGPU Space:
    1. Push to HF Hub
    2. Settings → Hardware → ZeroGPU
    3. The serve process will use dynamic GPU allocation
"""

import json
import os
import subprocess
import time
from urllib.request import Request, urlopen

import gradio as gr


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

def _is_zerogpu_space() -> bool:
    """Detect if we are running in a Hugging Face ZeroGPU Space.

    Checks for ``SPACE_ID`` env var, which is present in all HF Spaces.
    """
    return bool(os.environ.get("SPACE_ID"))


SERVE_URL = os.environ.get("SERVE_URL", "http://127.0.0.1:8000")
MODEL_ID = os.environ.get("TRANSFORMERS_ZEROGPU_MODEL", "google/gemma-4-26B-A4B-it")
SERVE_DEVICE = os.environ.get("SERVE_DEVICE", "cuda" if _is_zerogpu_space() else "auto")
SERVE_HOST = os.environ.get("SERVE_HOST", "127.0.0.1")
SERVE_PORT = int(os.environ.get("SERVE_PORT", 8000))


def _build_serve_cmd() -> list[str]:
    """Build the command line to launch ``transformers serve``.

    The model is passed as ``force_model`` (positional arg) so it eager-loads
    at startup. The pre-download ensures the timeout countdown only starts
    after the model is on disk.

    ZeroGPU detection is handled automatically by ``transformers serve``
    internally (it checks for ``SPACE_ID`` env var). No flags needed here.
    """
    return [
        "transformers", "serve",
        MODEL_ID,  # force_model (positional) → eager load instead of lazy
        "--device", SERVE_DEVICE,
        "--host", SERVE_HOST,
        "--port", str(SERVE_PORT),
        "--log-level", "warning",
    ]


# ---------------------------------------------------------------------------
# Serve lifecycle management
# ---------------------------------------------------------------------------
_serve_proc: subprocess.Popen | None = None


def _wait_for_serve(timeout: int = 180) -> bool:
    """Poll the serve health endpoint until it is ready or timeout."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            req = Request(f"{SERVE_URL}/health")
            resp = urlopen(req, timeout=1)
            return json.loads(resp.read()).get("status") == "ok"
        except Exception:
            time.sleep(0.5)
    return False


def _pre_download_model(model_id: str):
    """Pre-download the model so serve doesn't stall on first request.

    Without this the health-check timeout starts counting before the model
    is even downloaded — a large model can take minutes.
    """
    from huggingface_hub import snapshot_download

    print(f"📦 Downloading model '{model_id}' (may take a minute)... ")
    snapshot_download(repo_id=model_id)
    print("✅ Model downloaded")


def _start_serve():
    """Pre-download, then start ``transformers serve`` as a child process."""
    global _serve_proc
    if _serve_proc is not None and _serve_proc.poll() is None:
        if _wait_for_serve(5):
            return  # already running and healthy

    # Pre-download so the timeout countdown only starts AFTER the model is on disk
    _pre_download_model(MODEL_ID)

    cmd = _build_serve_cmd()
    print(f"Starting serve: {' '.join(cmd)}")
    _serve_proc = subprocess.Popen(cmd)  # inherits stdout/stderr so logs are visible
    if _wait_for_serve(timeout=180):
        print("✅ Serve API is ready")
    else:
        _serve_proc.kill()
        _serve_proc = None
        raise RuntimeError("Serve API did not start in time")


def _stop_serve():
    """Shut down the serve subprocess."""
    global _serve_proc
    if _serve_proc is not None and _serve_proc.poll() is None:
        print("Stopping serve...")
        _serve_proc.terminate()
        try:
            _serve_proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _serve_proc.kill()
            _serve_proc.wait(timeout=2)
        _serve_proc = None
        print("Serve stopped")


# ---------------------------------------------------------------------------
# HTTP client helpers for transformers serve
# ---------------------------------------------------------------------------

def _chat_completions(
    messages: list[dict],
    model: str,
    max_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
) -> str:
    """Call the ``/v1/chat/completions`` endpoint (non-streaming).

    Args:
        messages: List of chat messages (OpenAI format).
        model: Model name/ID.
        max_tokens: Maximum new tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling parameter.

    Returns:
        The generated text content.
    """
    body = json.dumps({
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "stream": False,
    }).encode()

    req = Request(
        f"{SERVE_URL}/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    resp = urlopen(req, timeout=300)
    data = json.loads(resp.read())
    return data["choices"][0]["message"]["content"]


def _chat_completions_stream(
    messages: list[dict],
    model: str,
    max_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    """Call the ``/v1/chat/completions`` endpoint (streaming, SSE).

    Yields each text chunk as it arrives from the server.

    Args:
        messages: List of chat messages (OpenAI format).
        model: Model name/ID.
        max_tokens: Maximum new tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling parameter.

    Yields:
        `str`: Each text chunk from the streaming response.
    """
    body = json.dumps({
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p,
        "stream": True,
    }).encode()

    req = Request(
        f"{SERVE_URL}/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    resp = urlopen(req, timeout=300)

    for line in resp:
        chunk = line.decode("utf-8").strip()
        if chunk.startswith("data: "):
            payload = chunk[6:]
            if payload == "[DONE]":
                break
            try:
                event = json.loads(payload)
                content = event["choices"][0]["delta"].get("content", "")
                if content:
                    yield content
            except (json.JSONDecodeError, KeyError):
                continue


# ---------------------------------------------------------------------------
# Inference functions — pure HTTP wrappers
# ---------------------------------------------------------------------------

def generate_non_streaming(
    prompt: str,
    max_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
) -> str:
    """Generate text by calling ``transformers serve`` (non-streaming).

    All inference goes through the serve HTTP API — no model is loaded here.
    The model is eager-loaded at serve startup (no download latency at request time).

    Args:
        prompt: The user's input text.
        max_tokens: Maximum new tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling parameter.

    Returns:
        The generated text.
    """
    messages = [{"role": "user", "content": prompt}]
    return _chat_completions(
        messages=messages,
        model=MODEL_ID,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )


def generate_streaming(
    prompt: str,
    max_tokens: int = 256,
    temperature: float = 0.7,
    top_p: float = 0.9,
):
    """Generate text by calling ``transformers serve`` (streaming).

    Yields chunks as they arrive from the server.
    The model is eager-loaded at serve startup (no download latency at request time).

    Args:
        prompt: The user's input text.
        max_tokens: Maximum new tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling parameter.

    Yields:
        `str`: Each text chunk from the server.
    """
    messages = [{"role": "user", "content": prompt}]
    yield from _chat_completions_stream(
        messages=messages,
        model=MODEL_ID,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )


# ---------------------------------------------------------------------------
# Gradio interface
# ---------------------------------------------------------------------------

def _is_server_ready() -> bool:
    """Check if the serve process is up and responding."""
    try:
        resp = urlopen(f"{SERVE_URL}/health", timeout=2)
        return json.loads(resp.read()).get("status") == "ok"
    except Exception:
        return False


def create_interface() -> gr.Blocks:
    """Create the Gradio Blocks interface.

    Both tabs call the ``transformers serve`` HTTP API — no local model
    is loaded by this Gradio app. The model is eager-loaded at startup.
    """
    with gr.Blocks(title="Transformers Serve — ZeroGPU Demo") as demo:
        mode = "ZeroGPU (dynamic GPU allocation)" if _is_zerogpu_space() else "Local (persistent GPU)"
        gr.Markdown(f"""
        # 🤗 Transformers Serve

        A Gradio frontend that calls ``transformers serve`` over HTTP.
        **Everything starts automatically — just run ``python app.py``.**

        - **Serve API**: `{SERVE_URL}`
        - **Model**: `{MODEL_ID}`
        - **Mode**: `{mode}`
        - **API endpoints**: ``/v1/chat/completions`` (streaming + non-streaming)

        ### Architecture

        This Gradio app pre-downloads the model, then starts ``transformers serve``
        as a child process with eager model loading. The model is ready before
        any request arrives:

        ```
        Gradio UI  ──HTTP──▶  transformers serve  ──▶  GPU
          (app.py)              (child process, eager load + pre-download)
        ```

        ZeroGPU Spaces are detected automatically by ``transformers serve``
        via the ``SPACE_ID`` environment variable.

        ### Quick start

        ```bash
        pip install -r requirements.txt
        python app.py        # ← starts both the Gradio UI AND serve
        ```
        """)

        with gr.Tabs():
            # ── Non-streaming tab ──
            with gr.Tab("Non-streaming"):
                with gr.Row():
                    with gr.Column(scale=3):
                        prompt_non_stream = gr.Textbox(
                            label="Prompt",
                            placeholder="Write something here...",
                            lines=3,
                        )
                        ns_max_tokens = gr.Slider(
                            label="Max new tokens",
                            minimum=32, maximum=1024, value=256, step=32,
                        )
                        ns_temperature = gr.Slider(
                            label="Temperature",
                            minimum=0.0, maximum=2.0, value=0.7, step=0.1,
                        )
                        ns_top_p = gr.Slider(
                            label="Top-p",
                            minimum=0.0, maximum=1.0, value=0.9, step=0.05,
                        )
                        ns_btn = gr.Button("Generate", variant="primary")
                        ns_clear_btn = gr.Button("Clear")

                    with gr.Column(scale=3):
                        output_non_stream = gr.Textbox(
                            label="Response",
                            lines=12,
                            interactive=False,
                        )

                ns_btn.click(
                    fn=generate_non_streaming,
                    inputs=[prompt_non_stream, ns_max_tokens, ns_temperature, ns_top_p],
                    outputs=output_non_stream,
                )
                ns_clear_btn.click(
                    fn=lambda: ("", ""),
                    inputs=None,
                    outputs=[prompt_non_stream, output_non_stream],
                )

            # ── Streaming tab ──
            with gr.Tab("Streaming"):
                with gr.Row():
                    with gr.Column(scale=3):
                        prompt_stream = gr.Textbox(
                            label="Prompt",
                            placeholder="Write something here...",
                            lines=3,
                        )
                        s_max_tokens = gr.Slider(
                            label="Max new tokens",
                            minimum=32, maximum=1024, value=256, step=32,
                        )
                        s_temperature = gr.Slider(
                            label="Temperature",
                            minimum=0.0, maximum=2.0, value=0.7, step=0.1,
                        )
                        s_top_p = gr.Slider(
                            label="Top-p",
                            minimum=0.0, maximum=1.0, value=0.9, step=0.05,
                        )
                        s_btn = gr.Button("Generate (streaming)", variant="primary")
                        s_clear_btn = gr.Button("Clear")

                    with gr.Column(scale=3):
                        output_stream = gr.Markdown(label="Response")

                s_btn.click(
                    fn=generate_streaming,
                    inputs=[prompt_stream, s_max_tokens, s_temperature, s_top_p],
                    outputs=output_stream,
                )
                s_clear_btn.click(
                    fn=lambda: "",
                    inputs=None,
                    outputs=output_stream,
                )

        # Status indicator
        gr.Markdown(
            f"*API status: {'✅ Connected' if _is_server_ready() else '❌ Not connected'} "
            f"| Serve URL: ``{SERVE_URL}``*"
        )

    return demo


if __name__ == "__main__":
    print(f"Model: {MODEL_ID}")
    print(f"Serve URL: {SERVE_URL}")
    print(f"Device: {SERVE_DEVICE}")

    if _is_zerogpu_space():
        print("🚀 ZeroGPU Space detected — serve will auto-detect and use dynamic GPU allocation")
        import spaces

        @spaces.GPU
        def fn():
            """required because there is a dummy check in spaces to fail early if the decorator is not present in the app.py file"""
            pass
    else:
        print("🖥️  Local mode — serve will use persistent GPU")

    # Start serve, launch Gradio, clean up on exit.
    # ZeroGPU detection and GPU allocation are handled automatically
    # inside the serve process — no decorator needed here.
    _start_serve()
    try:
        demo = create_interface()
        demo.launch()
    finally:
        _stop_serve()