# SPDX-License-Identifier: Apache-2.0
"""Workflow request lifecycle adapter. Never replaces a model or GPU decorator.

V22's raw thread drops ContextVars used by Gradio/Spaces. This bridge copies
them, calls the SAME decorated render_all_segments once, and watches its
existing progress file. A deadline stops the caller waiting, not the GPU job.
An unfinished SDK call retains its lease: no second model job may overlap it.
"""
from contextvars import copy_context
import json
import math
import os
from pathlib import Path
import threading
import time

from gradio.context import LocalContext
from workflow_diagnostics import _active


class ExecutionBusy(RuntimeError):
    pass


class ExecutionCancelled(RuntimeError):
    pass


def cancellation_check():
    """Capture the current queued event before cancellation removes it.

    These are pinned Gradio 6.30 internals, covered by a real queued HTTP test.
    Direct Python calls have no event and remain usable. No auth is synthesized.
    """
    blocks = LocalContext.blocks.get(None)
    ident = LocalContext.event_id.get(None)
    events = getattr(getattr(blocks, "_queue", None), "event_ids_to_events", {})
    event = events.get(ident)

    def check():
        if event is not None and (not event.alive or events.get(ident) is not event):
            raise ExecutionCancelled("Execution cancelled. No completed take was saved.")
    return check


class GPUBridge:
    def __init__(self, render, budget, *, grace_seconds=180, poll_seconds=.25):
        self.render = render
        self.budget = budget
        self.grace_seconds = float(grace_seconds)
        if not math.isfinite(self.grace_seconds) or self.grace_seconds < 0:
            raise ValueError("WORKFLOW_GPU_WAIT_GRACE_SECONDS must be finite and nonnegative.")
        self.poll_seconds = poll_seconds
        self.lock = threading.Lock()
        self.pending = None

    def ensure_idle(self):
        with self.lock:
            if self.pending is not None:
                raise ExecutionBusy(
                    "A previous Viggle GPU call is still ending. No new model job was started. "
                    "Wait for it to finish; if it never returns, restart the Space before retrying.")

    def __call__(self, refs, conds, frame_counts, context, new, width, height,
                 audios, seed, arrays, anchors, progress_path, progress):
        args = (refs, conds, frame_counts, context, new, width, height, audios,
                seed, arrays, anchors, progress_path)
        seconds = float(self.budget(*args))
        if not math.isfinite(seconds) or seconds <= 0:
            raise ValueError("Invalid V22 GPU budget.")
        check_cancelled = cancellation_check()
        check_cancelled()
        done = threading.Event()
        outcome = {}
        inherited = copy_context()
        trace = _active.get()
        # Reserve before launching: a cancelled coordinator cannot free this
        # lease while the SDK's blocking result-queue read is still in flight.
        with self.lock:
            if self.pending is not None:
                raise ExecutionBusy("A previous Viggle GPU call is still ending; retry after it finishes.")
            self.pending = done

        def worker():
            try:
                if trace:
                    trace.emit("gpu_call_enter", budget_s=seconds)
                outcome["value"] = self.render(*args)
            except BaseException as exc:
                outcome["error"] = exc
            finally:
                try:
                    if trace:
                        trace.emit("gpu_call_return", success="error" not in outcome)
                finally:
                    with self.lock:
                        self.pending = None
                        done.set()

        thread = threading.Thread(target=inherited.run, args=(worker,), daemon=True,
                                  name="viggle-workflow-gpu")
        try:
            thread.start()
        except BaseException:
            with self.lock:
                self.pending = None
            raise

        # This is an application wait limit, not a larger requested allocation.
        deadline = time.monotonic() + seconds + self.grace_seconds
        total = max(1, len(frame_counts))
        last = None
        phase = "waiting_worker"
        while not done.is_set():
            check_cancelled()
            try:
                data = json.loads(Path(progress_path).read_text())
                finished = int(data["done"])
                current = int(data["current"])
                denoise = int(data.get("denoise", 0))
                steps = max(1, int(data.get("denoise_total", 3)))
                phase = str(data["phase"])
            except (OSError, ValueError, KeyError, TypeError):
                finished, current, denoise, steps = 0, 1, 0, 3
                phase = "waiting_worker"
            token = (phase, finished, current, denoise)
            if token != last:
                if phase == "loading_gpu":
                    fraction, desc = .025, "Viggle — worker entered · loading model on GPU…"
                elif phase == "preparing_segment":
                    fraction = min(.995, (finished + .06) / total)
                    desc = f"Viggle — preparing segment {current}/{total}…"
                elif phase == "denoise":
                    fraction = min(.999, (finished + denoise / steps) / total)
                    desc = f"Viggle — segment {current}/{total} · denoise {denoise}/{steps}"
                else:
                    fraction, desc = .01, "Viggle — waiting for ZeroGPU worker; model entry not confirmed…"
                if progress is not None:
                    progress(fraction, desc=desc)
                last = token
            if time.monotonic() >= deadline:
                raise TimeoutError(
                    f"Viggle did not return within {seconds + self.grace_seconds:g}s "
                    f"(last worker phase: {phase}). No completed take was saved. "
                    "The GPU call may still be active; another model job is blocked until it ends. "
                    "If it never returns, restart the Space before retrying.")
            done.wait(self.poll_seconds)
        check_cancelled()
        if "error" in outcome:
            raise outcome["error"]
        if progress is not None:
            progress((total, total), desc="Viggle — done", unit="segment")
        return outcome["value"]


_bridge = None


def ensure_gpu_idle():
    if _bridge is not None:
        _bridge.ensure_idle()


def install(core, budget):
    """Install only the UI/thread bridge after canonical eager CPU model boot."""
    global _bridge
    if _bridge is not None:
        raise RuntimeError("Workflow execution adapter was already installed.")
    _bridge = GPUBridge(core.render_all_segments, budget,
                        grace_seconds=os.environ.get("WORKFLOW_GPU_WAIT_GRACE_SECONDS", "180"))
    core._render_gpu_with_ui_progress = _bridge
    return _bridge