# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Unit tests for the disk-backed video gallery: MP4 + JSON-sidecar round-trips,
listing order, safe id handling, orphan-pair skipping, and delete/clear."""

from __future__ import annotations

import json
import os
from pathlib import Path

import core.inference.video_gallery as gallery


import pytest


@pytest.fixture(autouse = True)
def _tmp_gallery(monkeypatch, tmp_path):
    # Point the gallery at a throwaway root instead of ~/.unsloth/studio.
    monkeypatch.setattr(gallery, "studio_root", lambda: tmp_path)


def _mp4(tag = b"\x00\x00\x00\x18ftypmp42"):
    # Not a real container; the gallery treats the bytes as opaque payload.
    return tag


def _meta(**over):
    base = {
        "prompt": "a sloth surfing",
        "negative_prompt": None,
        "width": 1024,
        "height": 576,
        "num_frames": 49,
        "fps": 24,
        "duration_s": 2.0,
        "steps": 30,
        "guidance": 6.0,
        "seed": 7,
        "model": "unsloth/some-video-model",
        "created_at": 100.0,
    }
    base.update(over)
    return base


def test_save_writes_pair_and_round_trips():
    record = gallery.save(_mp4(), _meta())
    assert record["id"] and record["url"].endswith(f"{record['id']}/file")

    # Both files of the pair exist: the mp4 payload and the json recipe sidecar.
    directory = gallery.gallery_dir()
    assert (directory / f"{record['id']}.mp4").is_file()
    sidecar = directory / f"{record['id']}.json"
    assert json.loads(sidecar.read_text(encoding = "utf-8"))["prompt"] == "a sloth surfing"

    listed = gallery.list_videos()
    assert len(listed) == 1
    assert listed[0]["prompt"] == "a sloth surfing" and listed[0]["seed"] == 7
    # Meta fields survive the sidecar round-trip untouched.
    assert listed[0]["num_frames"] == 49 and listed[0]["model"] == "unsloth/some-video-model"


def test_url_shape():
    record = gallery.save(_mp4(), _meta())
    assert record["url"] == f"/api/inference/video/gallery/{record['id']}/file"


def _save_with_mtime(prompt: str, t: float) -> dict:
    record = gallery.save(_mp4(), _meta(prompt = prompt, created_at = t))
    # Listing orders by mp4 mtime; set it explicitly so a tight test loop can't tie it.
    os.utime(gallery.gallery_dir() / f"{record['id']}.mp4", (t, t))
    return record


def test_list_is_newest_first():
    old = _save_with_mtime("old", 100.0)
    new = _save_with_mtime("new", 200.0)
    assert [r["id"] for r in gallery.list_videos()] == [new["id"], old["id"]]


def test_list_paginates_with_limit_offset():
    # 5 videos, newest (t=4) first.
    for i in range(5):
        _save_with_mtime(f"p{i}", float(i))
    page1 = gallery.list_videos(limit = 2, offset = 0)
    page2 = gallery.list_videos(limit = 2, offset = 2)
    assert [r["prompt"] for r in page1] == ["p4", "p3"]
    assert [r["prompt"] for r in page2] == ["p2", "p1"]
    # limit=None still returns everything from the offset.
    assert len(gallery.list_videos()) == 5
    assert len(gallery.list_videos(offset = 4)) == 1


def test_video_path_rejects_unsafe_ids():
    # Traversal / bad chars / absolute paths never resolve to a path.
    assert gallery.video_path("../../etc/passwd") is None
    assert gallery.video_path("/etc/passwd") is None
    assert gallery.video_path("a/b") is None
    assert gallery.video_path("missing") is None


def test_video_path_returns_mp4_for_saved_id():
    record = gallery.save(_mp4(), _meta())
    path = gallery.video_path(record["id"])
    assert path is not None and path.name == f"{record['id']}.mp4"


def test_owned_video_path_serves_only_owned_clips():
    # A hand-dropped orphan MP4 resolves via video_path (safe stem, on disk) but must NOT be
    # served: owned_video_path applies the same sidecar check as delete/clear, so the serve and
    # export routes can't stream/transcode a clip the listing hides.
    orphan = gallery.gallery_dir() / "recording.mp4"
    orphan.write_bytes(_mp4())
    assert gallery.video_path("recording") is not None  # resolvable...
    assert gallery.owned_video_path("recording") is None  # ...but not ours to serve

    ours = gallery.save(_mp4(), _meta(prompt = "ours"))
    assert gallery.owned_video_path(ours["id"]) is not None
    assert gallery.owned_video_path("../../etc/passwd") is None
    assert gallery.owned_video_path("missing") is None


def test_transcode_refuses_orphan_mp4():
    # Export starts from the same resolver as /file, so a guessed stem for an orphan MP4 (no
    # readable sidecar) must not be re-encoded out either.
    orphan = gallery.gallery_dir() / "recording.mp4"
    orphan.write_bytes(_real_mp4_bytes())
    assert gallery.transcode("recording", "gif") is None
    assert gallery.transcode("recording", "webm") is None


def test_delete_removes_both_files():
    record = gallery.save(_mp4(), _meta(prompt = "a"))
    gallery.save(_mp4(), _meta(prompt = "b"))
    directory = gallery.gallery_dir()
    assert gallery.delete(record["id"]) is True
    # Both halves of the pair are gone.
    assert not (directory / f"{record['id']}.mp4").exists()
    assert not (directory / f"{record['id']}.json").exists()
    assert gallery.delete(record["id"]) is False  # already gone
    assert len(gallery.list_videos()) == 1


def test_delete_keeps_sidecar_listable_when_mp4_unlink_fails(monkeypatch):
    # delete() must remove the MP4 FIRST: list_videos globs *.mp4 but requires a readable sidecar,
    # so if the sidecar were dropped first and the mp4 unlink then failed (a Windows lock from a
    # concurrent stream/transcode), the still-present mp4 would vanish from the gallery with no way
    # to retry. Simulate the mp4 unlink failing and assert the video stays listable (sidecar kept).
    record = gallery.save(_mp4(), _meta(prompt = "keep"))
    directory = gallery.gallery_dir()
    mp4 = directory / f"{record['id']}.mp4"
    sidecar = directory / f"{record['id']}.json"

    real_unlink = Path.unlink

    def _fail_on_mp4(self, *a, **k):
        if self.suffix == ".mp4":
            raise PermissionError("mp4 locked")
        return real_unlink(self, *a, **k)

    # Patch Path.unlink (what delete() actually calls) rather than os.unlink: on Python 3.10
    # Path.unlink dispatches through a cached _accessor bound to os.unlink at import, so patching
    # os.unlink there has no effect and the mp4 delete would wrongly succeed. Scope it to its own
    # context so undoing it does NOT also revert the autouse fixture's studio_root redirect (both
    # share the function-scoped monkeypatch); otherwise list_videos below would read the real home.
    with pytest.MonkeyPatch.context() as m:
        m.setattr(Path, "unlink", _fail_on_mp4)
        assert gallery.delete(record["id"]) is False  # mp4 unlink failed
    # The sidecar was NOT dropped, so the record is still listable and the user can retry.
    assert sidecar.exists() and mp4.exists()
    assert [r["prompt"] for r in gallery.list_videos()] == ["keep"]
    assert gallery.delete(record["id"]) is True  # retry now succeeds


def test_clear_returns_count():
    gallery.save(_mp4(), _meta(prompt = "a"))
    gallery.save(_mp4(), _meta(prompt = "b"))
    assert gallery.clear() == 2
    assert gallery.list_videos() == []
    # No stray sidecars left behind after a clear.
    assert list(gallery.gallery_dir().glob("*.json")) == []


def test_clear_preserves_orphan_mp4():
    # An orphan / foreign MP4 (no readable sidecar) is invisible to list_videos; clear must not
    # destroy it while removing the owned pair.
    foreign = gallery.gallery_dir() / "recording.mp4"
    foreign.write_bytes(_mp4())
    gallery.save(_mp4(), _meta(prompt = "ours"))
    assert gallery.clear() == 1
    assert foreign.exists()
    assert gallery.list_videos() == []


def test_delete_ignores_orphan_mp4():
    # A per-id delete must refuse an MP4 we do not own (no readable sidecar).
    foreign = gallery.gallery_dir() / "recording.mp4"
    foreign.write_bytes(_mp4())
    assert gallery.delete("recording") is False
    assert foreign.exists()


def test_list_skips_orphan_mp4_without_sidecar():
    # An MP4 with no readable json sidecar (a hand-dropped file) is not a record.
    orphan = gallery.gallery_dir() / "orphan.mp4"
    orphan.write_bytes(_mp4())
    gallery.save(_mp4(), _meta(prompt = "ours"))
    listed = gallery.list_videos()
    assert [r["prompt"] for r in listed] == ["ours"]


def test_list_skips_orphan_sidecar_without_mp4():
    # A json sidecar with no MP4 alongside it is never surfaced (listing globs mp4s).
    orphan = gallery.gallery_dir() / "lonely.json"
    orphan.write_text(json.dumps(_meta(prompt = "no video")), encoding = "utf-8")
    gallery.save(_mp4(), _meta(prompt = "ours"))
    listed = gallery.list_videos()
    assert [r["prompt"] for r in listed] == ["ours"]


def test_orphan_mp4_in_window_does_not_drop_valid_videos():
    # An orphan MP4 sorting INTO the requested page must not consume a window slot
    # and drop a valid video that sorts after it: paging is over readable records.
    _save_with_mtime("p2", 100.0)
    orphan = gallery.gallery_dir() / "zzz_orphan.mp4"
    orphan.write_bytes(_mp4())  # newest by mtime (set below), sorts first
    os.utime(orphan, (300.0, 300.0))
    _save_with_mtime("p1", 200.0)
    # First page of 2 must still return both real videos, not [p1] (orphan eating a slot).
    page1 = gallery.list_videos(limit = 2, offset = 0)
    assert [r["prompt"] for r in page1] == ["p1", "p2"]


def test_list_skips_corrupt_sidecar():
    # A sidecar that is not valid JSON is treated as a foreign/orphan mp4 and skipped.
    directory = gallery.gallery_dir()
    (directory / "broken.mp4").write_bytes(_mp4())
    (directory / "broken.json").write_text("{not json", encoding = "utf-8")
    gallery.save(_mp4(), _meta(prompt = "ours"))
    listed = gallery.list_videos()
    assert [r["prompt"] for r in listed] == ["ours"]


def test_list_skips_invalid_utf8_sidecar():
    # Invalid UTF-8 raises UnicodeDecodeError, which is not an OSError: one corrupt sidecar must
    # be skipped like any other, not 500 the whole gallery listing.
    directory = gallery.gallery_dir()
    (directory / "badbytes.mp4").write_bytes(_mp4())
    (directory / "badbytes.json").write_bytes(b"\xff\xfe{}")
    gallery.save(_mp4(), _meta(prompt = "ours"))
    assert [r["prompt"] for r in gallery.list_videos()] == ["ours"]


def test_clear_preserves_mp4_with_present_but_invalid_sidecar():
    # A hand-dropped MP4 whose sidecar parses but lacks the required recipe keys (e.g. "{}") is
    # hidden by list_videos, so clear must not destroy it while removing the owned pair.
    directory = gallery.gallery_dir()
    (directory / "foreign.mp4").write_bytes(_mp4())
    (directory / "foreign.json").write_text("{}", encoding = "utf-8")
    gallery.save(_mp4(), _meta(prompt = "ours"))
    assert gallery.clear() == 1
    assert (directory / "foreign.mp4").exists()


def test_delete_refuses_mp4_with_present_but_invalid_sidecar():
    # A per-id delete of an MP4 whose sidecar parses but is missing required keys must refuse it:
    # the gallery never surfaced it, so a guessed id must not destroy it.
    directory = gallery.gallery_dir()
    (directory / "foreign.mp4").write_bytes(_mp4())
    (directory / "foreign.json").write_text(
        json.dumps({"prompt": "x"}), encoding = "utf-8"
    )  # partial sidecar (no width/seed/...)
    assert gallery.delete("foreign") is False
    assert (directory / "foreign.mp4").exists()


def test_valid_callback_paginates_over_accepted_records():
    # ``valid`` must filter before pagination, so offset/limit/has_more count over accepted records;
    # else a leading bad record returns a short page with more remaining and stalls scroll.
    _save_with_mtime("BAD", 300.0)  # newest, sorts first
    _save_with_mtime("g1", 200.0)
    _save_with_mtime("g2", 100.0)

    def _valid(rec):
        return rec.get("prompt") != "BAD"

    page = gallery.list_videos(limit = 2, offset = 0, valid = _valid)
    assert [r["prompt"] for r in page] == ["g1", "g2"]
    assert len(gallery.list_videos(limit = 3, offset = 0, valid = _valid)) == 2


def test_valid_callback_leading_bad_records_do_not_stall_at_offset_zero():
    # Every record in the first window is schema-invalid: the pager must look past them and return
    # the good record so has_more is False and the client advances off offset 0.
    for i in range(3):
        _save_with_mtime(f"BAD{i}", 300.0 - i)
    _save_with_mtime("good", 10.0)

    def _valid(rec):
        return not str(rec.get("prompt", "")).startswith("BAD")

    records = gallery.list_videos(limit = 2, offset = 0, valid = _valid)
    assert [r["prompt"] for r in records] == ["good"]


def test_save_leaves_no_orphan_mp4_when_sidecar_publish_fails(monkeypatch):
    # If the sidecar (the pair's commit marker) fails to publish, the MP4 must not be left as an
    # invisible orphan. Fail the second os.replace and assert nothing is stranded.
    real_replace = gallery.os.replace
    calls = {"n": 0}

    def _replace(src, dst, *a, **k):
        calls["n"] += 1
        if calls["n"] == 2:  # the sidecar publish
            raise OSError("simulated sidecar failure")
        return real_replace(src, dst, *a, **k)

    monkeypatch.setattr(gallery.os, "replace", _replace)
    with pytest.raises(OSError, match = "simulated sidecar failure"):
        gallery.save(_mp4(), _meta())
    # No mp4, no sidecar, no temp files -- the whole record was rolled back.
    assert list(gallery.gallery_dir().iterdir()) == []
    assert gallery.list_videos() == []


def _real_mp4_bytes(
    frames: int = 8,
    size: int = 32,
    rate: int = 8,
) -> bytes:
    # A real (tiny) MP4 for the transcode tests: flat-color frames encoded with mpeg4 (bundled in
    # every PyAV build, unlike libx264). The GIF cap tests ask for more/larger frames.
    av = pytest.importorskip("av")
    np = pytest.importorskip("numpy")
    import io

    buf = io.BytesIO()
    with av.open(buf, "w", format = "mp4") as out:
        stream = out.add_stream("mpeg4", rate = rate)
        stream.width = size
        stream.height = size
        stream.pix_fmt = "yuv420p"
        for i in range(frames):
            frame = av.VideoFrame.from_ndarray(
                np.full((size, size, 3), (i * 30) % 256, dtype = np.uint8), format = "rgb24"
            )
            for packet in stream.encode(frame):
                out.mux(packet)
        for packet in stream.encode():
            out.mux(packet)
    return buf.getvalue()


def test_transcode_gif_and_webm_produce_real_containers():
    record = gallery.save(_real_mp4_bytes(), _meta())
    gif = gallery.transcode(record["id"], "gif")
    assert gif is not None and gif.startswith(b"GIF8")
    webm = gallery.transcode(record["id"], "webm")
    # EBML magic: WebM is a Matroska container.
    assert webm is not None and webm[:4] == b"\x1a\x45\xdf\xa3"


def test_transcode_unknown_id_and_bad_format():
    assert gallery.transcode("does-not-exist", "gif") is None
    record = gallery.save(_real_mp4_bytes(), _meta())
    with pytest.raises(ValueError):
        gallery.transcode(record["id"], "avi")


def test_gif_export_bounds_frames_and_edge(monkeypatch):
    """Every kept frame is held as a paletted image before the encoder runs, so an unbounded walk
    is a memory bomb: the generate request allows 2048x2048 for 1024 frames, and at the 12 fps
    target the step is 1, which is over 4 GB of frames plus the GIF buffer. Cap both axes."""
    import io as _io

    from core.inference import video_gallery as vg

    Image = pytest.importorskip("PIL.Image")
    monkeypatch.setattr(vg, "_GIF_MAX_EDGE", 16)
    monkeypatch.setattr(vg, "_GIF_MAX_FRAMES", 4)

    record = gallery.save(_real_mp4_bytes(frames = 24, size = 64, rate = 12), _meta())
    gif = vg._transcode_gif(vg.gallery_dir() / f"{record['id']}.mp4")

    assert gif.startswith(b"GIF8")
    with Image.open(_io.BytesIO(gif)) as im:
        assert max(im.size) <= 16, im.size
        frames = 1
        try:
            while True:
                im.seek(im.tell() + 1)
                frames += 1
        except EOFError:
            pass
    assert frames <= 4, frames