"""Tests for extraction module."""

import numpy as np
import pytest

from sltk.data.core import PoseSequence
from sltk.exceptions import ConfigurationError


def _check_mediapipe_available():
    """Check if MediaPipe is installed."""
    try:
        import mediapipe

        return True
    except ImportError:
        return False


class TestExtractionBase:
    """Tests for base extraction utilities."""

    def test_extraction_config(self):
        """Test ExtractionConfig creation."""
        from sltk.extraction import ExtractionConfig

        config = ExtractionConfig(
            device="cuda:0",
            batch_size=64,
            confidence_threshold=0.3,
        )

        assert config.device == "cuda:0"
        assert config.batch_size == 64
        assert config.confidence_threshold == 0.3

    def test_extraction_result(self):
        """Test ExtractionResult creation."""
        from sltk.extraction import ExtractionResult

        # Successful result
        data = np.random.randn(100, 21, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=25.0, format="test")

        result = ExtractionResult(
            poses=poses,
            num_frames=100,
            num_detections=100,
            fps=25.0,
        )

        assert result.success
        assert result.poses is not None
        assert result.num_frames == 100

    def test_extraction_result_failure(self):
        """Test ExtractionResult with errors."""
        from sltk.extraction import ExtractionResult

        result = ExtractionResult(
            errors=["Model not found"],
            fps=25.0,
        )

        assert not result.success
        assert result.poses is None
        assert len(result.errors) == 1


class TestMediaPipeExtractor:
    """Tests for MediaPipe extractor."""

    def test_config_creation(self):
        """Test MediaPipeConfig creation."""
        from sltk.extraction.mediapipe import MediaPipeConfig

        config = MediaPipeConfig(
            enable_body=True,
            enable_hands=True,
            enable_face=False,
            model_complexity=1,
        )

        assert config.enable_body
        assert config.enable_hands
        assert not config.enable_face

    def test_extractor_initialization(self):
        """Test MediaPipeExtractor initialization."""
        from sltk.extraction.mediapipe import MediaPipeExtractor

        extractor = MediaPipeExtractor()

        assert extractor.name == "mediapipe"
        assert extractor.output_format == "mediapipe"
        assert not extractor.is_initialized

    @pytest.mark.slow
    @pytest.mark.skipif(not _check_mediapipe_available(), reason="MediaPipe not installed")
    def test_extractor_load_model(self):
        """Test loading MediaPipe model."""
        from sltk.extraction.mediapipe import MediaPipeExtractor

        extractor = MediaPipeExtractor()
        extractor.load_model()

        assert extractor.is_initialized
        extractor.close()

    @pytest.mark.slow
    @pytest.mark.skipif(not _check_mediapipe_available(), reason="MediaPipe not installed")
    def test_extract_from_frames(self):
        """Test extraction from numpy frames."""
        from sltk.extraction.mediapipe import MediaPipeConfig, MediaPipeExtractor

        # Create dummy frames
        frames = np.random.randint(0, 255, (10, 480, 640, 3), dtype=np.uint8)

        config = MediaPipeConfig(
            enable_body=True,
            enable_hands=False,  # Faster test
            verbose=False,
        )
        extractor = MediaPipeExtractor(config)

        try:
            result = extractor.extract_from_frames(frames, fps=30.0)

            assert result.poses is not None
            assert result.poses.num_frames == 10
            assert result.fps == 30.0
        finally:
            extractor.close()


class TestWiLoRExtractor:
    """Tests for WiLoR extractor."""

    def test_config_creation(self):
        """Test WiLoRConfig creation."""
        from sltk.extraction.wilor import WiLoRConfig

        config = WiLoRConfig(
            checkpoint_path="/path/to/model.pth",
            rescale_factor=2.0,
            detection_confidence=0.3,
        )

        assert config.checkpoint_path == "/path/to/model.pth"
        assert config.rescale_factor == 2.0

    def test_extractor_initialization(self):
        """Test WiLoRExtractor initialization."""
        from sltk.extraction.wilor import WiLoRExtractor

        extractor = WiLoRExtractor()

        assert extractor.name == "wilor"
        assert extractor.output_format == "wilor"
        assert not extractor.is_initialized

    def test_bundled_weights_resolved(self):
        """Test that bundled weights are automatically resolved."""
        from sltk.extraction.weights import get_weight_path
        from sltk.extraction.wilor import WiLoRConfig

        bundled_checkpoint = get_weight_path("wilor_checkpoint")
        if not bundled_checkpoint:
            pytest.skip("WiLoR bundled weights not available")

        config = WiLoRConfig()
        assert config.checkpoint_path is not None, "Checkpoint should be auto-resolved"
        assert config.detector_path is not None, "Detector should be auto-resolved"
        assert config.checkpoint_path == str(bundled_checkpoint)


class TestNLFExtractor:
    """Tests for NLF extractor."""

    def test_config_creation(self):
        """Test NLFConfig creation."""
        from sltk.extraction.nlf import NLFConfig

        config = NLFConfig(
            model_path="/path/to/model.pt",
            smplx_path="/path/to/smplx.npz",
            detection_confidence=0.05,
        )

        assert config.model_path == "/path/to/model.pt"
        assert config.detection_confidence == 0.05

    def test_extractor_initialization(self):
        """Test NLFExtractor initialization."""
        from sltk.extraction.nlf import NLFExtractor

        extractor = NLFExtractor()

        assert extractor.name == "nlf"
        assert extractor.output_format == "smplx"
        assert not extractor.is_initialized

    def test_load_requires_model_path(self):
        """Test that load_model raises ConfigurationError when model_path is None.

        NLFConfig.__post_init__ auto-resolves bundled weights if they exist
        on disk, so we must explicitly set model_path=None after construction
        to simulate missing weights and test the error path.
        """
        from sltk.extraction.nlf import NLFConfig, NLFExtractor

        config = NLFConfig()
        # Override auto-resolved path to simulate missing weights
        config.model_path = None
        extractor = NLFExtractor(config)

        with pytest.raises(ConfigurationError, match="model_path required"):
            extractor.load_model()


class TestPoseUplifter:
    """Tests for 2D to 3D pose uplift."""

    def test_config_creation(self):
        """Test UpliftConfig creation."""
        from sltk.extraction.uplift import UpliftConfig

        config = UpliftConfig(
            method="mlp",
            ik_iterations=100,
        )

        assert config.method == "mlp"
        assert config.ik_iterations == 100

    def test_uplifter_initialization(self):
        """Test PoseUplifter initialization."""
        from sltk.extraction.uplift import PoseUplifter

        uplifter = PoseUplifter()

        assert uplifter.config.method == "mlp"

    def test_triangulation(self):
        """Test multi-view triangulation."""
        from sltk.extraction.uplift import PoseUplifter, UpliftConfig

        # Create simple 2D observations from 2 cameras
        # Camera 1: looking down Z axis
        P1 = np.array(
            [
                [1, 0, 0, 0],
                [0, 1, 0, 0],
                [0, 0, 1, 0],
            ],
            dtype=np.float32,
        )

        # Camera 2: rotated 90 degrees around Y
        P2 = np.array(
            [
                [0, 0, -1, 5],
                [0, 1, 0, 0],
                [1, 0, 0, 0],
            ],
            dtype=np.float32,
        )

        camera_matrices = np.stack([P1, P2])

        # A 3D point at (1, 2, 3)
        point_3d = np.array([1, 2, 3], dtype=np.float32)

        # Project to both cameras
        p1 = P1 @ np.append(point_3d, 1)
        p1 = p1[:2] / p1[2]

        p2 = P2 @ np.append(point_3d, 1)
        p2 = p2[:2] / p2[2]

        points_2d = np.stack([p1, p2])

        # Create uplifter and triangulate
        config = UpliftConfig(
            method="triangulation",
            camera_matrices=camera_matrices,
        )
        uplifter = PoseUplifter(config)

        recovered = uplifter._triangulate_point(points_2d, camera_matrices)

        # Should recover the 3D point
        np.testing.assert_array_almost_equal(recovered, point_3d, decimal=3)

    def test_normalize_poses_2d(self):
        """Test 2D pose normalization."""
        from sltk.extraction.uplift import PoseUplifter

        uplifter = PoseUplifter()

        # Create simple 2D poses
        poses_2d = np.array(
            [
                [[100, 200], [150, 250], [200, 300]],
                [[110, 210], [160, 260], [210, 310]],
            ],
            dtype=np.float32,
        )

        normalized = uplifter._normalize_poses_2d(poses_2d)

        # Check root is centered
        assert normalized.shape == poses_2d.shape
        # Check values are reasonable (normalized)
        assert np.abs(normalized).max() < 10


class TestExtractionIntegration:
    """Integration tests for extraction pipeline."""

    def test_lazy_import(self):
        """Test lazy import of extraction modules."""
        from sltk import extraction

        # These should not raise
        _ = extraction.PoseExtractor
        _ = extraction.ExtractionConfig
        _ = extraction.ExtractionResult

    def test_extraction_result_to_posesequence(self):
        """Test converting ExtractionResult to PoseSequence."""
        from sltk.extraction import ExtractionResult

        data = np.random.randn(50, 33, 3).astype(np.float32)
        poses = PoseSequence(data=data, fps=30.0, format="mediapipe")

        result = ExtractionResult(
            poses=poses,
            num_frames=50,
            fps=30.0,
        )

        # Should be able to access poses directly
        assert result.poses.num_frames == 50
        assert result.poses.fps == 30.0
        assert result.poses.format == "mediapipe"
