"""A whole training run through the module's own entry point, on a tiny real H3 on CPU.

The model is ComfyUI's own MiniMaxH3 class, a few hundred thousand parameters wide, with a stand-in text
encoder and VAEs that return latents of the right shape. What this proves is OUR side: the set, the
captions, the latents packed the way H3 wants, the two H3 repairs, the loss stream, and the saved file.
It does not prove anything about real weights, a real VAE or real hardware.
"""

import pytest


@pytest.fixture(autouse=True)
def _needs_comfy(comfyui):
    """Needs ComfyUI's source for the trainer."""


def tiny_h3_patcher():
    import torch
    import comfy.supported_models as sm
    from comfy.model_patcher import ModelPatcher
    unet = dict(image_model="minimax_h3", hidden_size=64, num_layers=2, token_refiner_num_layers=1, num_attention_heads=2,
                attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64,
                time_embed_dim=32, rope_inv_freq_len=4)
    cfg = sm.MiniMaxH3(dict(unet))
    cfg.set_inference_dtype(torch.float32, None)

    def make(disable_dynamic=False):
        base = cfg.get_model({}, "", device=torch.device("cpu"))
        torch.manual_seed(0)
        for p in base.parameters():
            torch.nn.init.normal_(p, std=0.05)
        base.diffusion_model.rope.inv_freq.copy_(torch.rand(4))
        return ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
    patcher = make()
    patcher.cached_patcher_init = (make, ())
    return patcher


class FakeClip:
    def tokenize(self, text):
        return text

    def encode_from_tokens_scheduled(self, tokens):
        import torch
        return [[torch.randn(1, 6, 48), {}]]


class FakeVae:
    def encode(self, pixels):
        import torch
        from comfy_extras.nodes_minimax_h3 import video_latent_t
        frames, h, w = pixels.shape[0], pixels.shape[1], pixels.shape[2]
        return torch.randn(1, 24, video_latent_t(frames), h // 16, w // 16)


class FakeAudioVae:
    audio_sample_rate = 44100

    def encode(self, wave):
        import torch
        seconds = wave.shape[1] / 44100
        return torch.randn(1, 32, 2, round(seconds * 40))


@pytest.fixture
def run_env(tmp_path, monkeypatch):
    import folder_paths
    from PIL import Image
    from modules.training.lora import jobs, store, train
    from core import config
    monkeypatch.setattr(config, "TRAINING_DIR", tmp_path / "training")
    monkeypatch.setattr(train, "roots", lambda: {"input": str(tmp_path)})
    loras = tmp_path / "loras"
    monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [str(loras)])
    from modules.models.minimax_h3 import training as h3
    monkeypatch.setattr(train, "registry", None, raising=False)
    import modules.training.lora.train as t
    class Reg:
        def ask(self, capability, *a, **k):
            return h3.adapt(*a, **k) if capability == "training_adapt" else None
    from modules import _core
    monkeypatch.setattr(_core.registry, "current", lambda: Reg())
    names = []
    for i in range(2):
        Image.new("RGB", (80, 80), (40 * i, 90, 160)).save(tmp_path / f"p{i}.png")
        names.append(f"p{i}.png")
    store.create("tiny")
    store.add("tiny", [{"src": {"kind": "output", "type": "input", "subfolder": "", "filename": n}, "media_kind": "image",
                        "label": "own", "caption": "a blue square, flat"} for n in names])
    return tmp_path, jobs


def settings(**over):
    base = {"name": "tiny_lora", "use": "normal", "steps": 6, "learning_rate": 3e-3, "rank": 4, "algorithm": "LoRA", "optimizer": "AdamW",
            "loss": "MSE", "batch_size": 4, "grad_accumulation": 1, "resolution": 64, "aspect_buckets": True, "frames": 5, "fps": 0.0,
            "start": "start", "trigger": "zq", "seed": 0, "training_dtype": "bf16", "lora_dtype": "fp32",
            "gradient_checkpointing": False, "offloading": False, "bypass": False}
    base.update(over)
    return base


def test_a_tiny_h3_trains_and_saves_a_lora(run_env):
    import safetensors.torch
    from modules.training.lora import train
    tmp_path, jobs = run_env
    file = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings())
    saved = safetensors.torch.load_file(str(tmp_path / "loras" / file))
    assert file == "tiny_lora.safetensors" and len(saved) > 20
    snap = jobs.snapshot()
    assert snap["state"] == "done" and snap["step"] == 6 and len(snap["loss"]) == 6
    assert all(x == x for x in snap["loss"]), "a NaN loss"


def test_a_second_run_never_overwrites_the_first(run_env):
    from modules.training.lora import train
    tmp_path, _ = run_env
    first = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
    second = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
    assert first != second and (tmp_path / "loras" / first).is_file()


def test_h3_without_its_audio_vae_says_what_to_add(run_env):
    from modules.training.lora import train
    with pytest.raises(ValueError, match="audio VAE"):
        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), None, "tiny", settings())


def test_a_model_the_trainer_cannot_copy_is_refused_before_training_with_the_reason(run_env):
    """A loader that converted the weights (MXFP8, int8) leaves no way to re-read the file: ComfyUI's own error must not be what is shown."""
    from modules.training.lora import train
    patcher = tiny_h3_patcher()
    patcher.cached_patcher_init = None
    with pytest.raises(ValueError, match="MXFP8"):
        train.run(patcher, FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))


def test_stop_ends_the_run_and_records_it(run_env):
    from modules.training.lora import train
    _, jobs = run_env
    real = jobs.step

    def stop_after_two(loss, seconds=None):
        real(loss, seconds)
        if jobs.snapshot()["step"] >= 2:
            jobs.request_stop()
    jobs.step = stop_after_two
    try:
        with pytest.raises(train.Stopped):
            train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=50))
    finally:
        jobs.step = real
    snap = jobs.snapshot()
    assert snap["state"] == "stopped" and snap["step"] < 50


def test_a_set_with_nothing_for_the_chosen_use_is_refused(run_env):
    from modules.training.lora import train
    with pytest.raises(ValueError, match="nothing to train on"):
        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(use="avoid"))


# --- a picture model: no adapter, plain latent tensors ------------------------------------------------

def tiny_sd_patcher():
    import torch
    import comfy.supported_models as sm
    from comfy.model_patcher import ModelPatcher
    unet = {"context_dim": 32, "model_channels": 32, "use_linear_in_transformer": False, "adm_in_channels": None, "use_temporal_attention": False,
            "in_channels": 4, "out_channels": 4, "image_size": 32, "num_res_blocks": [1, 1], "transformer_depth": [1, 1], "channel_mult": [1, 2],
            "transformer_depth_output": [1, 1, 1, 1], "transformer_depth_middle": 1, "use_spatial_transformer": True, "num_head_channels": 8}
    cfg = sm.SD15(dict(unet))
    cfg.set_inference_dtype(torch.float32, None)

    def make(disable_dynamic=False):
        base = cfg.get_model({}, "", device=torch.device("cpu"))
        torch.manual_seed(0)
        for p in base.parameters():
            torch.nn.init.normal_(p, std=0.05)
        return ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
    patcher = make()
    patcher.cached_patcher_init = (make, ())
    return patcher


class FakeClip32(FakeClip):
    def encode_from_tokens_scheduled(self, tokens):
        import torch
        return [[torch.randn(1, 6, 32), {}]]


class FakeVae4:
    def encode(self, pixels):
        import torch
        return torch.randn(pixels.shape[0], 4, pixels.shape[1] // 8, pixels.shape[2] // 8)


@pytest.fixture
def plain_env(run_env, monkeypatch):
    from modules import _core
    class NoAdapter:
        def ask(self, capability, *a, **k):
            return None
    monkeypatch.setattr(_core.registry, "current", lambda: NoAdapter())
    return run_env


def test_a_picture_model_trains_from_a_set_of_one_shape(plain_env):
    from modules.training.lora import train
    _, jobs = plain_env
    file = train.run(tiny_sd_patcher(), FakeClip32(), FakeVae4(), None, "tiny", settings(aspect_buckets=False, batch_size=2, steps=4))
    snap = jobs.snapshot()
    assert file.endswith(".safetensors") and snap["state"] == "done" and len(snap["loss"]) == 4
    assert all(x == x for x in snap["loss"])


def test_a_picture_model_trains_from_pictures_of_different_shapes(plain_env):
    from PIL import Image
    from modules.training.lora import store, train
    tmp_path, jobs = plain_env
    Image.new("RGB", (160, 64), (200, 30, 30)).save(tmp_path / "wide.png")
    store.add("tiny", [{"src": {"kind": "output", "type": "input", "subfolder": "", "filename": "wide.png"}, "media_kind": "image", "label": "own", "caption": "wide"}])
    train.run(tiny_sd_patcher(), FakeClip32(), FakeVae4(), None, "tiny", settings(aspect_buckets=True, batch_size=1, steps=3))
    assert jobs.snapshot()["state"] == "done"


def test_a_dynamic_model_is_copied_as_a_plain_one_for_the_trainer():
    from modules.training.lora import train

    class Dynamic:
        calls = []

        def is_dynamic(self):
            return True

        def clone(self, disable_dynamic=False, force_deepcopy=False):
            self.calls.append((disable_dynamic, force_deepcopy))
    m = Dynamic()
    train._train_on_a_plain_copy(m)
    m.clone(force_deepcopy=True)
    assert m.calls == [(True, True)]

    class Plain(Dynamic):
        def is_dynamic(self):
            return False
    p = Plain()
    train._train_on_a_plain_copy(p)
    assert "clone" not in vars(p)


def test_the_trainer_loads_the_model_file_once_and_its_copy_is_that_model(monkeypatch):
    import comfy.sd
    import folder_paths
    import torch
    from modules.training.lora import train
    seen = {}

    class Patcher:
        def __init__(self):
            self.weight = torch.zeros(2)                  # made where the loader runs: inference mode must not reach it

        def set_model_compute_dtype(self, dtype):
            seen["compute"] = dtype

    def load(path, model_options=None, disable_dynamic=False):
        seen.update(path=path, options=model_options, dynamic=disable_dynamic, loads=seen.get("loads", 0) + 1)
        return Patcher()
    monkeypatch.setattr(comfy.sd, "load_diffusion_model", load)
    monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/models/{name}")
    with torch.inference_mode():
        model = train.load_model("m.safetensors", "bf16", "fp32")
    assert not model.weight.is_inference()
    assert seen["path"] == "/models/m.safetensors" and seen["options"] == {"dtype": torch.bfloat16} and seen["dynamic"] and seen["compute"] == torch.float32
    factory, args = model.cached_patcher_init
    assert factory(*args, disable_dynamic=True) is model and seen["loads"] == 1       # the trainer's copy reads nothing again
    with pytest.raises(ValueError, match="GGUF"):
        train.load_model("m.gguf")
    with pytest.raises(ValueError, match="no model"):
        train.load_model("")


def test_h3_is_checkpointed_one_block_at_a_time_not_as_a_whole(run_env):
    import comfy_extras.nodes_train as nt
    from modules.models.minimax_h3 import training
    patcher = tiny_h3_patcher()
    depth = training.adapt(patcher)["checkpoint_depth"]
    net = patcher.model.diffusion_model
    pieces = nt.find_modules_at_depth(net, depth=depth)
    assert depth == 2 and len(pieces) > len(net.blocks) and all(b in pieces for b in net.blocks)
    assert net not in pieces


def test_what_was_made_in_inference_mode_is_copied_for_the_backward_pass():
    import torch
    from modules.training.lora import train
    with torch.inference_mode():
        made = [(torch.ones(2), {"k": torch.zeros(1), "n": 3})], [torch.ones(1)]
    assert made[1][0].is_inference()
    got = train._normal(made)
    assert not got[1][0].is_inference() and not got[0][0][0].is_inference() and not got[0][0][1]["k"].is_inference() and got[0][0][1]["n"] == 3
    x = got[1][0].requires_grad_(False)
    (x * torch.ones(1, requires_grad=True)).sum().backward()             # a saved-for-backward use no longer fails


def test_a_tiny_h3_trains_when_run_the_way_comfyui_runs_a_node(run_env):
    """ComfyUI runs every node in inference mode, and the models, text and latents here are made inside it."""
    import torch
    from modules.training.lora import train
    _, jobs = run_env
    with torch.inference_mode():
        patcher = tiny_h3_patcher()
        train.run(patcher, FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
    assert jobs.snapshot()["state"] == "done"


def test_block_swap_moves_the_last_blocks_to_system_memory_and_leaves_trainable_weights():
    import torch
    from modules.training.lora import swap

    class Layer(torch.nn.Linear):
        comfy_cast_weights = False

    def block():
        b = torch.nn.Sequential(Layer(4, 4), Layer(4, 4))
        for p in b[1].parameters():
            p.requires_grad_(True)               # a trainable weight is never moved
        return b

    blocks = [block() for _ in range(4)]
    done = swap.move(blocks, 2, torch)
    assert done["blocks"] == 2
    assert all(l.comfy_cast_weights for l in blocks[3].modules() if isinstance(l, Layer)), "the swapped layers fetch their weights"
    assert not any(l.comfy_cast_weights for l in blocks[0].modules() if isinstance(l, Layer))
    assert swap.move(blocks, 2, torch)["blocks"] == 0, "asking twice moves nothing again"
    assert swap.move(blocks, 0, torch)["blocks"] == 0


def test_a_tiny_h3_trains_with_block_swap_on(run_env):
    from modules.training.lora import train
    tmp_path, jobs = run_env
    file = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2, gradient_checkpointing=True, swap_blocks=1))
    snap = jobs.snapshot()
    assert file.endswith(".safetensors") and snap["state"] == "done" and all(x == x for x in snap["loss"])


def test_block_modulation_is_out_of_place_while_training_and_put_back_after():
    import torch
    from comfy.ldm.minimax import model as m
    from modules.models.minimax_h3 import training
    before = (m._mod_scale_shift, m._mod_gate)
    segs = [(0, 3, 0), (3, 5, 1)]
    shift, scale, gate = (torch.randn(2, 4) for _ in range(3))
    x, other = torch.randn(5, 4), torch.randn(5, 4)
    want_h = m._mod_scale_shift(x.clone(), shift, scale, segs)
    want_x = m._mod_gate(x.clone(), gate, other, segs)
    with training.shims():
        assert m._mod_scale_shift is not before[0]
        x0 = x.clone()
        assert torch.allclose(m._mod_scale_shift(x0, shift, scale, segs), want_h)
        assert torch.equal(x0, x)                           # the input is left alone
        assert torch.allclose(m._mod_gate(x0, gate, other, segs), want_x)
        assert torch.equal(x0, x)
    assert (m._mod_scale_shift, m._mod_gate) == before


def test_a_run_pauses_for_a_preview_and_carries_on_to_the_same_end(run_env):
    import safetensors.torch
    from modules.training.lora import train
    tmp_path, jobs = run_env
    first = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2))
    snap = jobs.snapshot()
    assert first == "" and snap["state"] == "paused"
    assert snap["paused"]["reason"] == "preview" and snap["paused"]["step"] == 2 and snap["step"] == 2
    preview = tmp_path / "loras" / snap["paused"]["preview"]
    assert preview.name == "tiny_lora_preview.safetensors"
    assert any(k.endswith(".lora_up.weight") for k in safetensors.torch.load_file(str(preview)))
    run = snap["run"]
    # the person looks at the preview, and training carries on from where it stopped
    second = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
    snap = jobs.snapshot()
    assert second == "" and snap["run"] == run and snap["paused"]["step"] == 4 and len(snap["loss"]) == 4
    done = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
    snap = jobs.snapshot()
    assert done == "tiny_lora.safetensors" and snap["state"] == "done" and len(snap["loss"]) == 6
    times = snap["times"]
    assert len(times) == 6 and [t is None for t in times] == [True, False, True, False, True, False], "each restart's first step also loads the model, so it is not timed"
    assert snap["step_seconds"] and snap["elapsed"] >= sum(t for t in times if t), "the time of every part of the run is counted"
    import json
    from modules.training.lora import jobs as j
    saved = json.loads((j._dir() / f"{run}.json").read_text())
    assert saved["seconds"] >= snap["elapsed"] - 1 and saved["step_seconds"] == snap["step_seconds"] and len(saved["times"]) == 6
    assert not preview.exists(), "the preview file stays after the run is over"
    assert not (tmp_path / "training" / "checkpoints" / f"{run}.pt").exists()


def test_a_failed_run_is_restored_from_its_saved_state_in_comfyuis_inference_mode(run_env, monkeypatch):
    """The optimizer's memory is read back in inference mode there; updating it in place used to fail on the first step."""
    import torch
    from modules.training.lora import jobs as j, loop, train
    tmp_path, jobs = run_env
    with torch.inference_mode():
        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2))
    run = jobs.snapshot()["run"]
    assert loop.has_checkpoint(run)
    real = j.step

    def fail(loss, seconds=None):
        real(loss, seconds)
        raise RuntimeError("the card fell over")
    monkeypatch.setattr(j, "step", fail)
    with torch.inference_mode(), pytest.raises(RuntimeError, match="fell over"):
        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
    assert jobs.snapshot()["state"] == "error" and loop.has_checkpoint(run), "a failed run keeps its saved state"
    monkeypatch.setattr(j, "step", real)
    with torch.inference_mode():
        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
        done = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
    snap = jobs.snapshot()
    assert done == "tiny_lora.safetensors" and snap["state"] == "done" and len(snap["loss"]) == 6
    assert not loop.has_checkpoint(run)


def test_a_second_run_on_the_same_data_does_not_encode_it_again(run_env, monkeypatch):
    from modules.training.lora import train
    tmp_path, jobs = run_env
    calls = {"video": 0, "text": 0}
    real_item, real_text = train.encode_item, train._encode_text

    def item(*a, **k):
        calls["video"] += 1
        return real_item(*a, **k)

    def text(*a, **k):
        calls["text"] += 1
        return real_text(*a, **k)
    monkeypatch.setattr(train, "encode_item", item)
    monkeypatch.setattr(train, "_encode_text", text)
    train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
    first = dict(calls)
    assert first["video"] > 0 and first["text"] > 0 and any((tmp_path / "training" / "cache").glob("latent-*.pt"))
    train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
    assert calls == first, "the clips and captions came from the cache"


def test_pause_button_waits_at_the_next_step_and_stop_ends_a_paused_run(run_env, monkeypatch):
    from modules.training.lora import jobs as j, train
    tmp_path, jobs = run_env
    real = j.step

    def press(loss, seconds=None):
        real(loss, seconds)
        j.request_pause()
    monkeypatch.setattr(j, "step", press)
    train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6))
    snap = jobs.snapshot()
    assert snap["state"] == "paused" and snap["paused"]["reason"] == "user" and snap["paused"]["step"] == 1
    assert j.request_stop() and jobs.snapshot()["state"] == "stopped"


def test_a_run_that_cannot_be_resumed_says_why(run_env):
    from modules.training.lora import train
    _, jobs = run_env
    with pytest.raises(ValueError, match="saved state is gone"):
        train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(resume="0123456789ab"))
    assert jobs.snapshot()["state"] == "error"


def test_copies_of_the_lora_so_far_are_saved_as_name_stepN(run_env):
    tmp_path, jobs = run_env
    from modules.training.lora import train
    train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, save_every=2))
    names = sorted(p.name for p in (tmp_path / "loras").glob("*.safetensors"))
    assert names == ["tiny_lora.safetensors", "tiny_lora_step2.safetensors", "tiny_lora_step4.safetensors"]