"""ComfyUI library-mode backend.

Single-process, single-implementation. The @spaces.GPU decorator is the only
divergence between local and HF Spaces deployment.
"""

from __future__ import annotations

import asyncio
import contextvars
import os
import pathlib
import sys
import threading
import traceback as tb_mod
from collections.abc import AsyncIterator, Iterable
from dataclasses import dataclass, field
from typing import Any

import models


@dataclass
class DownloadEvent:
    filename: str
    mb_done: float
    mb_total: float


@dataclass
class ProgressEvent:
    stage: int
    stage_label: str
    step: int
    total_steps: int


@dataclass
class OutputEvent:
    video_path: str
    audio_path: str | None = None
    meta: dict = field(default_factory=dict)


@dataclass
class ErrorEvent:
    category: str  # "oom" | "zerogpu_timeout" | "execution" | "interrupt" | "download"
    message: str
    stage: int | None = None
    traceback: str = ""


def _on_spaces() -> bool:
    return bool(os.environ.get("SPACES_ZERO_GPU"))


try:
    import spaces  # type: ignore
except ImportError:
    spaces = None  # type: ignore[assignment]


def _identity(fn):
    return fn


# --- Per-call ZeroGPU duration estimator -----------------------------------
# `duration` is a per-call timeout. Shorter declared duration → faster queue
# priority on the shared ZeroGPU pool. Estimating from (mode, preset, frames)
# instead of using a one-size-fits-all 600s cap means light T2V calls jump
# the queue while heavy modes (lipsync, style) reserve real headroom.

_BASE_DURATION_S: dict[str, int] = {
    # Rough sampler+decode time at ~120 frames, balanced preset, warm cache.
    "t2v": 90,
    "i2v": 90,
    "a2v": 120,
    "lipsync": 240,   # extra: audio encoder + audio VAE + extra LoRAs
    "keyframe": 180,
    "style": 360,     # extra: preprocessor (canny/dwpose/depth) + IC-LoRAs
}
_PRESET_MULT: dict[str, float] = {"fast": 1.0, "balanced": 1.5, "quality": 3.0}

# Highest *effective* (post-`duration_factor`) GPU duration this account is
# observed to accept. ZeroGPU multiplies the declared duration by a
# GPU-model-specific `duration_factor` (1.0 on H200, 1.5 on RTX PRO 6000
# Blackwell — see spaces/zero/configs.json) before testing it against the
# per-call max. Declaring our raw estimate meant the factored value (e.g.
# 240 -> 360 on a 1.5x GPU) overshot the cap and the scheduler rejected every
# call with "ZeroGPU illegal duration"; the handler then returned a failure
# card with no video, so the user saw progress then nothing — intermittently,
# since H200 (1.0x) landings still went through. As of 2026-06 the per-call max
# tightened: a declared 240s (factor 1.0 / H200) now trips "ZeroGPU illegal
# duration", so i2v could not schedule at all. 120s effective is the current
# standard ZeroGPU per-call ceiling and schedules reliably across GPU landings;
# fast/short i2v completes comfortably inside it on H200.
_EFFECTIVE_DURATION_CEILING_S = 120


def _zerogpu_duration_factor() -> float:
    """Live ZeroGPU duration multiplier for the GPU model this Space landed on.

    Mirrors what the spaces lib does at schedule time (client.py reads the same
    `duration_factor`). Falls back to the worst-case 1.5 if the config can't be
    read, so we never over-declare and trip the per-call maximum.
    """
    try:
        from spaces.zero.config import get_config

        factor = float(get_config().get("duration_factor", 1.0))
        return factor if factor > 0 else 1.0
    except Exception:
        return 1.5


def _frames_from_workflow(workflow: dict) -> int:
    """Read the frame count from the workflow's EmptyLTXVLatentVideo node."""
    for node in workflow.values():
        if isinstance(node, dict) and node.get("class_type") == "EmptyLTXVLatentVideo":
            try:
                return int((node.get("inputs") or {}).get("length", 121))
            except (TypeError, ValueError):
                return 121
    return 121


def _duration_for(
    executor: Any,
    workflow: dict,
    output_ids: list[str],
    mode: str,
    preset: str,
    multiplier: float = 1.0,
    progress: Any = None,
) -> int:
    """ZeroGPU duration estimator. Same signature as _execute_workflow.

    `progress` is a gr.Progress instance forwarded by the caller; we ignore it
    here (estimator doesn't emit progress) but must accept it positionally so
    ZeroGPU can call us with the same arg list it'll use for _execute_workflow.

    Estimate = (base × preset multiplier + cold-cache buffer + per-frame VAE
    decode time) × retry multiplier. The result is then clamped so that the
    value ZeroGPU actually charges against the per-call max — declared ×
    `duration_factor` (1.0 on H200, 1.5 on RTX PRO 6000) — stays at/under the
    known-good ceiling. Without dividing by the factor, a 1.5x GPU inflates the
    declared duration past the cap and the call is rejected with "ZeroGPU
    illegal duration" (client.py), leaving the user with progress then nothing.
    """
    base = _BASE_DURATION_S.get(mode, 180)
    mult = _PRESET_MULT.get(preset.lower(), 1.5)
    frames = _frames_from_workflow(workflow)
    est = int((base * mult + 60 + frames * 0.3) * multiplier)
    declared_ceiling = int(_EFFECTIVE_DURATION_CEILING_S / _zerogpu_duration_factor())
    return max(60, min(est, declared_ceiling))


# Decorate at module load time so ZeroGPU's startup analyzer detects it.
_GPU = (
    spaces.GPU(duration=_duration_for)
    if (spaces is not None and _on_spaces())
    else _identity
)


@_GPU
def _execute_workflow(
    executor: Any,
    workflow: dict,
    output_ids: list[str],
    mode: str,
    preset: str,
    multiplier: float = 1.0,
    progress: Any = None,
):
    """Run the workflow on GPU, streaming progress, ending with the video path.

    Generator — every yield is `("progress", step, total)` during sampling and
    the final yield is `("video", path)`. Yielding is the *only* way ComfyUI's
    per-step counter crosses the @spaces.GPU subprocess boundary on HF Spaces as
    discrete events: the `spaces` library streams a GPU generator's yields back
    to the parent, so the caller can turn them into ProgressEvents (and thus SSE
    `generating` frames). A plain return would run the whole diffusion inside the
    forked child with its progress hook firing on a queue the parent never sees.
    The yielded values are picklable tuples, fine across the boundary. The `mode`,
    `preset`, and `multiplier` args are consumed by `_duration_for` to size the
    GPU slot. `progress` (a `gr.Progress`) is still updated for the native bar.
    """
    import time as _time
    _gpu_start = _time.time()
    try:
        yield from _execute_workflow_inner(
            executor, workflow, output_ids, mode, preset, multiplier, progress,
        )
    finally:
        print(
            f"[backend] GPU time consumed: {_time.time() - _gpu_start:.2f}s "
            f"(mode={mode}, preset={preset}, multiplier={multiplier})",
            file=sys.stderr,
            flush=True,
        )


def _execute_workflow_inner(
    executor: Any,
    workflow: dict,
    output_ids: list[str],
    mode: str,
    preset: str,
    multiplier: float = 1.0,
    progress: Any = None,
):
    import comfy.utils as _cu
    import queue as _queue
    import threading as _threading

    step_q: _queue.Queue = _queue.Queue()

    def _gp_hook(value, total, _preview=None, **_kw):
        try:
            v, t = int(value), int(total)
            if progress is not None:
                progress(v / max(t, 1), desc=f"Sampling step {v}/{t}")
            step_q.put((v, t))
        except Exception:
            pass

    _cu.set_progress_bar_global_hook(_gp_hook)

    result: dict = {}

    def _run():
        try:
            executor.execute(
                workflow,
                prompt_id="ltx23-aio",
                extra_data={"client_id": "ltx23-aio"},
                execute_outputs=output_ids,
            )
            hist = getattr(executor, "history_result", {}) or {}
            outs = hist.get("outputs") or {}
            path = ""
            for output in outs.values():
                if not isinstance(output, dict):
                    continue
                for value in output.values():
                    if not isinstance(value, list):
                        continue
                    for item in value:
                        if isinstance(item, dict):
                            fn = item.get("filename") or ""
                            if fn.endswith((".mp4", ".webm", ".mov")):
                                path = item.get("fullpath") or fn
            result["path"] = path
        except Exception as exc:  # surfaced to the generator's main thread
            result["error"] = exc
        finally:
            step_q.put(None)

    worker = _threading.Thread(target=_run, daemon=True)
    worker.start()
    while True:
        item = step_q.get()
        if item is None:
            break
        yield ("progress", item[0], item[1])
    worker.join()
    if "error" in result:
        raise result["error"]
    yield ("video", result.get("path", ""))


class _StubServer:
    """Minimal stub matching the surface ComfyUI's PromptExecutor expects."""

    client_id: str | None = "ltx23-aio"
    last_node_id: str | None = None

    def send_sync(self, event: str, data: dict, sid: str | None = None) -> None:
        pass

    def queue_updated(self) -> None:
        pass


class _StubPromptQueue:
    """Stub matching the surface VideoHelperSuite + others touch."""

    currently_running: dict = {}
    history: dict = {}
    flags: dict = {}

    def get_current_queue(self) -> tuple[list, list]:
        return ([], [])

    def get_tasks_remaining(self) -> int:
        return 0

    def set_flag(self, name: str, data) -> None:
        pass

    def get_flags(self, *a, **kw) -> dict:
        return {}

    def task_done(self, *a, **kw) -> None:
        pass

    def put(self, *a, **kw) -> None:
        pass

    def wipe_queue(self) -> None:
        pass

    def delete_queue_item(self, *a, **kw) -> None:
        pass


class _StubPromptServerInstance:
    """Surface that ComfyUI's `server.PromptServer.instance` exposes to custom nodes.

    VideoHelperSuite, KJNodes, and others read this at import time. They mostly
    use it to register HTTP routes or send WS events or peek at the prompt queue.
    No-ops here are fine — we have no real server.
    """

    client_id: str | None = "ltx23-aio"
    # KJNodes' preview thread reads `last_node_id.encode('ascii')` directly.
    # ComfyUI's real server keeps it as a string per executing node and resets
    # to None at end-of-prompt — which races the preview thread. Keep it a
    # safe non-empty string so .encode() never NPEs.
    last_node_id: str = "ltx23-aio"
    web_root: str = ""

    class _Routes:
        def get(self, *a, **kw):
            return lambda fn: fn

        def post(self, *a, **kw):
            return lambda fn: fn

        def static(self, *a, **kw):
            return None

    routes = _Routes()
    sockets: dict = {}
    prompt_queue = _StubPromptQueue()
    # Custom-Scripts checks PromptServer.instance.supports — claim the
    # "custom_nodes_from_web" capability so it skips its JS install path.
    supports: list[str] = ["custom_nodes_from_web"]
    web_root: str = ""

    def add_routes(self) -> None:
        pass

    def send_sync(self, event: str, data: dict, sid: str | None = None) -> None:
        pass

    def send_progress_text(self, text: str, node_id=None, sid=None) -> None:
        # Comfy_extras nodes call this; we just no-op since we don't have a UI
        # to surface intermediate text on.
        pass

    def queue_updated(self) -> None:
        pass

    def get_node_class_def(self, *a, **kw):
        return None

    def __getattr__(self, name):
        # Anything else our custom nodes might reach for — give them a no-op.
        # This is a deliberate liberal catch-all so the inference path doesn't
        # die on cosmetic UI hooks. Inspection-style access (hasattr) gets True.
        def _noop(*a, **kw):
            return None
        return _noop


def _comfy_dir() -> pathlib.Path:
    if _on_spaces():
        return pathlib.Path.home() / "comfyui"
    return pathlib.Path(__file__).parent / "comfyui"


class ComfyUILibraryBackend:
    """Wraps PromptExecutor for in-process workflow execution."""

    def __init__(self) -> None:
        self._comfy_dir = _comfy_dir()
        if not self._comfy_dir.exists():
            raise RuntimeError(
                f"ComfyUI not found at {self._comfy_dir}. "
                f"Local: run `bash setup.sh`. Spaces: see app.py:_bootstrap()."
            )
        if str(self._comfy_dir) not in sys.path:
            sys.path.insert(0, str(self._comfy_dir))

        # Defer comfy imports until the path is set up.
        # NOTE: ComfyUI ships PromptExecutor in the top-level `execution.py`
        # module, NOT under `comfy.execution`. Same for `nodes`. Both must be
        # imported AFTER the sys.path insert above.
        import asyncio
        import threading

        import comfy.cli_args  # noqa: F401 — side-effect: registers CLI flags
        import execution  # top-level module — provides PromptExecutor
        import nodes  # top-level module — provides init_extra_nodes (async)

        # CRITICAL ordering fix: ComfyUI's nodes.py:24 inserts `comfyui/comfy/`
        # at sys.path[0]. That dir contains a module-style `utils.py`, which
        # shadows `comfyui/utils/` (a package containing install_util.py).
        # Some custom nodes (KJNodes, VideoHelperSuite via app.frontend_management)
        # do `from utils.install_util import …` and get `comfy/utils.py` instead,
        # raising "'utils' is not a package". Rewrite sys.path so comfy_dir is
        # ahead of comfy_dir/comfy and force-clear any cached `utils` binding.
        comfy_subdir = str(self._comfy_dir / "comfy")
        sys.path = [p for p in sys.path if p not in (str(self._comfy_dir), comfy_subdir)]
        sys.path.insert(0, comfy_subdir)
        sys.path.insert(0, str(self._comfy_dir))
        if "utils" in sys.modules and not getattr(sys.modules["utils"], "__path__", None):
            del sys.modules["utils"]

        # Some custom nodes (e.g. VideoHelperSuite) read `server.PromptServer.instance`
        # at import time. We don't run a real ComfyUI server, so install a stub
        # that exposes the attributes those nodes touch (sockets, send, etc.).
        import server as comfy_server

        if getattr(comfy_server.PromptServer, "instance", None) is None:
            comfy_server.PromptServer.instance = _StubPromptServerInstance()

        # `nodes.init_extra_nodes` is async. We may be called from within a
        # running event loop (Gradio's handler) — running `asyncio.run()` there
        # raises. Run the coroutine in a fresh loop on a worker thread instead.
        def _init_in_thread() -> None:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                loop.run_until_complete(nodes.init_extra_nodes())
            finally:
                loop.close()

        thread = threading.Thread(target=_init_in_thread, daemon=False)
        thread.start()
        thread.join()
        # PromptExecutor expects a `server` with client_id, send_sync, last_node_id,
        # queue_updated. A minimal stub no-ops all of them — we don't run a real
        # websocket server, we surface progress via comfy.utils.PROGRESS_BAR_HOOK.
        # cache_args["ram"] is read unconditionally inside execute_async even when
        # cache_type is the default false — provide a sensible default so it doesn't
        # NoneType-subscript at line 727.
        self._executor = execution.PromptExecutor(
            server=_StubServer(),
            cache_args={"ram": 16.0, "lru": 0},
        )

    def __repr__(self) -> str:
        return f"ComfyUILibraryBackend(comfy_dir={self._comfy_dir!r})"

    async def submit(
        self,
        mode: str,
        workflow: dict,
        *,
        preset: str = "balanced",
        duration_multiplier: float = 1.0,
        gpu_duration: int = 0,  # legacy, ignored (now derived from preset+frames)
        progress: Any = None,
    ) -> AsyncIterator[Any]:
        """Run a workflow end-to-end. Yields Download/Progress/Output/Error events.

        `preset` and `duration_multiplier` flow through to the @spaces.GPU
        duration estimator. The handler can re-call submit() with
        duration_multiplier=2.0 if the first attempt aborts on timeout.
        """
        # Pre-flight: ensure all model files exist.
        try:
            needed = models.walk_workflow_for_models(workflow)
            for download_event in models.ensure_models(needed):
                yield download_event
        except Exception as e:
            yield ErrorEvent(
                category="download",
                message=str(e),
                traceback=tb_mod.format_exc(),
            )
            return

        # Run the inference in a worker thread; pass progress events through a queue.
        queue: asyncio.Queue = asyncio.Queue()
        loop = asyncio.get_running_loop()

        def _push(event: Any) -> None:
            asyncio.run_coroutine_threadsafe(queue.put(event), loop)

        # Track stage progression. ComfyUI fires the progress hook from inside
        # samplers, so we advance the stage every time we observe a new sampler
        # starting (step==0 with a different total than before, or a "new run"
        # signal — value smaller than the running max for the same total).
        progress_state = {"stage": 0, "prev_total": -1, "max_step": -1}

        def _hook(value: int, total: int, _preview=None, **_kwargs: Any) -> None:
            v, t = int(value), int(total)
            # New sampler started (different total, or step rewound)
            if t != progress_state["prev_total"] or v < progress_state["max_step"]:
                progress_state["stage"] += 1
                progress_state["prev_total"] = t
                progress_state["max_step"] = v
            else:
                progress_state["max_step"] = max(progress_state["max_step"], v)
            _push(
                ProgressEvent(
                    stage=progress_state["stage"],
                    stage_label="diffusion",
                    step=v,
                    total_steps=t,
                )
            )

        def _worker() -> None:
            import comfy.utils

            saved_hook = getattr(comfy.utils, "PROGRESS_BAR_HOOK", None)
            try:
                # Workflow is already API-format (saved from ComfyUI editor's
                # "Save (API Format)"), so it can be handed to PromptExecutor
                # directly. The execute_outputs list pinpoints which output
                # nodes to evaluate — we let PromptExecutor walk the whole
                # graph by passing every output-class node id.
                output_ids = [
                    nid for nid, n in workflow.items()
                    if n.get("class_type", "").startswith(("SaveVideo", "VHS_VideoCombine", "PreviewAudio", "CreateVideo"))
                ]
                print(
                    f"[backend] submitting workflow: {len(workflow)} nodes, "
                    f"output_ids={output_ids}",
                    file=sys.stderr,
                    flush=True,
                )
                # _execute_workflow is module-level and decorated with a
                # @spaces.GPU(duration=callable) on Spaces — the callable
                # estimates per-call timeout from (mode, preset, frames) so
                # light calls get fast queue priority while heavy ones reserve
                # real headroom. Off-Spaces it's a plain call. It's now a
                # generator: it yields ("progress", step, total) as ComfyUI
                # samples (these survive the @spaces.GPU subprocess boundary,
                # unlike a global hook) and a final ("video", path). We feed
                # each progress yield to `_hook`, which does the stage tracking
                # and pushes a ProgressEvent (→ SSE) onto the queue.
                video_path = ""
                for _item in _execute_workflow(
                    self._executor, workflow, output_ids, mode, preset, duration_multiplier, progress,
                ):
                    if not isinstance(_item, (tuple, list)) or not _item:
                        continue
                    if _item[0] == "progress":
                        _hook(_item[1], _item[2])
                    elif _item[0] == "video":
                        video_path = _item[1]
                # Fallback: if history_result didn't surface a path (rare on
                # Spaces — happens when ZeroGPU's subprocess boundary drops
                # mutated state), scan the output dir for the newest mp4
                # written within the last 60 s.
                if not video_path:
                    video_path = _newest_recent_video(self._comfy_dir / "output") or ""
                print(
                    f"[backend] workflow done; video_path={video_path!r}",
                    file=sys.stderr,
                    flush=True,
                )
                _push(OutputEvent(video_path=video_path))
            except Exception as exc:
                tb_text = tb_mod.format_exc()
                print(f"[backend] worker exception:\n{tb_text}", file=sys.stderr, flush=True)
                _push(
                    ErrorEvent(
                        category=_classify(exc),
                        message=str(exc),
                        traceback=tb_text,
                    )
                )
            finally:
                comfy.utils.set_progress_bar_global_hook(saved_hook)
                _free_memory()
                _push(None)  # sentinel: stop the consumer

        # ZeroGPU's @spaces.GPU wrapper reads the user's identity from the
        # current Gradio request via gradio.context.LocalContext.request,
        # which is a contextvar. Plain threads don't inherit contextvars, so
        # without this the worker sees request=None, X-IP-Token never gets
        # read, and `client.schedule` raises "Space app has reached its GPU
        # limit" (token-is-None branch in spaces/zero/client.py:138). Copy
        # the calling task's context so the request — and therefore the Pro
        # user's quota attribution — survives the thread boundary.
        ctx = contextvars.copy_context()
        thread = threading.Thread(target=ctx.run, args=(_worker,), daemon=True)
        thread.start()

        while True:
            event = await queue.get()
            if event is None:
                return
            yield event

    def interrupt(self) -> None:
        """Cancel the currently running workflow (if any)."""
        try:
            import comfy.model_management as mm

            mm.interrupt_current_processing()
        except Exception:
            pass


def _classify(exc: Exception) -> str:
    name = type(exc).__name__.lower()
    msg = str(exc).lower()
    if "outofmemory" in name or "cuda out of memory" in msg:
        return "oom"
    if "expired zerogpu proxy token" in msg or "expired" in msg and "token" in msg:
        return "expired_token"
    if "illegal duration" in msg:
        return "illegal_duration"
    if "unlogged user" in msg:
        return "unlogged"
    if "exceeded your" in msg and "gpu" in msg:
        return "quota_exceeded"
    # ZeroGPU enforces the @spaces.GPU(duration=N) cap and re-raises as
    # gradio.exceptions.Error('GPU task aborted').
    if "gpu task aborted" in msg or ("gpu" in msg and "aborted" in msg):
        return "gpu_timeout"
    if "interrupt" in name:
        return "interrupt"
    return "execution"


def _free_memory() -> None:
    """Free VRAM after a workflow finishes (success or failure)."""
    try:
        import comfy.model_management as mm

        mm.unload_all_models()
    except Exception:
        pass
    try:
        import torch

        if torch.backends.mps.is_available():
            torch.mps.empty_cache()
    except Exception:
        pass
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _newest_recent_video(output_root: pathlib.Path, within_seconds: float = 60.0) -> str | None:
    """Filesystem fallback: return the newest .mp4/.webm/.mov under *output_root*
    that was modified within the last *within_seconds* seconds.

    Used when the executor's history_result didn't surface a path — typically
    happens when ZeroGPU's subprocess boundary drops the mutation. The disk
    is shared, so the file is there even when the in-memory state isn't.
    """
    import time

    if not output_root.exists():
        return None
    cutoff = time.time() - within_seconds
    candidates: list[tuple[float, pathlib.Path]] = []
    for ext in (".mp4", ".webm", ".mov"):
        for p in output_root.rglob(f"*{ext}"):
            try:
                mtime = p.stat().st_mtime
            except OSError:
                continue
            if mtime >= cutoff:
                candidates.append((mtime, p))
    if not candidates:
        return None
    candidates.sort(reverse=True)
    return str(candidates[0][1])