# SPDX-License-Identifier: Apache-2.0
"""Tests for the typed LTX-2 continuation state.

Covers:

* round-trip through :class:`ContinuationState` (inline and blob-backed)
* payload is JSON-serializable (Dynamo RPC / HTTP client constraint)
* kind / schema_version validation on deserialization
* compat-layer validation (known kinds, payload shape)
* round-trip through :func:`request_to_sampling_param` attaches the
  state to the resulting :class:`SamplingParam` without losing fidelity
"""
from __future__ import annotations

import json

import numpy as np
import pytest
import torch

# Importing compat first, then the LTX-2 module, exercises the
# self-registration side effect on import (important for the API
# test suite where the pipeline package isn't otherwise imported).
from fastvideo.api import compat as api_compat  # noqa: F401
from fastvideo.api.schema import (
    ContinuationState,
    GenerationRequest,
    OutputConfig,
)
from fastvideo.entrypoints.streaming.session_store import InMemoryBlobStore
from fastvideo.pipelines.basic.ltx2.continuation import (
    LTX2_CONTINUATION_KIND,
    LTX2_CONTINUATION_SCHEMA_VERSION,
    LTX2ContinuationState,
)


def _make_typed_state() -> LTX2ContinuationState:
    return LTX2ContinuationState(
        segment_index=3,
        video_frames=[(np.ones((64, 64, 3), dtype=np.uint8) * (i * 10)) for i in range(4)],
        video_conditioning_frame_idx=9,
        video_conditioning_strength=0.75,
        audio_latents=torch.randn(1, 4, 16, 64, dtype=torch.float32),
        audio_sample_rate=24000,
        audio_conditioning_num_frames=5,
        audio_conditioning_strength=0.5,
        video_position_offset_sec=0.125,
        metadata={"note": "unit-test"},
    )


class TestRoundTrip:
    """Round-trip through :class:`ContinuationState` preserves all fields."""

    def test_kind_and_schema_version(self):
        state = _make_typed_state().to_continuation_state()
        assert state.kind == LTX2_CONTINUATION_KIND
        assert state.payload["schema_version"] == LTX2_CONTINUATION_SCHEMA_VERSION

    def test_inline_roundtrip_preserves_scalars(self):
        original = _make_typed_state()
        envelope = original.to_continuation_state()
        restored = LTX2ContinuationState.from_continuation_state(envelope)
        assert restored.segment_index == original.segment_index
        assert restored.video_conditioning_frame_idx == (original.video_conditioning_frame_idx)
        assert restored.video_conditioning_strength == (original.video_conditioning_strength)
        assert restored.audio_sample_rate == original.audio_sample_rate
        assert restored.audio_conditioning_num_frames == (original.audio_conditioning_num_frames)
        assert restored.audio_conditioning_strength == (original.audio_conditioning_strength)
        assert restored.video_position_offset_sec == (original.video_position_offset_sec)
        assert restored.metadata == original.metadata

    def test_inline_roundtrip_preserves_video_frames(self):
        original = _make_typed_state()
        envelope = original.to_continuation_state()
        restored = LTX2ContinuationState.from_continuation_state(envelope)
        assert restored.video_frames is not None
        assert len(restored.video_frames) == len(original.video_frames)
        for before, after in zip(original.video_frames, restored.video_frames):
            np.testing.assert_array_equal(before, after)

    def test_inline_roundtrip_preserves_audio_latents(self):
        original = _make_typed_state()
        envelope = original.to_continuation_state()
        restored = LTX2ContinuationState.from_continuation_state(envelope)
        assert restored.audio_latents is not None
        assert tuple(restored.audio_latents.shape) == tuple(original.audio_latents.shape)
        assert restored.audio_latents.dtype == original.audio_latents.dtype
        torch.testing.assert_close(restored.audio_latents, original.audio_latents)

    def test_payload_is_json_serializable(self):
        envelope = _make_typed_state().to_continuation_state()
        # json.dumps must not raise — required for Dynamo RPC transport
        # and HTTP client round-trip.
        reserialized = json.loads(json.dumps(envelope.payload))
        restored = LTX2ContinuationState.from_continuation_state(
            ContinuationState(
                kind=envelope.kind,
                payload=reserialized,
            ))
        assert restored.segment_index == 3

    def test_bf16_audio_latents_preserved(self):
        """safetensors serialization must preserve bf16 dtype (numpy
        has no bf16, so a raw-bytes path would silently promote)."""
        state = LTX2ContinuationState(
            segment_index=0,
            audio_latents=torch.randn(1, 4, 16, 64, dtype=torch.bfloat16),
        )
        envelope = state.to_continuation_state()
        restored = LTX2ContinuationState.from_continuation_state(envelope)
        assert restored.audio_latents is not None
        assert restored.audio_latents.dtype == torch.bfloat16
        torch.testing.assert_close(restored.audio_latents, state.audio_latents)


class TestBlobIndirection:
    """Large tensors live in the :class:`BlobStore` instead of the payload."""

    def test_threshold_triggers_blob_path(self):
        blob_store = InMemoryBlobStore()
        state = _make_typed_state()
        envelope = state.to_continuation_state(
            blob_store=blob_store,
            inline_threshold_bytes=0,
        )
        assert "blob_id" in envelope.payload["video"]
        assert "blob_id" in envelope.payload["audio"]
        assert "frames_b64" not in envelope.payload["video"]
        assert "safetensors_b64" not in envelope.payload["audio"]
        assert len(blob_store) == 2

    def test_blob_roundtrip_reconstructs_tensors(self):
        blob_store = InMemoryBlobStore()
        original = _make_typed_state()
        envelope = original.to_continuation_state(
            blob_store=blob_store,
            inline_threshold_bytes=0,
        )
        restored = LTX2ContinuationState.from_continuation_state(envelope, blob_store=blob_store)
        assert restored.video_frames is not None
        assert len(restored.video_frames) == len(original.video_frames)
        torch.testing.assert_close(restored.audio_latents, original.audio_latents)

    def test_blob_id_held_when_store_unavailable(self):
        """Deserializing without a blob store preserves the blob id so
        the caller can fetch it later."""
        blob_store = InMemoryBlobStore()
        envelope = _make_typed_state().to_continuation_state(
            blob_store=blob_store,
            inline_threshold_bytes=0,
        )
        blob_id_video = envelope.payload["video"]["blob_id"]
        blob_id_audio = envelope.payload["audio"]["blob_id"]

        restored = LTX2ContinuationState.from_continuation_state(envelope)
        assert restored.video_frames is None
        assert restored.video_frames_blob_id == blob_id_video
        assert restored.audio_latents is None
        assert restored.audio_latents_blob_id == blob_id_audio

    def test_large_threshold_keeps_payload_inline(self):
        blob_store = InMemoryBlobStore()
        envelope = _make_typed_state().to_continuation_state(
            blob_store=blob_store,
            inline_threshold_bytes=10 * 1024 * 1024,  # 10 MiB
        )
        assert "frames_b64" in envelope.payload["video"]
        assert "safetensors_b64" in envelope.payload["audio"]
        assert len(blob_store) == 0


class TestValidation:
    """Invalid payloads error cleanly."""

    def test_wrong_kind_rejected(self):
        envelope = ContinuationState(kind="longcat.v1", payload={})
        with pytest.raises(ValueError, match="Expected ContinuationState.kind"):
            LTX2ContinuationState.from_continuation_state(envelope)

    def test_unsupported_schema_version_rejected(self):
        envelope = ContinuationState(
            kind=LTX2_CONTINUATION_KIND,
            payload={"schema_version": 999},
        )
        with pytest.raises(ValueError, match="Unsupported LTX-2 continuation schema"):
            LTX2ContinuationState.from_continuation_state(envelope)

    def test_non_png_frame_rejected(self):
        state = LTX2ContinuationState(video_frames=[np.ones((64, 64, 3), dtype=np.float32)], )
        with pytest.raises(ValueError, match="uint8 HxWx3"):
            state.to_continuation_state()


class TestCompatLayerWireUp:
    """The public compat layer accepts request.state without reverting
    to NotImplementedError and attaches it to the SamplingParam path."""

    def test_request_with_state_passes_through(self, tmp_path):
        # PR 7 removes the NotImplementedError for request.state; build a
        # minimal GenerationRequest carrying an LTX-2 state and make sure
        # the public boundary accepts it.
        from fastvideo.api.compat import (
            normalize_generation_request,
            _validate_continuation_state,
        )
        envelope = _make_typed_state().to_continuation_state()
        request = GenerationRequest(
            prompt="test",
            state=envelope,
        )
        normalized = normalize_generation_request(request)
        _validate_continuation_state(normalized.state)

    def test_unknown_kind_rejected_at_boundary(self):
        from fastvideo.api.compat import _validate_continuation_state
        with pytest.raises(ValueError, match="Unknown ContinuationState kind"):
            _validate_continuation_state(ContinuationState(kind="mystery.v1", payload={}))

    def test_empty_kind_rejected_at_boundary(self):
        from fastvideo.api.compat import _validate_continuation_state
        with pytest.raises(ValueError, match="non-empty string"):
            _validate_continuation_state(ContinuationState(kind="", payload={}))

    def test_output_return_state_flag(self):
        request = GenerationRequest(
            prompt="x",
            output=OutputConfig(return_state=True),
        )
        # The typed public surface exposes the flag directly.
        assert request.output.return_state is True