import statistics
from unittest.mock import MagicMock

import pytest

from redenlab_extract.asset import Asset
from redenlab_extract.exceptions import InferenceError
from redenlab_extract.workflows import (
    _default_aggregator_for,
    aggregate_results,
    run_chunked_inference,
    run_inference,
)


@pytest.fixture
def mock_session():
    """Create a mock RedenLabSession for testing."""
    session = MagicMock()
    session.api_key = "sk_live_test_key_1234567890123456"
    session.base_url = "https://api.test.example.com"
    return session


@pytest.fixture
def mock_client():
    """Create a mock InferenceClient for testing."""
    client = MagicMock()
    client.chunk_threshold = 50.0
    client.score_key = "intelligibility_score"
    client.model_name = "als-intelligibility-v1"
    return client


@pytest.fixture
def mock_asset():
    """Create a mock Asset for testing."""
    return Asset(
        asset_id="test-asset-123",
        filename="test_audio.wav",
        content_type="audio/wav",
    )


@pytest.fixture
def completed_result():
    """Standard completed inference result."""
    return {
        "job_id": "job-123",
        "status": "completed",
        "result": {"intelligibility_score": 0.85},
        "created_at": "2024-01-01T00:00:00Z",
        "completed_at": "2024-01-01T00:01:00Z",
    }


class TestRunInference:

    def test_run_inference_short_file(
        self, mock_session, mock_client, mock_asset, completed_result
    ):
        """Test run_inference with short file - simple upload + predict flow."""
        # Setup mocks
        mock_session._get_audio_duration.return_value = 30.0  # Under threshold
        mock_session.upload.return_value = mock_asset
        mock_client.predict.return_value = completed_result

        # Call run_inference
        result = run_inference(mock_session, mock_client, "short.wav")

        # Verify
        assert result["status"] == "completed"
        assert result["result"]["intelligibility_score"] == 0.85
        mock_session.upload.assert_called_once_with("short.wav")
        mock_client.predict.assert_called_once()
        mock_session.upload_chunked.assert_not_called()  # Should not chunk

    def test_run_inference_long_file_auto_chunks(self, mock_session, mock_client):
        """Test run_inference with long file - auto-chunks and aggregates."""
        # Setup mocks - file exceeds threshold
        mock_session._get_audio_duration.return_value = 120.0  # Over 50s threshold

        # Create mock chunks
        chunk1 = Asset(asset_id="chunk-1", filename="test_chunk_0.wav", content_type="audio/wav")
        chunk2 = Asset(asset_id="chunk-2", filename="test_chunk_1.wav", content_type="audio/wav")
        mock_session.upload_chunked.return_value = [chunk1, chunk2]

        # Mock submit/poll for chunks
        mock_client.submit.side_effect = ["job-1", "job-2"]
        mock_client.poll.side_effect = [
            {
                "job_id": "job-1",
                "status": "completed",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
            },
            {
                "job_id": "job-2",
                "status": "completed",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:10Z",
                "completed_at": "2024-01-01T00:01:10Z",
            },
        ]

        # Call run_inference
        result = run_inference(mock_session, mock_client, "long.wav")

        # Verify chunking happened
        mock_session.upload_chunked.assert_called_once()
        assert mock_client.submit.call_count == 2
        assert mock_client.poll.call_count == 2

        # Verify aggregation happened
        assert result["status"] == "completed"
        assert result["result"]["intelligibility_score"] == pytest.approx(
            0.85
        )  # Mean of 0.80 and 0.90
        assert result["result"]["num_chunks"] == 2
        assert result["result"]["chunk_scores"] == [0.80, 0.90]

    def test_run_inference_without_score_key_skips_chunking(
        self, mock_session, mock_client, mock_asset, completed_result
    ):
        """Test run_inference skips chunking when client has no score_key.

        Models without score_key (like TranscribeClient) return non-numeric results
        that can't be meaningfully aggregated, so chunking is skipped even for long files.
        """
        # Setup mocks
        mock_session._get_audio_duration.return_value = 120.0  # Over threshold
        mock_session.upload.return_value = mock_asset
        mock_client.predict.return_value = completed_result
        mock_client.score_key = None  # No score_key → no aggregator → no chunking

        # Call run_inference
        result = run_inference(mock_session, mock_client, "long.wav")

        # Verify - should NOT chunk even though file is long
        assert result["status"] == "completed"
        mock_session.upload.assert_called_once()  # Regular upload
        mock_session.upload_chunked.assert_not_called()  # No chunking
        mock_client.predict.assert_called_once()


class TestRunChunkedInference:

    def test_run_chunked_inference_single_chunk(self, mock_session, mock_client, completed_result):
        """Test run_chunked_inference with single chunk - returns list."""
        # Setup mocks
        chunk = Asset(asset_id="chunk-1", filename="test_chunk_0.wav", content_type="audio/wav")
        mock_session.upload_chunked.return_value = [chunk]
        mock_client.predict.return_value = completed_result

        # Call run_chunked_inference
        results = run_chunked_inference(mock_session, mock_client, "short.wav", chunk_duration=50.0)

        # Verify
        assert isinstance(results, list)
        assert len(results) == 1
        assert results[0]["_chunk_index"] == 0
        assert results[0]["_total_chunks"] == 1
        assert results[0]["result"]["intelligibility_score"] == 0.85

    def test_run_chunked_inference_multiple_chunks(self, mock_session, mock_client):
        """Test run_chunked_inference with multiple chunks - includes metadata."""
        # Setup mocks
        chunk1 = Asset(asset_id="chunk-1", filename="test_chunk_0.wav", content_type="audio/wav")
        chunk2 = Asset(asset_id="chunk-2", filename="test_chunk_1.wav", content_type="audio/wav")
        chunk3 = Asset(asset_id="chunk-3", filename="test_chunk_2.wav", content_type="audio/wav")
        mock_session.upload_chunked.return_value = [chunk1, chunk2, chunk3]

        # Mock submit/poll for chunks
        mock_client.submit.side_effect = ["job-1", "job-2", "job-3"]
        mock_client.poll.side_effect = [
            {
                "job_id": "job-1",
                "status": "completed",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
            },
            {
                "job_id": "job-2",
                "status": "completed",
                "result": {"intelligibility_score": 0.85},
                "created_at": "2024-01-01T00:00:10Z",
                "completed_at": "2024-01-01T00:01:10Z",
            },
            {
                "job_id": "job-3",
                "status": "completed",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:20Z",
                "completed_at": "2024-01-01T00:01:20Z",
            },
        ]

        # Call run_chunked_inference
        results = run_chunked_inference(mock_session, mock_client, "long.wav", chunk_duration=50.0)

        # Verify
        assert len(results) == 3
        assert results[0]["_chunk_index"] == 0
        assert results[1]["_chunk_index"] == 1
        assert results[2]["_chunk_index"] == 2
        assert all(r["_total_chunks"] == 3 for r in results)
        assert mock_client.submit.call_count == 3
        assert mock_client.poll.call_count == 3

    def test_run_chunked_inference_partial_failure(self, mock_session, mock_client):
        """Test run_chunked_inference handles some chunk failures."""
        # Setup mocks
        chunk1 = Asset(asset_id="chunk-1", filename="test_chunk_0.wav", content_type="audio/wav")
        chunk2 = Asset(asset_id="chunk-2", filename="test_chunk_1.wav", content_type="audio/wav")
        chunk3 = Asset(asset_id="chunk-3", filename="test_chunk_2.wav", content_type="audio/wav")
        mock_session.upload_chunked.return_value = [chunk1, chunk2, chunk3]

        # Mock submit/poll - chunk 1 fails
        mock_client.submit.side_effect = ["job-1", "job-2", "job-3"]
        mock_client.poll.side_effect = [
            {
                "job_id": "job-1",
                "status": "completed",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
            },
            InferenceError("Chunk 2 failed"),  # Chunk 1 (0-indexed) fails
            {
                "job_id": "job-3",
                "status": "completed",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:20Z",
                "completed_at": "2024-01-01T00:01:20Z",
            },
        ]

        # Call run_chunked_inference
        results = run_chunked_inference(mock_session, mock_client, "long.wav", chunk_duration=50.0)

        # Verify - should have 2 successful results
        assert len(results) == 2
        assert results[0]["result"]["intelligibility_score"] == 0.80
        assert results[1]["result"]["intelligibility_score"] == 0.90

        # Last result should contain failed chunks info
        assert "_failed_chunks" in results[-1]
        assert 1 in results[-1]["_failed_chunks"]  # Chunk index 1 failed

    def test_run_chunked_inference_all_fail_raises(self, mock_session, mock_client):
        """Test run_chunked_inference raises when all chunks fail."""
        # Setup mocks
        chunk1 = Asset(asset_id="chunk-1", filename="test_chunk_0.wav", content_type="audio/wav")
        chunk2 = Asset(asset_id="chunk-2", filename="test_chunk_1.wav", content_type="audio/wav")
        mock_session.upload_chunked.return_value = [chunk1, chunk2]

        # Mock submit/poll - all fail
        mock_client.submit.side_effect = ["job-1", "job-2"]
        mock_client.poll.side_effect = [
            InferenceError("Chunk 1 failed"),
            InferenceError("Chunk 2 failed"),
        ]

        # Verify InferenceError is raised
        with pytest.raises(InferenceError, match="All 2 chunks failed"):
            run_chunked_inference(mock_session, mock_client, "long.wav", chunk_duration=50.0)


class TestAggregateResults:

    def test_aggregate_results_mean(self):
        """Test aggregate_results with default mean aggregation."""
        # Create chunk results
        results = [
            {
                "job_id": "job-1",
                "status": "completed",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
                "_chunk_index": 0,
                "_total_chunks": 3,
            },
            {
                "job_id": "job-2",
                "status": "completed",
                "result": {"intelligibility_score": 0.85},
                "created_at": "2024-01-01T00:00:10Z",
                "completed_at": "2024-01-01T00:01:10Z",
                "_chunk_index": 1,
                "_total_chunks": 3,
            },
            {
                "job_id": "job-3",
                "status": "completed",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:20Z",
                "completed_at": "2024-01-01T00:01:20Z",
                "_chunk_index": 2,
                "_total_chunks": 3,
            },
        ]

        # Aggregate
        aggregated = aggregate_results(results, "intelligibility_score")

        # Verify
        assert aggregated["status"] == "completed"
        assert aggregated["result"]["intelligibility_score"] == pytest.approx(0.85)  # Mean
        assert aggregated["result"]["num_chunks"] == 3
        assert aggregated["result"]["chunk_scores"] == [0.80, 0.85, 0.90]
        assert aggregated["result"]["aggregation_method"] == "mean"
        assert "std_dev" in aggregated["result"]

    def test_aggregate_results_custom_reduce(self):
        """Test aggregate_results with custom reduce function."""
        # Create chunk results
        results = [
            {
                "job_id": "job-1",
                "status": "completed",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
                "_chunk_index": 0,
                "_total_chunks": 3,
            },
            {
                "job_id": "job-2",
                "status": "completed",
                "result": {"intelligibility_score": 0.85},
                "created_at": "2024-01-01T00:00:10Z",
                "completed_at": "2024-01-01T00:01:10Z",
                "_chunk_index": 1,
                "_total_chunks": 3,
            },
            {
                "job_id": "job-3",
                "status": "completed",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:20Z",
                "completed_at": "2024-01-01T00:01:20Z",
                "_chunk_index": 2,
                "_total_chunks": 3,
            },
        ]

        # Aggregate with median
        aggregated = aggregate_results(
            results, "intelligibility_score", reduce_fn=statistics.median
        )

        # Verify
        assert aggregated["result"]["intelligibility_score"] == 0.85  # Median
        assert aggregated["result"]["aggregation_method"] == "custom"

    def test_aggregate_results_with_failed_chunks(self):
        """Test aggregate_results with partial status when chunks failed."""
        # Create chunk results with failed chunks info
        results = [
            {
                "job_id": "job-1",
                "status": "completed",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
                "_chunk_index": 0,
                "_total_chunks": 3,
            },
            {
                "job_id": "job-3",
                "status": "completed",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:20Z",
                "completed_at": "2024-01-01T00:01:20Z",
                "_chunk_index": 2,
                "_total_chunks": 3,
                "_failed_chunks": [1],  # Chunk 1 failed
            },
        ]

        # Aggregate
        aggregated = aggregate_results(results, "intelligibility_score")

        # Verify
        assert aggregated["status"] == "partial"  # Not all chunks succeeded
        assert aggregated["result"]["failed_chunks"] == [1]
        assert aggregated["result"]["success_rate"] == "2/3"
        assert aggregated["result"]["total_chunks"] == 3
        assert aggregated["result"]["num_chunks"] == 2  # Only successful chunks


class TestDefaultAggregator:

    def test_default_aggregator_for_client(self, mock_client):
        """Test _default_aggregator_for returns correct aggregator."""
        # Setup client with score_key
        mock_client.score_key = "intelligibility_score"

        # Get aggregator
        aggregator = _default_aggregator_for(mock_client)

        # Verify aggregator works
        assert aggregator is not None

        # Test the aggregator
        results = [
            {
                "job_id": "job-1",
                "result": {"intelligibility_score": 0.80},
                "created_at": "2024-01-01T00:00:00Z",
                "completed_at": "2024-01-01T00:01:00Z",
                "_chunk_index": 0,
                "_total_chunks": 2,
            },
            {
                "job_id": "job-2",
                "result": {"intelligibility_score": 0.90},
                "created_at": "2024-01-01T00:00:10Z",
                "completed_at": "2024-01-01T00:01:10Z",
                "_chunk_index": 1,
                "_total_chunks": 2,
            },
        ]

        aggregated = aggregator(results)
        assert aggregated["result"]["intelligibility_score"] == pytest.approx(0.85)

    def test_default_aggregator_none_without_score_key(self, mock_client):
        """Test _default_aggregator_for returns None without score_key."""
        # Setup client without score_key
        mock_client.score_key = None

        # Get aggregator
        aggregator = _default_aggregator_for(mock_client)

        # Verify
        assert aggregator is None
