#!/usr/bin/env python3
"""Offline NVIDIA Maxine VSR post-process for MiniMax H3 videos.

The default path is designed for a native H3 1280x736 render:
1280x736 -> Maxine 1.5x (1920x1104) -> centered crop -> 1920x1088.
Audio packets are copied from the source file without re-encoding.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import subprocess
import sys
import time
import traceback


WORKER_ENV = "H3_MAXINE_VSR_WORKER"
_EFFECT_GUARD: list[object] = []

# Conservative GPU prefilters. Flat regions receive a small separable blur and
# low-motion pixels receive a small previous-frame blend. Edge and motion masks
# protect waves, waterfalls, mist, architecture and camera movement.
DENOISE_PRESETS = {
    "mild": {
        "spatial_strength": 0.28,
        "edge_threshold": 0.035,
        "temporal_strength": 0.10,
        "temporal_threshold": 0.025,
    },
    "medium": {
        "spatial_strength": 0.42,
        "edge_threshold": 0.045,
        "temporal_strength": 0.16,
        "temporal_threshold": 0.035,
    },
    # Conservative finishing pass for footage that has already received the
    # medium prefilter. It concentrates temporal cleanup in stable flat areas
    # while using a tighter edge mask to protect architecture and wave detail.
    "finish": {
        "spatial_strength": 0.36,
        "edge_threshold": 0.032,
        "temporal_strength": 0.18,
        "temporal_threshold": 0.030,
    },
}


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Upscale an H3 video with NVIDIA Maxine VideoSuperRes."
    )
    parser.add_argument("input", nargs="?", type=Path, help="Source MP4/MOV file")
    parser.add_argument("output", nargs="?", type=Path, help="Destination MP4 file")
    parser.add_argument("--scale", type=float, default=1.5, help="Maxine scale before crop")
    parser.add_argument("--width", type=int, default=1920, help="Final output width")
    parser.add_argument("--height", type=int, default=1088, help="Final output height")
    parser.add_argument(
        "--quality",
        default="HIGHBITRATE_HIGH",
        help="nvvfx quality enum name; HIGHBITRATE_HIGH is tuned for clean H3 output",
    )
    parser.add_argument("--device", type=int, default=0, help="CUDA device index")
    parser.add_argument(
        "--encoder",
        default="h264_nvenc",
        choices=("h264_nvenc", "hevc_nvenc", "libx264"),
        help="Video encoder used after upscaling",
    )
    parser.add_argument("--bitrate", type=int, default=20_000_000, help="Video bitrate")
    parser.add_argument("--cq", type=int, default=18, help="NVENC constant-quality target")
    parser.add_argument(
        "--denoise",
        choices=("none", *DENOISE_PRESETS),
        default="none",
        help="Adaptive temporal denoise applied before Maxine VSR",
    )
    parser.add_argument(
        "--denoise-only",
        action="store_true",
        help="Keep the source resolution and run only the GPU denoise stage",
    )
    parser.add_argument("--force", action="store_true", help="Replace an existing output")
    parser.add_argument("--check", action="store_true", help="Check imports and GPU only")
    return parser


def even_scaled(value: int, scale: float) -> int:
    return int(round(value * scale / 2.0) * 2)


def launch_isolated_worker() -> int:
    env = os.environ.copy()
    env[WORKER_ENV] = "1"
    command = [sys.executable, "-u", str(Path(__file__).resolve()), *sys.argv[1:]]
    return subprocess.run(command, env=env, check=False).returncode


def environment_report(torch, VideoSuperRes) -> dict[str, object]:
    return {
        "python": sys.version.split()[0],
        "torch": torch.__version__,
        "torch_cuda": torch.version.cuda,
        "cuda_available": torch.cuda.is_available(),
        "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
        "driver_requirement": "Windows NVIDIA driver >= 570.65",
        "nvvfx_quality_default": int(VideoSuperRes.QualityLevel.HIGHBITRATE_HIGH),
    }


def run_check(torch, VideoSuperRes, device: int) -> None:
    report = environment_report(torch, VideoSuperRes)
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is unavailable in the current H3 Python environment")
    if device >= torch.cuda.device_count():
        raise RuntimeError(f"CUDA device {device} does not exist")

    sample = torch.zeros((3, 64, 64), device=f"cuda:{device}", dtype=torch.float32)
    effect = VideoSuperRes(
        quality=VideoSuperRes.QualityLevel.HIGHBITRATE_HIGH,
        device=device,
    )
    _EFFECT_GUARD.append(effect)
    effect.output_width = 96
    effect.output_height = 96
    started = time.perf_counter()
    effect.load()
    result = torch.from_dlpack(effect.run(sample).image).clone()
    torch.cuda.synchronize(device)
    report["smoke_shape"] = list(result.shape)
    report["smoke_seconds"] = round(time.perf_counter() - started, 4)
    print(json.dumps(report, ensure_ascii=False, indent=2), flush=True)


def gpu_denoise_frame(tensor, preset: str, state: dict, torch):
    if preset == "none":
        return tensor

    import torch.nn.functional as functional

    settings = DENOISE_PRESETS[preset]
    image = tensor.unsqueeze(0)
    base = torch.tensor(
        [1.0, 4.0, 6.0, 4.0, 1.0], device=image.device, dtype=image.dtype
    ).div_(16.0)
    kernel_x = base.view(1, 1, 1, 5).repeat(3, 1, 1, 1)
    kernel_y = base.view(1, 1, 5, 1).repeat(3, 1, 1, 1)
    blurred = functional.conv2d(
        functional.pad(image, (2, 2, 0, 0), mode="replicate"),
        kernel_x,
        groups=3,
    )
    blurred = functional.conv2d(
        functional.pad(blurred, (0, 0, 2, 2), mode="replicate"),
        kernel_y,
        groups=3,
    )

    luma = (
        image[:, 0:1].mul(0.2126)
        + image[:, 1:2].mul(0.7152)
        + image[:, 2:3].mul(0.0722)
    )
    grad_x = functional.pad(
        (luma[:, :, :, 1:] - luma[:, :, :, :-1]).abs(),
        (0, 1, 0, 0),
        mode="replicate",
    )
    grad_y = functional.pad(
        (luma[:, :, 1:, :] - luma[:, :, :-1, :]).abs(),
        (0, 0, 0, 1),
        mode="replicate",
    )
    gradient = torch.maximum(grad_x, grad_y)
    flat_mask = torch.exp(
        -torch.square(gradient / settings["edge_threshold"])
    )
    spatial_weight = flat_mask.mul(settings["spatial_strength"])
    cleaned = image + (blurred - image) * spatial_weight

    previous = state.get("previous")
    if previous is not None:
        difference = (cleaned - previous).abs().mean(dim=1, keepdim=True)
        stable_mask = flat_mask * torch.exp(
            -torch.square(difference / settings["temporal_threshold"])
        )
        temporal_weight = stable_mask.mul(settings["temporal_strength"])
        cleaned = cleaned + (previous - cleaned) * temporal_weight

    state["previous"] = cleaned.detach()
    return cleaned.squeeze(0)


def process_video(args, av, torch, VideoSuperRes) -> None:
    if args.input is None or args.output is None:
        raise ValueError("input and output are required unless --check is used")

    source = args.input.resolve()
    destination = args.output.resolve()
    if not source.is_file():
        raise FileNotFoundError(f"Input video not found: {source}")
    if source == destination:
        raise ValueError("Input and output paths must be different")
    if destination.exists() and not args.force:
        raise FileExistsError(f"Output already exists (use --force): {destination}")
    if not args.denoise_only and args.scale <= 1.0:
        raise ValueError("--scale must be greater than 1.0")
    if args.denoise_only and args.denoise == "none":
        raise ValueError("--denoise-only requires a non-none --denoise preset")
    if not args.denoise_only and (
        args.width <= 0 or args.height <= 0 or args.width % 2 or args.height % 2
    ):
        raise ValueError("Final width and height must be positive even numbers")

    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.exists():
        destination.unlink()

    input_container = None
    output_container = None
    try:
        input_container = av.open(str(source))
        if not input_container.streams.video:
            raise RuntimeError("Input contains no video stream")

        input_video = input_container.streams.video[0]
        input_width = int(input_video.codec_context.width)
        input_height = int(input_video.codec_context.height)
        if args.denoise_only:
            output_width = input_width
            output_height = input_height
            sr_width = input_width
            sr_height = input_height
            crop_left = crop_top = 0
            crop_right = output_width
            crop_bottom = output_height
            quality_name = None
            quality = None
        else:
            output_width = args.width
            output_height = args.height
            sr_width = even_scaled(input_width, args.scale)
            sr_height = even_scaled(input_height, args.scale)
            if output_width > sr_width or output_height > sr_height:
                raise ValueError(
                    f"Final {output_width}x{output_height} does not fit inside "
                    f"{args.scale:g}x Maxine output {sr_width}x{sr_height}"
                )

            crop_left = (sr_width - output_width) // 2
            crop_top = (sr_height - output_height) // 2
            crop_right = crop_left + output_width
            crop_bottom = crop_top + output_height
            if crop_left % 2 or crop_top % 2:
                raise ValueError("Centered crop offsets must be even for clean YUV encoding")

            quality_name = args.quality.upper().replace("-", "_")
            try:
                quality = getattr(VideoSuperRes.QualityLevel, quality_name)
            except AttributeError as exc:
                choices = ", ".join(member.name for member in VideoSuperRes.QualityLevel)
                raise ValueError(
                    f"Unknown quality {args.quality!r}; choose one of: {choices}"
                ) from exc

        fps = input_video.average_rate or input_video.base_rate or 24
        output_container = av.open(str(destination), mode="w")
        output_video = output_container.add_stream(args.encoder, rate=fps)
        output_video.width = output_width
        output_video.height = output_height
        output_video.pix_fmt = "yuv420p"
        output_video.bit_rate = args.bitrate
        if args.encoder.endswith("_nvenc"):
            output_video.options = {
                "preset": "p5",
                "tune": "hq",
                "rc": "vbr",
                "cq": str(args.cq),
            }
        else:
            output_video.options = {"preset": "medium", "crf": str(args.cq)}

        audio_map = {}
        for input_audio in input_container.streams.audio:
            audio_map[input_audio.index] = output_container.add_stream_from_template(input_audio)

        print(
            json.dumps(
                {
                    "mode": "denoise-only" if args.denoise_only else "maxine-vsr",
                    "input": str(source),
                    "input_size": [input_width, input_height],
                    "maxine_size": None if args.denoise_only else [sr_width, sr_height],
                    "crop": None
                    if args.denoise_only
                    else [crop_left, crop_top, crop_right, crop_bottom],
                    "output": str(destination),
                    "output_size": [output_width, output_height],
                    "quality": quality_name,
                    "encoder": args.encoder,
                    "denoise": args.denoise,
                    "denoise_settings": DENOISE_PRESETS.get(args.denoise),
                    "audio_streams_copied": len(audio_map),
                },
                ensure_ascii=False,
                indent=2,
            ),
            flush=True,
        )

        effect = None
        if not args.denoise_only:
            effect = VideoSuperRes(quality=quality, device=args.device)
            _EFFECT_GUARD.append(effect)
            effect.output_width = sr_width
            effect.output_height = sr_height
            load_started = time.perf_counter()
            effect.load()
            print(f"Maxine loaded in {time.perf_counter() - load_started:.3f}s", flush=True)

        frame_count = 0
        process_started = time.perf_counter()
        denoise_state = {}

        def enhance_and_encode(frame) -> None:
            nonlocal frame_count
            rgb = frame.to_ndarray(format="rgb24")
            tensor = (
                torch.from_numpy(rgb)
                .to(device=f"cuda:{args.device}", dtype=torch.float32)
                .permute(2, 0, 1)
                .contiguous()
                .div_(255.0)
            )
            tensor = gpu_denoise_frame(
                tensor, args.denoise, denoise_state, torch
            )
            if args.denoise_only:
                enhanced = tensor
            else:
                enhanced = torch.from_dlpack(effect.run(tensor).image).clone()
                enhanced = enhanced[:, crop_top:crop_bottom, crop_left:crop_right]
            enhanced = (
                enhanced.clamp_(0.0, 1.0)
                .mul_(255.0)
                .round_()
                .to(dtype=torch.uint8)
                .permute(1, 2, 0)
                .contiguous()
                .cpu()
                .numpy()
            )
            output_frame = av.VideoFrame.from_ndarray(enhanced, format="rgb24")
            output_frame.pts = frame.pts
            output_frame.time_base = frame.time_base
            for output_packet in output_video.encode(output_frame):
                output_container.mux(output_packet)

            frame_count += 1
            if frame_count == 1 or frame_count % 24 == 0:
                elapsed = time.perf_counter() - process_started
                print(
                    f"frames={frame_count} elapsed={elapsed:.2f}s "
                    f"throughput={frame_count / max(elapsed, 1e-6):.2f}fps",
                    flush=True,
                )

        streams = [input_video, *input_container.streams.audio]
        for packet in input_container.demux(*streams):
            if packet.stream.index in audio_map:
                if packet.dts is not None:
                    packet.stream = audio_map[packet.stream.index]
                    output_container.mux(packet)
                continue
            if packet.stream.index != input_video.index:
                continue

            for frame in packet.decode():
                enhance_and_encode(frame)

        for output_packet in output_video.encode():
            output_container.mux(output_packet)

        output_container.close()
        output_container = None
        input_container.close()
        input_container = None
        elapsed = time.perf_counter() - process_started
        print(
            json.dumps(
                {
                    "status": "ok",
                    "frames": frame_count,
                    "seconds": round(elapsed, 3),
                    "average_fps": round(frame_count / max(elapsed, 1e-6), 3),
                    "output_bytes": destination.stat().st_size,
                },
                ensure_ascii=False,
                indent=2,
            ),
            flush=True,
        )
    except Exception:
        if output_container is not None:
            try:
                output_container.close()
            except Exception:
                pass
        if input_container is not None:
            try:
                input_container.close()
            except Exception:
                pass
        if destination.exists():
            try:
                destination.unlink()
            except OSError:
                pass
        raise


def worker_main() -> int:
    args = build_parser().parse_args()
    import av
    import torch
    from nvvfx import VideoSuperRes

    if args.check:
        run_check(torch, VideoSuperRes, args.device)
    else:
        process_video(args, av, torch, VideoSuperRes)
    return 0


def main() -> int:
    if os.environ.get(WORKER_ENV) != "1":
        return launch_isolated_worker()

    exit_code = 0
    try:
        exit_code = worker_main()
    except Exception:
        traceback.print_exc()
        exit_code = 1

    # nvidia-vfx 0.1.0.1 can block during Windows interpreter teardown.
    # PyAV files are closed above, so a hard worker exit is safe and keeps the
    # caller/ComfyUI process isolated from that teardown issue.
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(exit_code)


if __name__ == "__main__":
    raise SystemExit(main())
