"""
CreateSaveVideo — Single-File Standalone ComfyUI Custom Node
============================================================
Drop this file into:  ComfyUI/custom_nodes/create_save_video.py
"""

import os
import wave
import time
import tempfile
import subprocess
import shutil
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import torch
import folder_paths

# ---------------------------------------------------------------------------
# FFmpeg binary resolution
# ---------------------------------------------------------------------------

def _find_ffmpeg() -> str:
    try:
        import imageio_ffmpeg
        exe = imageio_ffmpeg.get_ffmpeg_exe()
        if exe and os.path.isfile(exe):
            print(f"[CreateSaveVideo] FFmpeg via imageio-ffmpeg: {exe}")
            return exe
    except ImportError:
        pass

    exe = shutil.which("ffmpeg")
    if exe:
        print(f"[CreateSaveVideo] FFmpeg on PATH: {exe}")
        return exe

    from pathlib import Path
    candidates = [
        r"C:\ffmpeg\bin\ffmpeg.exe",
        r"C:\Program Files\ffmpeg\bin\ffmpeg.exe",
        r"C:\Program Files (x86)\ffmpeg\bin\ffmpeg.exe",
        os.path.expanduser(r"~\ffmpeg\bin\ffmpeg.exe"),
    ]
    try:
        comfy_root = Path(folder_paths.base_path)
        for rel in ["ffmpeg/bin/ffmpeg.exe", "ffmpeg/ffmpeg.exe",
                    "../ffmpeg/bin/ffmpeg.exe"]:
            candidates.append(str(comfy_root / rel))
    except Exception:
        pass

    for path in candidates:
        if os.path.isfile(path):
            print(f"[CreateSaveVideo] FFmpeg found at: {path}")
            return path

    raise FileNotFoundError(
        "[CreateSaveVideo] FFmpeg not found!\n"
        "Fix: pip install imageio-ffmpeg"
    )


FFMPEG_EXE = None

# ---------------------------------------------------------------------------
# Codec / container compatibility matrix
# ---------------------------------------------------------------------------
CODEC_MAP = {
    ("mp4",  "h264"):    ("libx264",    ["-movflags", "+faststart"]),
    ("mp4",  "h265"):    ("libx265",    ["-movflags", "+faststart", "-tag:v", "hvc1"]),
    ("mp4",  "av1"):     ("libaom-av1", ["-movflags", "+faststart", "-cpu-used", "4"]),
    ("mp4",  "prores"):  ("prores_ks",  ["-profile:v", "3"]),
    ("mov",  "h264"):    ("libx264",    ["-movflags", "+faststart"]),
    ("mov",  "h265"):    ("libx265",    ["-movflags", "+faststart", "-tag:v", "hvc1"]),
    ("mov",  "prores"):  ("prores_ks",  ["-profile:v", "3"]),
    ("webm", "vp9"):     ("libvpx-vp9", ["-b:v", "0", "-deadline", "realtime"]),
    ("webm", "av1"):     ("libaom-av1", ["-b:v", "0", "-cpu-used", "4"]),
    ("avi",  "h264"):    ("libx264",    []),
    ("mkv",  "h264"):    ("libx264",    []),
    ("mkv",  "h265"):    ("libx265",    []),
    ("mkv",  "vp9"):     ("libvpx-vp9", ["-b:v", "0", "-deadline", "realtime"]),
    ("mkv",  "av1"):     ("libaom-av1", ["-b:v", "0", "-cpu-used", "4"]),
}

LOSSLESS_CODECS = {"prores"}

# ---------------------------------------------------------------------------
# Audio helper — stdlib only, no torchaudio / torchcodec
# ---------------------------------------------------------------------------

def _save_wav_stdlib(path: str, waveform: torch.Tensor, sample_rate: int):
    """
    Write a WAV file using Python's built-in `wave` module.
    No torchaudio, no torchcodec, no external dependencies.

    waveform: (C, T) or (1, C, T) float32, range [-1.0, 1.0]
    """
    # normalise batch dim
    if waveform.dim() == 3:
        waveform = waveform.squeeze(0)          # (1,C,T) → (C,T)

    audio_np = waveform.cpu().numpy()           # (C, T)
    n_channels, n_samples = audio_np.shape

    # float32 [-1,1] → int16
    pcm = (audio_np * 32767.0).clip(-32768, 32767).astype(np.int16)

    # interleave channels: (C, T) → (T, C) → flatten
    interleaved = pcm.T.flatten()

    with wave.open(path, "wb") as wf:
        wf.setnchannels(n_channels)
        wf.setsampwidth(2)           # 16-bit PCM
        wf.setframerate(sample_rate)
        wf.writeframes(interleaved.tobytes())


# ---------------------------------------------------------------------------
# Frame conversion helper
# ---------------------------------------------------------------------------

def _tensor_to_bytes(frame: torch.Tensor) -> bytes:
    return (frame.numpy() * 255.0).clip(0, 255).astype(np.uint8).tobytes()


# ---------------------------------------------------------------------------
# Node
# ---------------------------------------------------------------------------

class CreateSaveVideo:
    CATEGORY     = "video/optimized"
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("filepath",)
    OUTPUT_NODE  = True
    FUNCTION     = "encode_and_save"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "fps": (
                    "FLOAT",
                    {"default": 24.0, "min": 1.0, "max": 120.0,
                     "step": 0.5, "display": "number"},
                ),
                "filename_prefix": ("STRING", {"default": "video/output"}),
                "format": (["mp4", "mov", "mkv", "webm", "avi"], {"default": "mp4"}),
                "codec":  (["h264", "h265", "vp9", "av1", "prores"], {"default": "h264"}),
                "crf": (
                    "INT",
                    {"default": 18, "min": 0, "max": 63, "step": 1,
                     "display": "slider",
                     "tooltip": "Quality — lower = better. Ignored for prores."},
                ),
                "preset": (
                    ["ultrafast", "superfast", "veryfast", "faster",
                     "fast", "medium", "slow", "veryslow"],
                    {"default": "veryfast"},
                ),
                "worker_threads": (
                    "INT",
                    {"default": 4, "min": 1, "max": 32, "step": 1,
                     "display": "number"},
                ),
            },
            "optional": {
                "audio": ("AUDIO",),
                "ffmpeg_path": (
                    "STRING",
                    {"default": "",
                     "tooltip": "Absolute path to ffmpeg.exe — leave blank for auto-detect."},
                ),
                "ffmpeg_extra_args": (
                    "STRING",
                    {"default": "",
                     "tooltip": "Raw extra FFmpeg flags, space-separated (advanced)."},
                ),
            },
        }

    # ------------------------------------------------------------------

    def encode_and_save(
        self,
        images: torch.Tensor,
        fps: float,
        filename_prefix: str,
        format: str,
        codec: str,
        crf: int,
        preset: str,
        worker_threads: int,
        audio=None,
        ffmpeg_path: str = "",
        ffmpeg_extra_args: str = "",
    ):
        global FFMPEG_EXE
        t_start = time.perf_counter()

        # ── resolve ffmpeg binary ────────────────────────────────────────────
        if ffmpeg_path.strip() and os.path.isfile(ffmpeg_path.strip()):
            ffmpeg_exe = ffmpeg_path.strip()
        else:
            if FFMPEG_EXE is None:
                FFMPEG_EXE = _find_ffmpeg()
            ffmpeg_exe = FFMPEG_EXE

        # ── resolve output path ──────────────────────────────────────────────
        output_dir, prefix, _, _, _ = folder_paths.get_save_image_path(
            filename_prefix, folder_paths.get_output_directory()
        )
        os.makedirs(output_dir, exist_ok=True)
        existing = [f for f in os.listdir(output_dir) if f.startswith(prefix)]
        counter  = len(existing) + 1
        filename = f"{prefix}_{counter:05d}.{format}"
        out_path = os.path.join(output_dir, filename)

        # ── frame geometry ───────────────────────────────────────────────────
        n_frames, H, W, C = images.shape
        pix_fmt_in  = "rgb24" if C == 3 else "rgba"
        pix_fmt_out = "yuv422p10le" if codec == "prores" else "yuv420p"

        # ── codec lookup ─────────────────────────────────────────────────────
        vcodec, extra_flags = CODEC_MAP.get(
            (format, codec), ("libx264", ["-movflags", "+faststart"])
        )

        # ── write audio to temp WAV using stdlib only ────────────────────────
        tmp_audio_path = None
        has_audio = False
        if audio is not None:
            try:
                waveform  = audio["waveform"]
                sample_rate = audio["sample_rate"]
                tmp_fd, tmp_audio_path = tempfile.mkstemp(suffix=".wav")
                os.close(tmp_fd)
                _save_wav_stdlib(tmp_audio_path, waveform, sample_rate)
                has_audio = True
                print(f"[CreateSaveVideo] Audio written (stdlib WAV): "
                      f"{waveform.shape} @ {sample_rate}Hz")
            except Exception as exc:
                print(f"[CreateSaveVideo] Audio error — {exc}")
                tmp_audio_path = None
                has_audio = False

        # ── build FFmpeg command ─────────────────────────────────────────────
        # Input 0: raw video frames from stdin pipe
        cmd = [
            ffmpeg_exe, "-y",
            "-f", "rawvideo",
            "-vcodec", "rawvideo",
            "-s", f"{W}x{H}",
            "-pix_fmt", pix_fmt_in,
            "-r", str(fps),
            "-i", "pipe:0",
        ]

        # Input 1: audio WAV (if present)
        if has_audio:
            cmd += ["-i", tmp_audio_path]

        # Video codec + quality
        cmd += ["-vcodec", vcodec]
        if codec not in LOSSLESS_CODECS:
            cmd += ["-crf", str(crf)]
        if codec in ("h264", "h265"):
            cmd += ["-preset", preset]
        cmd += ["-pix_fmt", pix_fmt_out]
        cmd += extra_flags

        # Explicit stream mapping — required when multiple inputs exist
        if has_audio:
            cmd += [
                "-map", "0:v:0",
                "-map", "1:a:0",
                "-c:a", "aac",
                "-b:a", "192k",
                "-shortest",
            ]
        else:
            cmd += ["-map", "0:v:0"]

        if ffmpeg_extra_args.strip():
            cmd += ffmpeg_extra_args.strip().split()

        cmd.append(out_path)

        # ── parallel frame → bytes ────────────────────────────────────────────
        images_cpu = images.contiguous().cpu()

        def convert(i):
            return _tensor_to_bytes(images_cpu[i])

        print(f"[CreateSaveVideo] Converting {n_frames} frames ({W}x{H}) "
              f"with {worker_threads} threads …")
        t_conv = time.perf_counter()
        with ThreadPoolExecutor(max_workers=worker_threads) as pool:
            raw_frames = list(pool.map(convert, range(n_frames)))
        print(f"[CreateSaveVideo] Frame conversion : {time.perf_counter() - t_conv:.2f}s")

        # ── stream into FFmpeg ────────────────────────────────────────────────
        print(f"[CreateSaveVideo] Encoding → {out_path}")
        t_enc = time.perf_counter()

        proc = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        for raw in raw_frames:
            proc.stdin.write(raw)
        proc.stdin.close()

        _, stderr = proc.communicate()
        if proc.returncode != 0:
            raise RuntimeError(
                f"[CreateSaveVideo] FFmpeg failed (rc={proc.returncode}):\n"
                + stderr.decode(errors="replace")
            )

        print(f"[CreateSaveVideo] Encode time  : {time.perf_counter() - t_enc:.2f}s")
        print(f"[CreateSaveVideo] Total time   : {time.perf_counter() - t_start:.2f}s")
        print(f"[CreateSaveVideo] Saved → {out_path}")

        if tmp_audio_path and os.path.exists(tmp_audio_path):
            os.unlink(tmp_audio_path)

        return {"ui": {"video": [filename]}, "result": (out_path,)}


# ---------------------------------------------------------------------------
# ComfyUI registration
# ---------------------------------------------------------------------------

NODE_CLASS_MAPPINGS = {
    "CreateSaveVideo": CreateSaveVideo,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "CreateSaveVideo": "Create & Save Video (Optimized)",
}