tests/test_wan_video.py

from __future__ import annotations

import io
import json
import struct
import sys
import threading
import time
from pathlib import Path

import pytest

from adam.assets import AssetRegistry
from adam.config import ConfigManager
from adam.executor import ToolCancelled, ToolContext, ToolExecutionError
from adam.models import ExecutionPlan, PlanStep
from adam.registry import ToolRegistry, ToolSpec
from adam.tools import wan_video_adapter as adapter
from adam.video_lora import (
    MODEL_FILES, checkpoint_info, dataset_errors, discover_assets, frames_for_seconds,
    generation_plan, settings_for, trainer_root, training_plan,
)


def checkpoint(path: Path, *, base="wan2.1_t2v_1.3B_bf16.safetensors") -> Path:
    header = json.dumps({"__metadata__": {"ss_network_module": "networks.lora_wan", "ss_sd_model_name": base,
                                         "ss_epoch": "2", "ss_network_dim": "16"},
                         "lora_unet_blocks_0.lora_down.weight": {"dtype": "F32", "shape": [1, 1], "data_offsets": [0, 4]}}).encode()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_bytes(struct.pack("<Q", len(header)) + header + bytes(4))
    return path


@pytest.fixture
def workspace(tmp_path):
    external = tmp_path / "trainer"
    needed = [".venv/Scripts/python.exe", "tools/validate_dataset.py", "tools/split_video.py"]
    needed += [f"models/{name}" for name in MODEL_FILES]
    needed += [f"trainer-engine/src/musubi_tuner/{name}.py" for name in ("wan_cache_latents", "wan_cache_text_encoder_outputs", "wan_train_network", "wan_generate_video")]
    for relative in needed:
        path = external / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        path.touch()
    config = ConfigManager(tmp_path)
    config.update({"tool_folders": {"wan_video_trainer": str(external)}})
    (tmp_path / "config" / "tools.json").write_text((Path(__file__).resolve().parents[1] / "config" / "tools.json").read_text(encoding="utf-8"), encoding="utf-8")
    dataset = tmp_path / "clips"
    dataset.mkdir()
    (dataset / "clip.mp4").write_bytes(b"video")
    (dataset / "clip.txt").write_text("subject_token turns toward the camera", encoding="utf-8")
    return tmp_path, external, dataset


def context(root: Path, tool="wan_video_trainer"):
    run = threading.Event()
    run.set()
    return ToolContext(root, "TEST", ToolSpec(tool, "Wan", "test", "Training", "train"), threading.Event(), run,
                       lambda *_args, **_kwargs: None, lambda *_args: None)


def test_external_registration_resolves_without_personal_path(workspace):
    root, external, _ = workspace
    ConfigManager(root).update({"tool_folders": {}})
    (root / "config" / "external_tools.json").write_text(json.dumps({"tools": [{"id": "external_loravideotrainer", "backend": {"root": str(external)}}]}))
    assert trainer_root(root) == external


def test_wan_registry_does_not_enter_image_generation(workspace):
    from adam.generations import generation_tools
    root, _, _ = workspace
    registry = ToolRegistry(root)
    assert "wan_video_generator" not in {tool.id for tool in generation_tools(registry)}
    assert registry.get("wan_video_generator").required_arguments == ("model_path", "prompt")
    assert "image_count" not in registry.get("wan_video_generator").arguments
    assert "pause" in registry.get("wan_video_trainer").capabilities


def test_video_chat_request_never_routes_to_sdxl(workspace):
    from adam.planner import Planner
    root, _, _ = workspace
    config = ConfigManager(root)
    config.update({"provider": "manual"})
    planner = Planner(root, ToolRegistry(root), config)
    for request in ("Train a video LoRA from the clips dataset for 10 epochs", "Train a Wan 2.1 LoRA", "Generate a LoRA video"):
        plan = planner.plan(request)
        assert plan.steps == []
        assert plan.project_name == "Video LoRA"
        assert "sidebar" in plan.summary


@pytest.mark.parametrize("values", [{"learning_rate": float("nan")}, {"rank": 1.2}, {"rank": True}, {"blocks_to_swap": 30}, {"resolution": "1024x1024"}, {"target_frames": "24"}, {"trigger_word": ""}])
def test_training_settings_reject_bad_values(values):
    with pytest.raises(ValueError):
        settings_for("training", values)


@pytest.mark.parametrize("base", ["sdxl.safetensors", "wan2.1_t2v_14B.safetensors", "wan2.1_i2v_1.3B.safetensors"])
def test_checkpoint_rejects_incompatible_model(tmp_path, base):
    with pytest.raises(ValueError, match="incompatible"):
        checkpoint_info(checkpoint(tmp_path / "model.safetensors", base=base))


def test_truncated_checkpoint_rejected(tmp_path):
    path = checkpoint(tmp_path / "model.safetensors")
    path.write_bytes(path.read_bytes()[:-2])
    with pytest.raises(ValueError, match="incomplete"):
        checkpoint_info(path)


def test_dataset_rejects_missing_trigger_and_duplicate_stems(workspace):
    _, _, dataset = workspace
    (dataset / "clip.mov").write_bytes(b"video")
    errors = dataset_errors(dataset, "wrong_token")
    assert any("Duplicate" in error for error in errors)
    assert any("exact trigger" in error for error in errors)


def test_plan_requires_confirmation_and_separate_outputs(workspace):
    root, _, dataset = workspace
    plan = training_plan(root, str(dataset), "Video model", 2, {})
    assert plan.requires_confirmation
    assert "image" not in plan.summary
    with pytest.raises(ValueError, match="already exists"):
        training_plan(root, str(dataset), "Video model", 2, {}, output=str(dataset))
    with pytest.raises(ValueError, match="separate"):
        training_plan(root, str(dataset), "Video model", 2, {}, output=str(dataset / "new"))
    with pytest.raises(ValueError, match="reserved"):
        training_plan(root, str(dataset), "../oops", 2, {})


def test_pipeline_order_isolated_cache_and_model_registration(workspace, monkeypatch):
    root, external, dataset = workspace
    calls = []
    monkeypatch.setattr(adapter, "dataset_errors", lambda *_args, **_kwargs: [])
    output = root / "data" / "video_models" / "test"
    def run(_ctx, args, cwd, label, start, end, **kwargs):
        calls.append(args)
        assert cwd == external
        assert args[0] == str(external / ".venv" / "Scripts" / "python.exe")
        if "--max_train_epochs" in args:
            checkpoint(output / "test.safetensors")
    monkeypatch.setattr(adapter, "run_stage", run)
    result = adapter.train(context(root), str(dataset), "test", 2, str(output))
    assert "validate_dataset.py" in calls[0][1]
    assert "wan_cache_latents.py" in calls[1][1]
    assert "wan_cache_text_encoder_outputs.py" in calls[2][1]
    assert "--max_train_epochs" in calls[3]
    assert result["assets"][1]["trainer"] == "wan_video"
    assert (output / "cache").is_dir()
    assert not (external / "dataset_config.toml").exists()
    assert json.loads((output / "model_info.json").read_text())["status"] == "finished"
    with pytest.raises(ToolExecutionError, match="already exists"):
        adapter.train(context(root), str(dataset), "test", 2, str(output))


def test_failure_stops_pipeline_and_records_run(workspace, monkeypatch):
    root, _, dataset = workspace
    monkeypatch.setattr(adapter, "dataset_errors", lambda *_args, **_kwargs: [])
    calls = []
    def run(*args, **kwargs):
        calls.append(args)
        raise ToolExecutionError("cache failed")
    monkeypatch.setattr(adapter, "run_stage", run)
    output = root / "output"
    with pytest.raises(ToolExecutionError, match="cache failed"):
        adapter.train(context(root), str(dataset), "test", 2, str(output))
    assert len(calls) == 1
    assert json.loads((output / "model_info.json").read_text())["status"] == "failed"


def test_changed_caption_stops_before_cache_stage(workspace, monkeypatch):
    root, _, dataset = workspace
    monkeypatch.setattr(adapter, "dataset_errors", lambda *_args, **_kwargs: [])
    calls = []
    def run(*args, **kwargs):
        calls.append(args)
        (dataset / "clip.txt").write_text("edited", encoding="utf-8")
    monkeypatch.setattr(adapter, "run_stage", run)
    with pytest.raises(ToolExecutionError, match="Dataset changed"):
        adapter.train(context(root), str(dataset), "test", 2, str(root / "output"))
    assert len(calls) == 1


@pytest.mark.parametrize("format,width,height", [("Landscape 832x480", 832, 480), ("Portrait 480x832", 480, 832)])
def test_generation_height_width_order_and_metadata(workspace, monkeypatch, format, width, height):
    root, _, _ = workspace
    model = checkpoint(root / "adapter.safetensors")
    calls = []
    def run(_ctx, args, *_args, **kwargs):
        calls.append(args)
        output = Path(args[args.index("--save_path") + 1])
        (output / "clip.mp4").write_bytes(b"video")
    monkeypatch.setattr(adapter, "run_stage", run)
    monkeypatch.setattr(adapter, "inspect_clip", lambda path: {"width": width, "height": height, "frames": 25, "fps": 12})
    result = adapter.generate(context(root, "wan_video_generator"), str(model), "Test", prompt="subject_token walking", format=format)
    args = calls[0]
    index = args.index("--video_size")
    assert args[index + 1:index + 3] == [str(height), str(width)]
    info = json.loads((Path(result["output_folder"]) / "generation.json").read_text())
    assert info["frames"] == 25
    assert info["status"] == "finished"
    assert info["settings"]["seed"] == 1701
    assert result["assets"][0]["kind"] == "video"


def test_generated_wrong_dimensions_fail(workspace, monkeypatch):
    root, _, _ = workspace
    model = checkpoint(root / "adapter.safetensors")
    def run(_ctx, args, *_args, **kwargs):
        (Path(args[args.index("--save_path") + 1]) / "clip.mp4").write_bytes(b"video")
    monkeypatch.setattr(adapter, "run_stage", run)
    monkeypatch.setattr(adapter, "inspect_clip", lambda path: {"width": 480, "height": 832, "frames": 25, "fps": 12})
    with pytest.raises(ToolExecutionError, match="does not match"):
        adapter.generate(context(root), str(model), prompt="subject_token")
    assert json.loads(next((root / "data" / "video_generations").glob("*/generation.json")).read_text())["status"] == "failed"


def test_video_preflight_and_orion_do_not_use_image_estimates(workspace):
    from adam.training_assistant import append_preflight_summary
    root, _, dataset = workspace
    plan = training_plan(root, str(dataset), "test", 2, {})
    append_preflight_summary(plan, ConfigManager(root))
    assert "no top-level images" not in plan.summary
    assert "Wan video: 1 clips" in plan.summary
    assert plan.orion_review["estimated_high_minutes"] == 0


def test_discovery_labels_external_and_native_models_as_video(workspace):
    root, external, _ = workspace
    checkpoint(external / "output" / "old.safetensors")
    checkpoint(root / "data" / "video_models" / "new" / "new.safetensors")
    assets = AssetRegistry(root)
    discover_assets(assets, root, ConfigManager(root))
    models = [asset for asset in assets.assets if asset.kind == "model"]
    assert len(models) == 2
    assert all(asset.trainer == "wan_video" for asset in models)


def test_real_subprocess_output_is_drained(tmp_path):
    ctx = context(tmp_path)
    logs = []
    ctx.log_callback = logs.append
    adapter.run_stage(ctx, [sys.executable, "-c", "print('worker complete')"], tmp_path, "Worker", 0, 100)
    assert "worker complete" in logs


def test_wan_primary_progress_ignores_model_loading_bars():
    assert adapter._progress_counter(
        "Loading checkpoint shards: 98%|#########8| 819/825", "steps"
    ) is None
    assert adapter._progress_counter(
        "steps: 66%|######6| 106/160 [38:04<19:23, 21.55s/it]", "steps"
    ) == (106, 160)


def test_cancel_terminates_worker(tmp_path):
    ctx = context(tmp_path)
    timer = threading.Timer(0.3, ctx.cancel_event.set)
    timer.start()
    started = time.monotonic()
    try:
        with pytest.raises(ToolCancelled):
            adapter.run_stage(ctx, [sys.executable, "-c", "import time; time.sleep(20)"], tmp_path, "Worker", 0, 100)
    finally:
        timer.cancel()
    assert time.monotonic() - started < 8


def test_cancel_while_paused_resumes_then_terminates_worker(tmp_path, monkeypatch):
    ctx = context(tmp_path)
    transitions = []
    original = adapter.set_process_tree_paused
    def pause(process, paused):
        transitions.append(paused)
        return original(process, paused)
    monkeypatch.setattr(adapter, "set_process_tree_paused", pause)
    timer = threading.Timer(0.6, ctx.cancel_event.set)
    def log(line):
        if "ready" in line:
            ctx.run_event.clear()
            timer.start()
    ctx.log_callback = log
    try:
        with pytest.raises(ToolCancelled):
            adapter.run_stage(ctx, [sys.executable, "-u", "-c", "import time; print('ready'); time.sleep(20)"], tmp_path, "Worker", 0, 100)
    finally:
        timer.cancel()
    assert transitions == [True, False]


def test_frames_follow_wan_rule():
    assert frames_for_seconds(2, 12) == 25
    for fps in (4, 12, 24, 60):
        for duration in (2, 4.5, 15):
            frames = frames_for_seconds(duration, fps)
            assert (frames - 1) % 4 == 0
            assert abs(frames - duration * fps) <= 2


def test_desktop_workspace_smoke_and_settings_roundtrip(workspace, monkeypatch):
    monkeypatch.setenv("QT_QPA_PLATFORM", "offscreen")
    from PySide6.QtWidgets import QApplication
    from PySide6.QtCore import QObject, Signal
    from adam.ui.video_lora import VideoLoraPage
    class Jobs(QObject):
        job_updated = Signal(object)
        active_job = None
    root, external, _ = workspace
    checkpoint(external / "output" / "existing.safetensors")
    app = QApplication.instance() or QApplication([])
    page = VideoLoraPage(root, ToolRegistry(root), Jobs(), AssetRegistry(root), ConfigManager(root))
    page.resize(1120, 760)
    page.show()
    deadline = time.monotonic() + 5
    while page.models.count() != 1 and time.monotonic() < deadline:
        app.processEvents()
        time.sleep(0.005)
    assert page.tabs.count() == 4
    assert page.models.count() == 1
    page.generation.set_values({"duration": 4.5, "format": "Portrait 480x832"})
    page._save()
    saved = ConfigManager(root).get("video_lora_settings")
    assert saved["generation"]["duration"] == 4.5
    assert saved["generation"]["format"] == "Portrait 480x832"

    from PySide6.QtWidgets import QFileDialog, QInputDialog
    monkeypatch.setattr(QFileDialog, "getOpenFileName", lambda *_args, **_kwargs: (str(root / "source.mp4"), ""))
    monkeypatch.setattr(QInputDialog, "getInt", lambda *_args, **_kwargs: (8, True))
    monkeypatch.setattr(QInputDialog, "getText", lambda *_args, **_kwargs: ("Joe_Bartolozzi", True))
    monkeypatch.setattr(QInputDialog, "getMultiLineText", lambda *_args, **_kwargs: ("speaking directly to the camera", True))
    submitted = []
    monkeypatch.setattr(page, "_submit", submitted.append)
    page._split()
    arguments = submitted[0].steps[0].arguments
    assert arguments["trigger_word"] == "Joe_Bartolozzi"
    assert arguments["caption"] == "Joe_Bartolozzi, speaking directly to the camera"
    assert page.training.value("trigger_word") == "Joe_Bartolozzi"
    assert "Joe_Bartolozzi" in page.trigger_hint.text()
    page.close()
    page.deleteLater()
    app.processEvents()


def test_desktop_archive_current_dataset_starts_a_new_one(workspace, monkeypatch):
    monkeypatch.setenv("QT_QPA_PLATFORM", "offscreen")
    from PySide6.QtWidgets import QApplication, QMessageBox, QInputDialog
    from PySide6.QtCore import QObject, Signal
    from adam.ui.video_lora import VideoLoraPage

    class Jobs(QObject):
        job_updated = Signal(object)
        active_job = None

    root, _, dataset = workspace
    app = QApplication.instance() or QApplication([])
    page = VideoLoraPage(root, ToolRegistry(root), Jobs(), AssetRegistry(root), ConfigManager(root))
    page.show()
    app.processEvents()
    page.dataset_folder.setText(str(dataset))
    names = iter(["kept_clips", "fresh_clips"])
    monkeypatch.setattr(QInputDialog, "getText", lambda *_args, **_kwargs: (next(names), True))
    monkeypatch.setattr(QMessageBox, "question", lambda *_args, **_kwargs: QMessageBox.Yes)

    page._archive_and_new_dataset()

    archived = root / "ADAM_Datasets" / "Video" / "Archived" / "kept_clips"
    fresh = root / "ADAM_Datasets" / "Video" / "fresh_clips"
    assert (archived / "clip.mp4").is_file()
    assert (archived / "clip.txt").is_file()
    assert fresh.is_dir()
    assert page.dataset_folder.text() == str(fresh.resolve())
    assert page.saved_datasets.findData(str(archived.resolve())) >= 0
    assert str(archived.resolve()) in ConfigManager(root).get("video_lora_settings")["dataset_history"]
    page.close()
    page.deleteLater()
    app.processEvents()


def test_desktop_caption_shortcut_and_batch_captioning(workspace, monkeypatch):
    monkeypatch.setenv("QT_QPA_PLATFORM", "offscreen")
    from PySide6.QtWidgets import QApplication
    from PySide6.QtCore import QObject, Signal, Qt
    from PySide6.QtTest import QTest
    from adam.ui.video_lora import VideoLoraPage

    class Jobs(QObject):
        job_updated = Signal(object)
        active_job = None

    root, _, dataset = workspace
    second = dataset / "second.mp4"
    second.write_bytes(b"video")
    app = QApplication.instance() or QApplication([])
    page = VideoLoraPage(root, ToolRegistry(root), Jobs(), AssetRegistry(root), ConfigManager(root))
    page.show()
    app.processEvents()
    page.dataset_folder.setText(str(dataset))
    page._refresh_clips()
    page.clips.setCurrentRow(0)
    page.clips.item(1).setSelected(True)
    page.caption.setPlainText("subject_token smiles at the camera")

    page._apply_caption_to_selected()

    assert (dataset / "clip.txt").read_text(encoding="utf-8").strip() == "subject_token smiles at the camera"
    assert (dataset / "second.txt").read_text(encoding="utf-8").strip() == "subject_token smiles at the camera"
    page.clips.setCurrentRow(0)
    page.caption.setPlainText("subject_token waves")
    page.caption.document().setModified(True)
    page.caption.setFocus()
    QTest.keyClick(page.caption, Qt.Key_Return, Qt.ShiftModifier)
    app.processEvents()
    assert page.clips.currentRow() == 1
    assert (dataset / "clip.txt").read_text(encoding="utf-8").strip() == "subject_token waves"
    page.close()
    page.deleteLater()
    app.processEvents()