import os

import av
import numpy as np
import pytest

from modules.output.save import nodes


def _clip(n=30):
    rng = np.random.default_rng(0)
    return rng.integers(0, 255, (n, 64, 96, 3), dtype=np.uint8)


def test_cpu_encode_writes_a_playable_h264_mp4_with_sound_and_metadata(tmp_path):
    import torch
    path = str(tmp_path / "a.mp4")
    audio = {"waveform": torch.zeros(1, 2, 44100 * 2), "sample_rate": 44100}
    used = nodes.encode(path, _clip(), 24.0, audio, {"prompt": {"1": "x"}}, encoder="cpu")
    assert "libx264" in used
    with av.open(path) as f:
        v = f.streams.video[0]
        assert v.codec_context.name == "h264" and v.frames == 30 and f.streams.audio
        assert "prompt" in f.metadata
    head = open(path, "rb").read(4096)
    assert head.find(b"moov") != -1 and (head.find(b"mdat") == -1 or head.find(b"moov") < head.find(b"mdat")), "index first: seekable while streaming"


def test_auto_falls_back_to_the_cpu_when_there_is_no_gpu_encoder(tmp_path, monkeypatch):
    monkeypatch.setattr(nodes, "NVENC", {"h264": ("no_such_encoder_x", {})})
    monkeypatch.setattr(av, "codecs_available", av.codecs_available | {"no_such_encoder_x"})
    path = str(tmp_path / "b.mp4")
    assert "libx264" in nodes.encode(path, _clip(10), 24.0, encoder="auto")


def test_gpu_choice_still_saves_on_the_cpu_when_the_encoder_is_missing(tmp_path, monkeypatch):
    monkeypatch.setattr(nodes, "NVENC", {"h264": ("no_such_encoder_x", {})})
    path = str(tmp_path / "c.mp4")
    assert "libx264" in nodes.encode(path, _clip(5), 24.0, encoder="gpu")
    assert os.path.getsize(path) > 0


def test_a_refused_gpu_open_leaves_no_half_file_and_the_cpu_saves_it(tmp_path, monkeypatch):
    # a codec that exists but refuses to open with these options, as NVENC does on a driver mismatch
    monkeypatch.setattr(nodes, "NVENC", {"h264": ("libx264", {"preset": "no-such-preset-x", "crf": "abc"})})
    alerts = []
    monkeypatch.setattr(nodes.log, "once", lambda key, level, who, msg: alerts.append(msg))
    path = str(tmp_path / "e.mp4")
    assert "libx264" in nodes.encode(path, _clip(5), 24.0, encoder="gpu")
    assert os.path.getsize(path) > 0
    assert alerts and "Saved as H.264 on the CPU" in alerts[0]


def test_h265_is_tagged_hvc1_so_a_mac_browser_plays_it(tmp_path, monkeypatch):
    monkeypatch.setattr(nodes, "NVENC", {"h265": ("libx265", {"crf": "23"})})       # the software HEVC encoder stands in for NVENC here
    if "libx265" not in av.codecs_available:
        pytest.skip("no libx265 in this PyAV build")
    path = str(tmp_path / "d.mp4")
    assert "H265" in nodes.encode(path, _clip(10), 24.0, encoder="gpu", codec="h265")
    with av.open(path) as f:
        assert f.streams.video[0].codec_context.name == "hevc" and f.streams.video[0].codec_tag == "hvc1"


def test_nan_sound_is_dropped_and_the_video_still_saves(tmp_path):
    import torch
    path = str(tmp_path / "n.mp4")
    wave = torch.zeros(1, 2, 44100 * 2)
    wave[0, 0, 5] = float("nan")
    nodes.encode(path, _clip(), 24.0, {"waveform": wave, "sample_rate": 44100}, None, encoder="cpu")
    with av.open(path) as f:
        assert f.streams.video[0].frames == 30 and not f.streams.audio


def test_quality_sets_the_constant_quality_number_and_high_is_the_default():
    for name, key in (("libx264", "crf"), ("h264_nvenc", "cq")):
        base = (name, {key: "23"})
        assert nodes._with_quality(base, "standard")[1][key] == "23"
        assert int(nodes._with_quality(base, "high")[1][key]) < 23
        assert int(nodes._with_quality(base, "maximum")[1][key]) < int(nodes._with_quality(base, "high")[1][key])
        assert nodes._with_quality(base, "nonsense")[1][key] == nodes._with_quality(base, "high")[1][key]
    assert nodes.X264[1]["crf"] == "16"


def test_higher_quality_keeps_more_of_the_picture(tmp_path):
    sizes = {}
    for q in ("standard", "high"):
        path = str(tmp_path / f"{q}.mp4")
        nodes.encode(path, _clip(20), 24.0, encoder="cpu", quality=q)
        sizes[q] = os.path.getsize(path)
    assert sizes["high"] > sizes["standard"]


def test_the_last_decoded_picture_is_kept_exactly_for_the_next_shot(tmp_path, monkeypatch):
    from PIL import Image
    monkeypatch.setattr(nodes.folder_paths, "get_temp_directory", lambda: str(tmp_path / "temp"))
    frame = np.random.default_rng(1).integers(0, 256, (12, 16, 3), dtype=np.uint8)
    nodes.keep_last_frame("FunUI_00007_.mp4", frame)
    name = tmp_path / "temp" / nodes.last_frame_name("FunUI_00007_.mp4")
    assert name.name == "funui_last_FunUI_00007_.png"
    assert np.array_equal(np.array(Image.open(name).convert("RGB")), frame)           # lossless: every value as the sampler made it
    monkeypatch.setattr(nodes.folder_paths, "get_temp_directory", lambda: (_ for _ in ()).throw(OSError("no temp")))
    nodes.keep_last_frame("x.mp4", frame)           # a failure here never fails the save