"""
Tests for H5 save/load functionality.
"""

import tempfile
from pathlib import Path

import numpy as np
import pytest

from sltk.data.core import PoseSequence
from sltk.io.h5 import get_h5_metadata, load_poses_h5, save_poses_h5


@pytest.fixture
def h5_sample_poses():
    """Create sample pose data for H5 round-trip tests.

    Uses seeded RNG and includes confidence + metadata fields
    needed by H5 save/load tests. Named to avoid shadowing conftest
    sample_poses fixture.
    """
    rng = np.random.default_rng(42)
    return PoseSequence(
        data=rng.standard_normal((100, 33, 3)).astype(np.float32),
        fps=25.0,
        format="mediapipe",
        confidence=rng.random((100, 33)).astype(np.float32),
        metadata={"test_key": "test_value", "num_kpts": 33},
    )


@pytest.fixture
def temp_h5_path():
    """Create a temporary H5 file path."""
    with tempfile.TemporaryDirectory() as tmpdir:
        yield Path(tmpdir) / "test_poses.h5"


class TestSavePosesH5:
    """Tests for save_poses_h5 function."""

    def test_save_basic(self, h5_sample_poses, temp_h5_path):
        """Test basic pose saving."""
        result_path = save_poses_h5(h5_sample_poses, temp_h5_path)
        assert result_path.exists()
        assert result_path == temp_h5_path

    def test_save_creates_parent_dirs(self, h5_sample_poses, temp_h5_path):
        """Test that parent directories are created."""
        nested_path = temp_h5_path.parent / "nested" / "dir" / "poses.h5"
        result_path = save_poses_h5(h5_sample_poses, nested_path)
        assert result_path.exists()

    def test_save_with_compression(self, h5_sample_poses, temp_h5_path):
        """Test saving with compression."""
        result_path = save_poses_h5(h5_sample_poses, temp_h5_path, compress=True)
        assert result_path.exists()

    def test_save_without_compression(self, h5_sample_poses, temp_h5_path):
        """Test saving without compression."""
        result_path = save_poses_h5(h5_sample_poses, temp_h5_path, compress=False)
        assert result_path.exists()

    def test_save_with_extra_metadata(self, h5_sample_poses, temp_h5_path):
        """Test saving with additional metadata."""
        metadata = {"source": "test", "version": 1}
        save_poses_h5(h5_sample_poses, temp_h5_path, metadata=metadata)

        loaded_meta = get_h5_metadata(temp_h5_path)
        assert loaded_meta.get("source") == "test"


class TestLoadPosesH5:
    """Tests for load_poses_h5 function."""

    def test_load_basic(self, h5_sample_poses, temp_h5_path):
        """Test basic pose loading."""
        save_poses_h5(h5_sample_poses, temp_h5_path)
        loaded = load_poses_h5(temp_h5_path)

        assert loaded.num_frames == h5_sample_poses.num_frames
        assert loaded.num_keypoints == h5_sample_poses.num_keypoints
        assert loaded.fps == h5_sample_poses.fps
        assert loaded.format == h5_sample_poses.format

    def test_load_data_matches(self, h5_sample_poses, temp_h5_path):
        """Test that loaded data matches original."""
        save_poses_h5(h5_sample_poses, temp_h5_path)
        loaded = load_poses_h5(temp_h5_path)

        np.testing.assert_array_almost_equal(loaded.data, h5_sample_poses.data)

    def test_load_confidence_matches(self, h5_sample_poses, temp_h5_path):
        """Test that loaded confidence matches original."""
        save_poses_h5(h5_sample_poses, temp_h5_path)
        loaded = load_poses_h5(temp_h5_path)

        assert loaded.confidence is not None
        np.testing.assert_array_almost_equal(loaded.confidence, h5_sample_poses.confidence)

    def test_load_with_fps_override(self, h5_sample_poses, temp_h5_path):
        """Test loading with FPS override."""
        save_poses_h5(h5_sample_poses, temp_h5_path)
        loaded = load_poses_h5(temp_h5_path, fps=30.0)

        assert loaded.fps == 30.0

    def test_load_without_confidence(self, temp_h5_path):
        """Test loading poses saved without confidence."""
        poses = PoseSequence(
            data=np.random.default_rng(42).standard_normal((50, 21, 3)).astype(np.float32),
            fps=25.0,
            format="wilor",
        )
        save_poses_h5(poses, temp_h5_path)
        loaded = load_poses_h5(temp_h5_path)

        assert loaded.num_frames == 50
        # Confidence might be None if not saved
        # This depends on implementation


class TestGetH5Metadata:
    """Tests for get_h5_metadata function."""

    def test_get_metadata_basic(self, h5_sample_poses, temp_h5_path):
        """Test getting metadata from H5 file."""
        save_poses_h5(h5_sample_poses, temp_h5_path)
        meta = get_h5_metadata(temp_h5_path)

        assert meta["fps"] == h5_sample_poses.fps
        assert meta["format"] == h5_sample_poses.format
        assert meta["num_frames"] == h5_sample_poses.num_frames

    def test_get_metadata_lists_features(self, h5_sample_poses, temp_h5_path):
        """Test that metadata includes feature list."""
        save_poses_h5(h5_sample_poses, temp_h5_path)
        meta = get_h5_metadata(temp_h5_path)

        assert "features" in meta
        assert isinstance(meta["features"], list)
        assert "poses" in meta["features"]


class TestRoundTrip:
    """Tests for complete save/load round-trips."""

    def test_roundtrip_preserves_shape(self, temp_h5_path):
        """Test that round-trip preserves data shape."""
        shapes = [(50, 33, 3), (100, 42, 3), (200, 75, 3)]

        for shape in shapes:
            poses = PoseSequence(
                data=np.random.default_rng(42).standard_normal(shape).astype(np.float32),
                fps=25.0,
                format="test",
            )
            save_poses_h5(poses, temp_h5_path)
            loaded = load_poses_h5(temp_h5_path)

            assert loaded.data.shape == shape, f"Shape mismatch for {shape}"

    def test_roundtrip_different_formats(self, temp_h5_path):
        """Test round-trip with different pose formats."""
        formats = ["mediapipe", "wilor", "smplx", "custom"]

        for fmt in formats:
            poses = PoseSequence(
                data=np.random.default_rng(42).standard_normal((10, 21, 3)).astype(np.float32),
                fps=30.0,
                format=fmt,
            )
            save_poses_h5(poses, temp_h5_path)
            loaded = load_poses_h5(temp_h5_path)

            assert loaded.format == fmt, f"Format mismatch for {fmt}"
