from __future__ import annotations

import math

import pytest
import torch

from music3lab.codec.native_state_multiclip import (
    CAPABILITY_LABEL,
    LiteNativeStateBatch,
    ScaleUpMetrics,
    assert_no_residual_loss_path,
    deterministic_split_plan,
    derive_feedback_deterministically,
    extract_global_from_fused,
    scale_up_gate,
    validate_lite_native_state_batch,
)


FRAMES = 25
WIDTH = 4096


def _batch(n: int = 2) -> LiteNativeStateBatch:
    global_hidden = torch.randn(n, FRAMES, WIDTH).bfloat16()
    return LiteNativeStateBatch(
        audio=torch.randn(n, 2, 44_032, dtype=torch.float32).clamp(-1, 1),
        continuous_latent=torch.randn(n, 128, 86).bfloat16(),
        semantic=torch.randint(0, 16_384, (n, FRAMES), dtype=torch.int64),
        residual=torch.randint(0, 2048, (n, FRAMES, 7), dtype=torch.int64),
        priming=torch.randint(0, 16_384, (n, 8), dtype=torch.int64),
        global_hidden=global_hidden,
        feedback=torch.randn(n, FRAMES, WIDTH).bfloat16(),
        seeds=torch.arange(n, dtype=torch.int64),
    )


def test_lite_schema_accepts_exact_tensors_and_rejects_bad_dtype_shape_nonfinite() -> None:
    batch = _batch()
    validate_lite_native_state_batch(batch)
    with pytest.raises(ValueError, match="audio.*float32"):
        validate_lite_native_state_batch(LiteNativeStateBatch(**{**batch.__dict__, "audio": batch.audio.bfloat16()}))
    with pytest.raises(ValueError, match="continuous_latent.*128, 86"):
        validate_lite_native_state_batch(LiteNativeStateBatch(**{**batch.__dict__, "continuous_latent": batch.continuous_latent[..., :-1]}))
    bad_feedback = batch.feedback.clone()
    bad_feedback[0, 0, 0] = float("nan")
    with pytest.raises(ValueError, match="feedback.*finite"):
        validate_lite_native_state_batch(LiteNativeStateBatch(**{**batch.__dict__, "feedback": bad_feedback}))


def test_split_plan_is_deterministic_disjoint_and_has_the_bounded_composition() -> None:
    first = deterministic_split_plan()
    assert first == deterministic_split_plan()
    assert len(first.train) == 80
    assert len(first.validation) == len(first.heldout) == 12
    assert len(first.diagnostic_train) == len(first.diagnostic_heldout) == 16
    assert sum(item.source == "reused_french_house" for item in first.train) == 64
    assert sum(item.source == "fresh" for item in first.train) == 16
    assert len({item.prompt_id for item in first.train}) == 2
    primary = first.train + first.validation + first.heldout
    assert not ({item.prompt_id for item in first.validation} & {item.prompt_id for item in first.train})
    assert not ({item.prompt_id for item in first.heldout} & {item.prompt_id for item in first.train})
    assert not ({item.prompt_id for item in first.validation} & {item.prompt_id for item in first.heldout})
    assert len({(item.prompt_id, item.seed) for item in primary}) == len(primary)
    assert all(item.frames == FRAMES for item in primary)


def test_global_extraction_is_exact_fused_prefix_and_feedback_callback_consumes_no_rng() -> None:
    batch = _batch()
    fused = torch.randn(batch.audio.shape[0], FRAMES, 32_768).bfloat16()
    fused[..., :WIDTH] = batch.global_hidden
    assert torch.equal(extract_global_from_fused(fused), batch.global_hidden)
    state_before = torch.random.get_rng_state()
    observed = []
    def callback(global_hidden: torch.Tensor) -> torch.Tensor:
        observed.append(global_hidden.shape)
        return global_hidden + 1
    feedback = derive_feedback_deterministically(batch.global_hidden, callback)
    assert observed == [batch.global_hidden.shape]
    assert feedback.shape == (2, FRAMES, WIDTH)
    assert torch.equal(feedback, batch.global_hidden + 1)
    assert torch.equal(state_before, torch.random.get_rng_state())


def test_residual_is_provenance_only_not_a_loss_target() -> None:
    targets = assert_no_residual_loss_path(_batch())
    assert set(targets) == {"semantic", "feedback", "global_hidden"}
    assert all("residual" not in name for name in targets)


def test_scaleup_gate_requires_all_generalization_and_capability_conditions() -> None:
    baseline = ScaleUpMetrics(1.0, 1.0, 1.0, 1.0, math.log(16_384))
    passing = ScaleUpMetrics(0.80, 0.94, 0.80, 0.94, math.log(16_384) - 0.11)
    assert scale_up_gate(baseline, passing).accepted is True
    bad_gap = ScaleUpMetrics(0.40, 0.94, 0.40, 0.94, math.log(16_384) - 0.11)
    result = scale_up_gate(baseline, bad_gap)
    assert result.accepted is False
    assert "train_vs_heldout_gap" in result.reasons


def test_capability_is_truthfully_labeled_as_captured_multiclip_only() -> None:
    assert CAPABILITY_LABEL == "captured_music3_multiclip_only"
    assert CAPABILITY_LABEL != "native_tokenizer"
    assert CAPABILITY_LABEL != "arbitrary_external_music"