"""Hidden half-res output restore for Moxie MultiTrack projects.

Single-pass project sampling conditions and samples at half the
editor-selected resolution (see nodes/project.py). This module owns the
other half of that contract: Moxie Project Video Combine restores the exact
editor-selected resolution right before its images output, with no
user-visible wiring anywhere.

Hardware path adapted from NVIDIA's RTX Video Super Resolution ComfyUI node
(Apache License 2.0):

    https://github.com/Comfy-Org/Nvidia_RTX_Nodes_ComfyUI
    Copyright 2025 NVIDIA Corporation

It drives the ``nvidia-vfx`` pip package (RTX Video Super Resolution) per
frame via DLPack. When the package or a supported GPU is unavailable, the
module falls back to ComfyUI's Lanczos resize so the delivered frame size is
always exact and rendering never fails because of the upscaler.
"""

from __future__ import annotations

from typing import Any

__all__ = [
    "rtx_vsr_upscale",
    "restore_project_video_resolution",
    "upscale_images_to_project_resolution",
]


def _log(message: str) -> None:
    try:
        from .log import log_node_info
    except Exception:  # pragma: no cover - standalone tooling without the pack
        print(f"[Moxie] {message}", flush=True)  # noqa: T201
        return
    log_node_info("MultiTrack Output", message)


_FALLBACK_NOTICE_SHOWN = False


def _notice_fallback(message: str) -> None:
    global _FALLBACK_NOTICE_SHOWN
    if _FALLBACK_NOTICE_SHOWN:
        return
    _FALLBACK_NOTICE_SHOWN = True
    _log(message)


def _lanczos_upscale(images: Any, output_width: int, output_height: int) -> Any:
    import comfy.utils

    samples = images[..., :3].movedim(-1, 1)
    upscaled = comfy.utils.common_upscale(
        samples, output_width, output_height, "lanczos", "disabled"
    )
    return upscaled.movedim(1, -1)


def rtx_vsr_upscale(
    images: Any,
    output_width: int,
    output_height: int,
    quality: str = "ULTRA",
) -> Any:
    """Upscale an IMAGE batch to the exact target size, hardware first.

    Frames are processed per frame through RTX Video Super Resolution
    (adapted from Comfy-Org/Nvidia_RTX_Nodes_ComfyUI, Apache 2.0). Any
    failure - missing package, unsupported GPU, VRAM pressure - degrades to
    the Lanczos fallback instead of breaking the render, and the delivered
    size is corrected to the exact target whenever the hardware snap misses.
    """
    import torch

    frame_count, in_height, in_width, channels = images.shape
    if frame_count == 0 or (in_height, in_width) == (output_height, output_width):
        return images

    hardware_ready = channels == 3
    if hardware_ready:
        try:
            import nvvfx
        except Exception:
            hardware_ready = False
            _notice_fallback(
                "Hardware upscale unavailable (install nvidia-vfx for RTX "
                "acceleration); using Lanczos fallback."
            )

    if hardware_ready:
        # /8 snap from the NVIDIA node; editor selections are 8-aligned so
        # this is an identity operation for real project resolutions.
        snapped_width = max(8, round(output_width / 8) * 8)
        snapped_height = max(8, round(output_height / 8) * 8)
        max_batch_pixels = 1024 * 1024 * 16
        batch_size = max(
            1, max_batch_pixels // max(1, snapped_width * snapped_height)
        )
        quality_levels = {
            "LOW": nvvfx.effects.QualityLevel.LOW,
            "MEDIUM": nvvfx.effects.QualityLevel.MEDIUM,
            "HIGH": nvvfx.effects.QualityLevel.HIGH,
            "ULTRA": nvvfx.effects.QualityLevel.ULTRA,
        }
        selected_quality = quality_levels.get(
            str(quality).upper(), nvvfx.effects.QualityLevel.HIGH
        )
        out_tensor = torch.empty(
            (frame_count, snapped_height, snapped_width, channels),
            device=images.device,
            dtype=images.dtype,
        )
        try:
            with nvvfx.VideoSuperRes(selected_quality) as super_res:
                super_res.output_width = snapped_width
                super_res.output_height = snapped_height
                super_res.load()
                for start in range(0, frame_count, batch_size):
                    batch = images[start : start + batch_size]
                    batch_cuda = (
                        batch.cuda().permute(0, 3, 1, 2).float().contiguous()
                    )
                    for offset in range(batch_cuda.shape[0]):
                        result = super_res.run(batch_cuda[offset]).image
                        out_tensor[start + offset : start + offset + 1] = (
                            torch.from_dlpack(result).movedim(0, -1).unsqueeze(0)
                        )
        except Exception as error:
            hardware_ready = False
            _notice_fallback(
                f"Hardware upscale failed ({error}); using Lanczos fallback."
            )
        else:
            if (snapped_height, snapped_width) != (output_height, output_width):
                return _lanczos_upscale(out_tensor, output_width, output_height)
            return out_tensor

    return _lanczos_upscale(images, output_width, output_height)


def _project_output_resolution(project_name: Any) -> tuple[int, int] | None:
    """Editor-selected resolution stored in the project manifest, or ``None``.

    The manifest is written from the editor selection at the start of every
    run (nodes/project.py), so reading it fresh here can never see a stale
    frontend widget value.
    """
    try:
        from .h3_project import load_h3_project_data, safe_h3_project_name

        data = load_h3_project_data(safe_h3_project_name(project_name))
        width = int(data.get("width") or 0)
        height = int(data.get("height") or 0)
    except Exception:
        return None
    if width <= 0 or height <= 0 or (width, height) == (32, 32):
        return None
    return width, height


def upscale_images_to_project_resolution(
    images: Any,
    project_name: Any,
) -> Any:
    """Restore the project's editor resolution on combined frames.

    Used by Moxie Project Video Combine right before the ``images`` output.
    Frames already at the project resolution - full-res renders, dual-pass,
    SelfLift and passthrough outputs - pass through untouched; only frames
    whose size differs from the manifest (the hidden half-res sampling
    signature) are restored to the exact editor-selected resolution.
    """
    try:
        import torch
    except ImportError:  # pragma: no cover - torch is required in ComfyUI
        return images
    if (
        not isinstance(images, torch.Tensor)
        or images.ndim != 4
        or images.shape[0] == 0
    ):
        return images
    target = _project_output_resolution(project_name)
    if target is None:
        return images
    output_width, output_height = target
    if (int(images.shape[2]), int(images.shape[1])) == (
        output_width,
        output_height,
    ):
        return images
    return rtx_vsr_upscale(images, output_width, output_height)


def restore_project_video_resolution(
    video: Any,
    project_name: Any,
    source_path: str | None = None,
) -> Any:
    """Restore the manifest resolution on the composed project video.

    The combine composes the timeline at the segments' own size, which the
    hidden half-res sampling contract makes smaller than the manifest. This
    helper returns the video object untouched when the composed file already
    matches the manifest (full-res renders, dual-pass, SelfLift, passthrough);
    otherwise it decodes the frames once, upscales them to the exact manifest
    resolution, and rebuilds the video so every combine output (VIDEO,
    images, audio, fps) carries the editor-selected size. Any failure keeps
    the original video - rendering must never break because of the restore.
    """
    try:
        target = _project_output_resolution(project_name)
        if target is None:
            return video
        output_width, output_height = target
        if source_path:
            # Cheap ffprobe first: only pay the frame decode when the composed
            # file is genuinely smaller than the project resolution.
            from .video import ffprobe_info

            info = ffprobe_info(str(source_path))
            probed_width = info.get("width") if isinstance(info, dict) else None
            probed_height = info.get("height") if isinstance(info, dict) else None
            if (
                isinstance(probed_width, int)
                and isinstance(probed_height, int)
                and (probed_width, probed_height) == (output_width, output_height)
            ):
                return video
        components = video.get_components()
        images = upscale_images_to_project_resolution(
            components.images, project_name
        )
        if images is components.images:
            return video
        from comfy_api.latest import InputImpl, Types

        return InputImpl.VideoFromComponents(
            Types.VideoComponents(
                images=images,
                audio=components.audio,
                frame_rate=components.frame_rate,
            )
        )
    except Exception as error:
        _log(
            f"Output restore skipped ({error}); the composed video keeps its "
            "original resolution."
        )
        return video