# SPDX-License-Identifier: Apache-2.0
"""Per-method GPU smoke test: ``WanCausalModel`` + ``StreamingLongTuningMethod``.

Mirrors ``test_wan_causal_cd.py``. LongLive-style streaming long tuning runs a
short self-forcing stage first, then switches to a persistent streaming
sequence trained chunk-by-chunk with the causal student KV cache. The GPU test
walks a single method instance through both stages: the short-stage step, the
first streaming chunk, an overlapped chunk, the chunk that fills the sequence,
and the post-reset chunk, asserting finite losses, student/critic gradient
flow, and the streaming-length bookkeeping at every step.

The pure-tensor helpers (stage selection, schedule validation, first-frame
padding, latent truncation, chunk sizing) are covered by CPU-only tests at the
bottom — they need no GPU, no mocks, and run wherever this file is collected.
"""

from __future__ import annotations

import os

os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29523")

from pathlib import Path
from types import SimpleNamespace

import pytest
import torch

from fastvideo.train.methods.distribution_matching.streaming_long_tuning import (
    DistillStage,
    StreamingLongTuningMethod,
    parse_multi_phased_distill_schedule,
    select_distill_stage,
)
from fastvideo.train.models.wan import WanCausalModel, WanModel
from fastvideo.train.utils.config import load_run_config

from .grad_norm_regression import check_grad_norm_regression

_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_streaming_long_min.yaml")


def _build_synthetic_batch(
    device: torch.device,
    dtype: torch.dtype,
) -> dict[str, torch.Tensor]:
    batch_size = 1
    return {
        "text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
        "text_attention_mask": torch.ones(batch_size, 16, device=device),
        "vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
    }


def _assert_layer0_grads(model, *, role: str) -> None:
    blocks = model.transformer.blocks
    assert blocks is not None and len(blocks) > 0
    trainable = [p for p in blocks[0].parameters() if p.requires_grad]
    assert len(trainable) > 0, f"{role} layer 0 has no trainable parameters"
    for i, p in enumerate(trainable):
        assert p.grad is not None, f"{role} layer 0 param[{i}] has None grad"
        assert torch.isfinite(p.grad).all().item(), (f"{role} layer 0 param[{i}] grad contains NaN/Inf")
    assert any(p.grad.detach().float().norm().item() > 0.0
               for p in trainable), (f"all {role} layer-0 grads are exactly zero; the loss did not "
                                     "reach the first transformer block")


def _assert_finite_losses(loss_map: dict[str, torch.Tensor]) -> None:
    for name in ("total_loss", "generator_loss", "fake_score_loss"):
        loss = loss_map[name]
        assert torch.is_tensor(loss), f"{name} must be a torch.Tensor"
        assert torch.isfinite(loss).item(), (f"{name} is not finite: {loss.item()}")


@pytest.mark.usefixtures("distributed_setup")
def test_streaming_long_tuning_multi_stage_train_steps(monkeypatch: pytest.MonkeyPatch) -> None:
    if not torch.cuda.is_available():
        pytest.skip("requires CUDA")

    cfg = load_run_config(_FIXTURE)

    device = torch.device("cuda:0")
    dtype = torch.bfloat16

    monkeypatch.setattr(
        "fastvideo.train.utils.dataloader."
        "build_parquet_t2v_train_dataloader",
        lambda *args, **kwargs: None,
    )

    student = WanCausalModel(
        init_from=cfg.models["student"]["init_from"],
        training_config=cfg.training,
        trainable=True,
    )
    student.transformer = student.transformer.to(device=device, dtype=dtype)

    teacher = WanModel(
        init_from=cfg.models["teacher"]["init_from"],
        training_config=cfg.training,
        trainable=False,
    )
    teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)

    critic = WanModel(
        init_from=cfg.models["critic"]["init_from"],
        training_config=cfg.training,
        trainable=True,
    )
    critic.transformer = critic.transformer.to(device=device, dtype=dtype)

    method = StreamingLongTuningMethod(
        cfg=cfg,
        role_models={
            "student": student,
            "teacher": teacher,
            "critic": critic,
        },
    )
    method.on_train_start()

    batch = _build_synthetic_batch(device, dtype)

    # Iteration 0 — short self-forcing stage (stage 0, num_latent_t=3).
    loss_map, outputs, metrics = method.single_train_step(dict(batch), 0)
    _assert_finite_losses(loss_map)
    assert metrics["distill_stage_index"] == 0.0
    assert metrics["active_num_latent_t"] == 3.0
    assert metrics["update_student"] == 1.0
    method.backward(loss_map, outputs, grad_accum_rounds=1)
    _assert_layer0_grads(student, role="student")
    _assert_layer0_grads(critic, role="critic")
    assert all(not p.requires_grad for p in teacher.transformer.parameters()), ("teacher must stay frozen")

    # Iteration 1 — first streaming chunk: 3 fresh latents, no overlap.
    method.optimizers_zero_grad(1)
    loss_map, outputs, metrics = method.single_train_step(dict(batch), 1)
    _assert_finite_losses(loss_map)
    assert metrics["distill_stage_index"] == 1.0
    assert metrics["streaming_current_length"] == 3.0
    assert metrics["streaming_max_length"] == 6.0
    method.backward(loss_map, outputs, grad_accum_rounds=1)
    _assert_layer0_grads(student, role="student")
    _assert_layer0_grads(critic, role="critic")
    vis = outputs["dmd_latent_vis_dict"]
    assert int(vis["streaming_chunk_mask"].sum().item()) == 3

    # Iteration 2 — overlapped chunk: 2 new latents + 1 overlap anchor.
    method.optimizers_zero_grad(2)
    loss_map, outputs, metrics = method.single_train_step(dict(batch), 2)
    _assert_finite_losses(loss_map)
    assert metrics["streaming_current_length"] == 5.0
    method.backward(loss_map, outputs, grad_accum_rounds=1)
    mask = outputs["dmd_latent_vis_dict"]["streaming_chunk_mask"]
    assert mask.dtype == torch.bool
    assert int(mask.sum().item()) == 2, ("overlap latents must be masked out of the training chunk")

    # Iteration 3 — final chunk fills the sequence (1 new latent).
    method.optimizers_zero_grad(3)
    loss_map, outputs, metrics = method.single_train_step(dict(batch), 3)
    _assert_finite_losses(loss_map)
    assert metrics["streaming_current_length"] == 6.0
    method.backward(loss_map, outputs, grad_accum_rounds=1)

    # Iteration 4 — the exhausted sequence resets and a fresh chunk starts.
    method.optimizers_zero_grad(4)
    loss_map, outputs, metrics = method.single_train_step(dict(batch), 4)
    _assert_finite_losses(loss_map)
    assert metrics["streaming_current_length"] == 3.0, (
        "streaming state must reset once the sequence reaches max length")
    method.backward(loss_map, outputs, grad_accum_rounds=1)
    _assert_layer0_grads(student, role="student")

    # Grad-norm regression on the post-reset chunk (grads zeroed at the top
    # of this iteration, so the reference captures exactly one step). Keep
    # this last: under FASTVIDEO_GRADNORM_UPDATE=1 it records and skips.
    check_grad_norm_regression(
        "test_streaming_long_tuning",
        student.transformer,
    )


# ---------------------------------------------------------------------------
# CPU-only helper tests (no GPU, no mocks)
# ---------------------------------------------------------------------------


def _method_without_init() -> StreamingLongTuningMethod:
    # Bare instance for exercising stateless helpers; no attributes beyond
    # the ones each test sets are touched.
    return StreamingLongTuningMethod.__new__(StreamingLongTuningMethod)


def _make_stage(**overrides) -> DistillStage:
    defaults = dict(
        name="streaming_long",
        start_step=0,
        end_step=None,
        num_latent_t=6,
        streaming_training=True,
        streaming_chunk_size=3,
        streaming_max_length=6,
    )
    defaults.update(overrides)
    return DistillStage(**defaults)


class TestStageSelection:

    def test_boundaries_between_stages(self) -> None:
        short = _make_stage(
            name="self_forcing",
            start_step=0,
            end_step=10,
            streaming_training=False,
            streaming_chunk_size=None,
            streaming_max_length=None,
        )
        long = _make_stage(start_step=10)
        stages = [short, long]

        assert select_distill_stage(stages, 0) is short
        assert select_distill_stage(stages, 9) is short
        assert select_distill_stage(stages, 10) is long
        assert select_distill_stage(stages, 10_000) is long

    def test_past_last_closed_stage_falls_back_to_last(self) -> None:
        short = _make_stage(
            name="self_forcing",
            start_step=0,
            end_step=10,
            streaming_training=False,
            streaming_chunk_size=None,
            streaming_max_length=None,
        )
        assert select_distill_stage([short], 99) is short


class TestScheduleParsing:

    def test_rejects_overlapping_stages(self) -> None:
        with pytest.raises(ValueError, match="ordered and non-overlapping"):
            parse_multi_phased_distill_schedule(
                [
                    {
                        "stage": "self_forcing",
                        "start_step": 0,
                        "end_step": 10
                    },
                    {
                        "stage": "streaming_long",
                        "start_step": 5
                    },
                ],
                default_num_latent_t=6,
                default_streaming_chunk_size=3,
            )

    def test_rejects_fixed_overlap_not_smaller_than_chunk(self) -> None:
        with pytest.raises(ValueError, match="smaller than"):
            parse_multi_phased_distill_schedule(
                [{
                    "stage": "streaming_long",
                    "streaming_chunk_size": 3,
                    "streaming_fixed_overlap_latents": 3,
                }],
                default_num_latent_t=6,
            )


class TestPadFirstFrameLatent:

    def test_pads_short_condition_with_zeros(self) -> None:
        method = _method_without_init()
        first_frame = torch.ones(1, 4, 1, 2, 2)

        batch = method._with_padded_first_frame_latent(
            {"first_frame_latent": first_frame},
            target_latent_frames=3,
        )

        padded = batch["first_frame_latent"]
        assert padded.shape == (1, 4, 3, 2, 2)
        assert torch.equal(padded[:, :, :1], first_frame)
        assert padded[:, :, 1:].abs().sum().item() == 0.0

    def test_passthrough_without_condition_or_when_long_enough(self) -> None:
        method = _method_without_init()
        raw: dict = {"other": 1}
        assert method._with_padded_first_frame_latent(raw, target_latent_frames=3) is raw

        full = {"first_frame_latent": torch.zeros(1, 4, 3, 2, 2)}
        assert method._with_padded_first_frame_latent(full, target_latent_frames=3) is full

    def test_rejects_non_5d_condition(self) -> None:
        method = _method_without_init()
        with pytest.raises(ValueError, match=r"\[B, C, T, H, W\]"):
            method._with_padded_first_frame_latent(
                {"first_frame_latent": torch.zeros(4, 1, 2, 2)},
                target_latent_frames=3,
            )


class TestTruncateBatchLatents:

    def test_slices_every_temporal_field(self) -> None:
        method = _method_without_init()
        # No _build_attention_metadata on the student → metadata rebuild is
        # a no-op, so only the tensor slicing is exercised.
        method.student = object()
        batch = SimpleNamespace(
            latents=torch.randn(1, 6, 2, 4, 4),
            noise_latents=torch.randn(1, 6, 2, 4, 4),
            noisy_model_input=torch.randn(1, 6, 2, 4, 4),
            noise=torch.randn(1, 6, 2, 4, 4),
            timesteps=torch.zeros(1, 6),
            sigmas=torch.linspace(0, 1, 6),
            raw_latent_shape=(1, 6, 2, 4, 4),
            attn_metadata=None,
            attn_metadata_vsa=None,
        )

        method._truncate_batch_latents(batch, 3)

        assert batch.latents.shape == (1, 3, 2, 4, 4)
        assert batch.noise_latents.shape == (1, 3, 2, 4, 4)
        assert batch.noisy_model_input.shape == (1, 3, 2, 4, 4)
        assert batch.noise.shape == (1, 3, 2, 4, 4)
        assert batch.timesteps.shape == (1, 3)
        assert batch.sigmas.shape == (3, )
        assert batch.raw_latent_shape == (1, 2, 3, 4, 4)


class TestStageChunkSize:

    def test_prefers_stage_then_config_then_rollout_block(self) -> None:
        method = _method_without_init()
        method.method_config = {"streaming_chunk_size": 5}
        method._chunk_size = 7

        assert method._stage_chunk_size(_make_stage()) == 3
        assert method._stage_chunk_size(_make_stage(streaming_chunk_size=None)) == 5

        method.method_config = {}
        assert method._stage_chunk_size(_make_stage(streaming_chunk_size=None)) == 7

    def test_rejects_non_positive_chunk_size(self) -> None:
        method = _method_without_init()
        method.method_config = {}
        method._chunk_size = 0
        with pytest.raises(ValueError, match="must be positive"):
            method._stage_chunk_size(_make_stage(streaming_chunk_size=None))