#!/usr/bin/env python3
"""GPU-vendor-neutral temporal cleanup and Lanczos video finishing.

This is the safe fallback used when NVIDIA Maxine VSR is unavailable.  It only
depends on PyAV, NumPy, SciPy and Pillow, keeps the source audio stream, and
encodes H.264 with libx264.  It deliberately favors stable geometry over
aggressive invented detail.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import time

import av
import numpy as np
from PIL import Image
from scipy.ndimage import gaussian_filter


PRESETS = {
    "mild": {"sigma": 0.48, "spatial": 0.24, "edge": 0.038, "temporal": 0.08, "motion": 0.030},
    "medium": {"sigma": 0.62, "spatial": 0.36, "edge": 0.048, "temporal": 0.14, "motion": 0.040},
    "finish": {"sigma": 0.56, "spatial": 0.30, "edge": 0.034, "temporal": 0.16, "motion": 0.034},
}


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser(description="Universal H3 video cleanup and Lanczos upscale")
    result.add_argument("input", type=Path)
    result.add_argument("output", type=Path)
    result.add_argument("--denoise", choices=("none", *PRESETS), default="medium")
    result.add_argument("--scale", type=float, default=1.5)
    result.add_argument("--width", type=int)
    result.add_argument("--height", type=int)
    result.add_argument("--denoise-only", action="store_true")
    result.add_argument("--bitrate", type=int, default=20_000_000)
    result.add_argument("--crf", type=int, default=18)
    result.add_argument("--force", action="store_true")
    return result


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


def clean_frame(rgb: np.ndarray, preset: str, state: dict[str, np.ndarray]) -> np.ndarray:
    if preset == "none":
        return rgb
    settings = PRESETS[preset]
    current = rgb.astype(np.float32) / 255.0
    blurred = gaussian_filter(current, sigma=(settings["sigma"], settings["sigma"], 0.0), mode="nearest")
    luma = current[..., 0] * 0.2126 + current[..., 1] * 0.7152 + current[..., 2] * 0.0722
    gradient_x = np.abs(np.diff(luma, axis=1, append=luma[:, -1:]))
    gradient_y = np.abs(np.diff(luma, axis=0, append=luma[-1:, :]))
    flat = np.exp(-np.square(np.maximum(gradient_x, gradient_y) / settings["edge"]))[..., None]
    cleaned = current + (blurred - current) * flat * settings["spatial"]
    previous = state.get("previous")
    if previous is not None:
        difference = np.mean(np.abs(cleaned - previous), axis=2, keepdims=True)
        stable = flat * np.exp(-np.square(difference / settings["motion"]))
        cleaned += (previous - cleaned) * stable * settings["temporal"]
    state["previous"] = cleaned.copy()
    return np.clip(np.rint(cleaned * 255.0), 0, 255).astype(np.uint8)


def process(args: argparse.Namespace) -> None:
    source = args.input.resolve()
    destination = args.output.resolve()
    if not source.is_file():
        raise FileNotFoundError(source)
    if source == destination:
        raise ValueError("input and output must be different files")
    if destination.exists() and not args.force:
        raise FileExistsError(destination)
    if not args.denoise_only and args.scale <= 1.0:
        raise ValueError("scale must be greater than 1")
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.exists():
        destination.unlink()

    input_container = 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, output_height = input_width, input_height
            scaled_width, scaled_height = input_width, input_height
        else:
            scaled_width = even_scaled(input_width, args.scale)
            scaled_height = even_scaled(input_height, args.scale)
            output_width = args.width or scaled_width
            output_height = args.height or scaled_height
            if output_width > scaled_width or output_height > scaled_height:
                raise ValueError("final output does not fit inside the scaled frame")
        crop_left = (scaled_width - output_width) // 2
        crop_top = (scaled_height - output_height) // 2

        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("libx264", rate=fps)
        output_video.width = output_width
        output_video.height = output_height
        output_video.pix_fmt = "yuv420p"
        output_video.bit_rate = args.bitrate
        output_video.options = {"preset": "medium", "crf": str(args.crf)}
        audio_map = {
            stream.index: output_container.add_stream_from_template(stream)
            for stream in input_container.streams.audio
        }
        print(json.dumps({
            "mode": "universal-finish",
            "input": str(source),
            "input_size": [input_width, input_height],
            "scaled_size": [scaled_width, scaled_height],
            "output_size": [output_width, output_height],
            "denoise": args.denoise,
            "scaler": "Pillow Lanczos",
            "encoder": "libx264",
            "audio_streams_copied": len(audio_map),
        }, ensure_ascii=False, indent=2), flush=True)

        state: dict[str, np.ndarray] = {}
        started = time.perf_counter()
        frames = 0
        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():
                rgb = clean_frame(frame.to_ndarray(format="rgb24"), args.denoise, state)
                if not args.denoise_only:
                    image = Image.fromarray(rgb, "RGB").resize(
                        (scaled_width, scaled_height), Image.Resampling.LANCZOS
                    )
                    if output_width != scaled_width or output_height != scaled_height:
                        image = image.crop((crop_left, crop_top, crop_left + output_width, crop_top + output_height))
                    rgb = np.asarray(image)
                output_frame = av.VideoFrame.from_ndarray(rgb, 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)
                frames += 1
                if frames == 1 or frames % 12 == 0:
                    elapsed = time.perf_counter() - started
                    print(f"frames={frames} elapsed={elapsed:.2f}s throughput={frames / max(elapsed, 1e-6):.2f}fps", flush=True)

        for packet in output_video.encode():
            output_container.mux(packet)
        output_container.close()
        output_container = None
        input_container.close()
        input_container = None
        elapsed = time.perf_counter() - started
        print(json.dumps({
            "status": "ok",
            "frames": frames,
            "seconds": round(elapsed, 3),
            "average_fps": round(frames / 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:
            output_container.close()
        if input_container is not None:
            input_container.close()
        if destination.exists():
            destination.unlink()
        raise


def main() -> int:
    process(parser().parse_args())
    return 0


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