"""ZeroGPU-safe decorator.

spaces.GPU attaches a device for the duration of the call and enforces a
wall-clock ceiling. Wrapping the whole pipeline in one call blows that ceiling
on anything longer than a short clip, so each GPU stage gets its own call and
all CPU work stays outside. Off ZeroGPU this is a no-op.

Adds over the bare decorator:
  * the worker traceback is printed and appended to DWCT_GPU_TRACE, because the
    ZeroGPU wrapper otherwise reduces a failure to a bare exception class name;
"""
from __future__ import annotations

import functools
import os
import sys
import traceback

TRACE_FILE = os.environ.get("DWCT_GPU_TRACE", "/tmp/dwct_gpu_trace.txt")

try:
    import spaces

    _GPU = spaces.GPU
    HAVE_ZEROGPU = True
except Exception:
    HAVE_ZEROGPU = False


def _loud(fn):
    @functools.wraps(fn)
    def inner(*a, **k):
        try:
            return fn(*a, **k)
        except BaseException as exc:
            text = "[gpu] %s raised %s: %s\n%s" % (
                fn.__name__, type(exc).__name__, exc, traceback.format_exc())
            print(text, file=sys.stderr)
            try:
                with open(TRACE_FILE, "a") as fh:
                    fh.write(text + "\n" + "=" * 60 + "\n")
            except Exception:
                pass
            sys.stderr.flush()
            sys.stdout.flush()
            raise

    return inner


def gpu(duration: int = 120):
    if HAVE_ZEROGPU:
        def deco(fn):
            return _GPU(duration=int(duration))(_loud(fn))
        return deco

    def deco(fn):
        return _loud(fn)

    return deco