from unittest.mock import MagicMock, patch

import pytest

from redenlab_extract import (
    InferenceClient,
    TranscribeClient,
    intelligibility_client,
    naturalness_client,
)
from redenlab_extract.asset import Asset
from redenlab_extract.exceptions import (
    APIError,
    InferenceError,
    TimeoutError,
    ValidationError,
)
from redenlab_extract.payload import IntelligibilityPayloadBuilder, TranscribePayloadBuilder
from redenlab_extract.session import RedenLabSession


class TestInferenceClient:

    @patch("redenlab_extract.client.validate_model_name_format")
    @patch("redenlab_extract.session.get_merged_config")
    def test_client_init_with_session(self, mock_config, mock_validate_model_format):
        """Test client initialisation logic"""

        # setup mocks
        mock_config.return_value = {
            "api_key": "sk_test_12345",
            "base_url": "https://api.test.example.com",
        }
        mock_validate_model_format.return_value = "als-intelligibility-mtpa"

        # Test - create session and client
        session = RedenLabSession(api_key="sk_test_12345", base_url="https://api.test.example.com")
        client = InferenceClient(session, model_name="als-intelligibility-mtpa")

        # Verify client state
        assert client._session.api_key == "sk_test_12345"
        assert client.model_name == "als-intelligibility-mtpa"
        assert client._session.base_url == "https://api.test.example.com"

    def test_client_init_raises_if_model_name_format_is_incorrect(self):
        "Test client checks model name format"
        session = RedenLabSession(api_key="sk_test_12345", base_url="https://api.test.example.com")

        with pytest.raises(ValidationError):
            InferenceClient(session, model_name="Abs")

    def test_client_init_raises_if_no_payload_template_exist(self):
        "Test client checks if the template for payload exist"
        session = RedenLabSession(api_key="sk_test_12345", base_url="https://api.test.example.com")

        with pytest.raises(ValueError):
            InferenceClient(session, model_name="als-intelligibility")

    @patch("redenlab_extract.client.api.submit_inference_job")
    def test_submit_returns_job_id(self, mock_submit, mock_session, mock_asset):
        """Test that submit() returns a job_id string."""
        # Setup mock response
        mock_submit.return_value = {"job_id": "job-123", "status": "submitted"}

        # Create client
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Call submit
        job_id = client.submit(mock_asset)

        # Verify
        assert job_id == "job-123"
        assert isinstance(job_id, str)
        mock_submit.assert_called_once()

    @patch("redenlab_extract.client.api.get_job_status")
    def test_poll_completed_returns_result(self, mock_status, mock_session):
        """Test that poll() returns result when job completes."""
        # Setup mock - return completed on first call
        mock_status.return_value = {
            "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",
        }

        # Create client
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Call poll
        result = client.poll(job_id="job-123")

        # Verify
        assert result["status"] == "completed"
        assert result["result"]["intelligibility_score"] == 0.85
        assert result["job_id"] == "job-123"

    @patch("redenlab_extract.client.api.get_job_status")
    def test_poll_failed_raises_inference_error(self, mock_status, mock_session):
        """Test that poll() raises InferenceError when job fails."""
        # Setup mock - return failed status
        mock_status.return_value = {
            "job_id": "job-123",
            "status": "failed",
            "error": "Audio duration exceeds maximum allowed",
        }

        # Create client
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Verify InferenceError is raised
        with pytest.raises(InferenceError, match="Audio duration exceeds maximum allowed"):
            client.poll(job_id="job-123")

    @patch("redenlab_extract.client.api.get_job_status")
    def test_poll_timeout_raises_timeout_error(self, mock_status, mock_session):
        """Test that poll() raises TimeoutError when exceeds timeout."""
        # Setup mock - always return processing status
        mock_status.return_value = {
            "job_id": "job-123",
            "status": "processing",
        }

        # Create client with very short timeout
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Verify TimeoutError is raised (use very short timeout for test)
        with pytest.raises(TimeoutError):
            client.poll(job_id="job-123", timeout=1)

    @patch("redenlab_extract.client.api.get_job_status")
    @patch("redenlab_extract.client.api.submit_inference_job")
    def test_predict_submits_and_polls(self, mock_submit, mock_status, mock_session, mock_asset):
        """Test that predict() = submit() + poll()."""
        # Setup mocks
        mock_submit.return_value = {"job_id": "job-123", "status": "submitted"}
        mock_status.return_value = {
            "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",
        }

        # Create client
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Call predict with wait_for_upload=False to test direct flow
        result = client.predict(mock_asset, wait_for_upload=False)

        # Verify both submit and poll were called
        assert result["status"] == "completed"
        assert result["result"]["intelligibility_score"] == 0.85
        mock_submit.assert_called_once()
        mock_status.assert_called()

    @patch("redenlab_extract.client.api.get_job_status")
    @patch("redenlab_extract.client.api.submit_inference_job")
    def test_predict_with_progress_callback(
        self, mock_submit, mock_status, mock_session, mock_asset
    ):
        """Test that predict() with progress_callback receives status updates."""
        # Setup mocks - simulate multiple polling iterations
        mock_submit.return_value = {"job_id": "job-123", "status": "submitted"}

        # Create a side_effect that returns processing, then completed
        status_responses = [
            {"job_id": "job-123", "status": "processing"},
            {"job_id": "job-123", "status": "processing"},
            {
                "job_id": "job-123",
                "status": "completed",
                "result": {"intelligibility_score": 0.85},
            },
        ]
        mock_status.side_effect = status_responses

        # Create client and asset
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Track callback invocations
        callback_statuses = []

        def progress_callback(status):
            callback_statuses.append(status["status"])

        # Call predict with callback
        result = client.predict(
            mock_asset, progress_callback=progress_callback, wait_for_upload=False
        )

        # Verify callback was called
        assert result["status"] == "completed"
        assert len(callback_statuses) >= 1  # At least one callback
        assert "processing" in callback_statuses or "completed" in callback_statuses
        assert mock_status.call_count == 3

    @patch("redenlab_extract.client.api.submit_inference_job")
    def test_submit_with_upload_wait_retries(self, mock_submit, mock_session, mock_asset):
        """Test that _submit_with_upload_wait() retries on 409 upload_pending."""
        # Setup mock - fail twice with upload pending, then succeed
        mock_submit.side_effect = [
            APIError("Asset upload status is still upload_pending"),
            APIError("Asset upload status is still upload_pending"),
            {"job_id": "job-123", "status": "submitted"},
        ]

        # Create client and asset
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Call _submit_with_upload_wait
        job_id = client._submit_with_upload_wait(mock_asset, max_retries=3, initial_wait=0.1)

        # Verify retries worked
        assert job_id == "job-123"
        assert mock_submit.call_count == 3

    @patch("redenlab_extract.client.api.check_health")
    def test_health_returns_status(self, mock_health, mock_session):
        """Test that health() returns backend status."""
        # Setup mock
        mock_health.return_value = {
            "status": "ready",
            "ready": True,
            "response_time": 1.2,
            "is_warming": False,
        }

        # Create client
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Call health
        health = client.health()

        # Verify
        assert health["status"] == "ready"
        assert health["ready"] is True
        assert health["is_warming"] is False
        mock_health.assert_called_once()

    @patch("redenlab_extract.client.api.check_health")
    def test_warmup_already_ready_skips(self, mock_health, mock_session):
        """Test that warmup() skips if backend is already ready."""
        # Setup mock - backend is already ready
        mock_health.return_value = {
            "status": "ready",
            "ready": True,
            "response_time": 1.2,
            "is_warming": False,
        }

        # Create client
        client = InferenceClient(mock_session, model_name="als-intelligibility-mtpa")

        # Call warmup
        result = client.warmup(show_progress=False)

        # Verify warmup was skipped
        assert result["status"] == "already_ready"
        assert "already ready" in result["message"].lower()
        mock_health.assert_called_once()


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


@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",
    )


class TestTranscribeClient:

    def test_transcribe_client_with_language_code(self, mock_session):
        """Test that TranscribeClient stores default language_code."""
        client = TranscribeClient(session=mock_session, language_code="en-US")
        assert client.language_code == "en-US"

    @patch("redenlab_extract.client.api.submit_inference_job")
    def test_transcribe_client_submit_uses_language(self, mock_submit, mock_session, mock_asset):
        """Test that TranscribeClient submit() includes language_code in payload."""
        # Setup mock
        mock_submit.return_value = {"job_id": "job-123", "status": "submitted"}

        # Create client with default language_code
        client = TranscribeClient(session=mock_session, language_code="en-US")

        # Submit without explicit language_code
        job_id = client.submit(mock_asset)

        # Verify submit was called and payload includes language_code
        assert job_id == "job-123"
        mock_submit.assert_called_once()

        # Check that the payload builder received language_code
        call_args = mock_submit.call_args
        payload = call_args.kwargs["payload"]
        assert "language_code" in payload
        assert payload["language_code"] == "en-US"

    @patch("redenlab_extract.client.api.submit_inference_job")
    def test_transcribe_client_submit_overrides_language(
        self, mock_submit, mock_session, mock_asset
    ):
        """Test that TranscribeClient submit() can override default language_code."""
        # Setup mock
        mock_submit.return_value = {"job_id": "job-123", "status": "submitted"}

        # Create client with default language_code
        client = TranscribeClient(session=mock_session, language_code="en-US")

        # Submit with explicit language_code override
        job_id = client.submit(mock_asset, language_code="es-ES")

        # Verify submit was called with overridden language_code
        assert job_id == "job-123"
        mock_submit.assert_called_once()

        # Check that the payload builder received overridden language_code
        call_args = mock_submit.call_args
        payload = call_args.kwargs["payload"]
        assert "language_code" in payload
        assert payload["language_code"] == "es-ES"

    def test_transcribe_client_initialises_with_transcribe_builder(self, mock_session):
        """Test that TranscribeClient uses the TranscribePayloadBuilder."""
        client = TranscribeClient(session=mock_session, language_code="en-US")
        assert isinstance(client._payload_builder, TranscribePayloadBuilder)
        assert client.language_code == "en-US"


class TestIntelligibilityClient:

    def test_intelligibility_client_valid_model(self, mock_session):
        """Test that intelligibility_client returns InferenceClient for valid model."""
        client = intelligibility_client(session=mock_session, model_name="als-intelligibility-v1")
        assert isinstance(client, InferenceClient)
        assert client.model_name == "als-intelligibility-v1"

    def test_intelligibility_client_invalid_model_raises(self, mock_session):
        """Test that intelligibility_client raises ValidationError for invalid model."""
        with pytest.raises(ValidationError, match="Invalid intelligibility model"):
            intelligibility_client(session=mock_session, model_name="invalid-model")

    def test_intelligibility_client_sets_score_key(self, mock_session):
        """Test that intelligibility_client sets score_key for aggregation."""
        client = intelligibility_client(session=mock_session, model_name="als-intelligibility-v1")
        assert client.score_key == "intelligibility_score"

    def test_intelligibility_client_maps_alias(self, mock_session):
        """Test that intelligibility_client maps v1/v2 aliases to backend names."""
        # Test v1 mapping
        client_v1 = intelligibility_client(
            session=mock_session, model_name="als-intelligibility-v1"
        )
        assert client_v1.model_name == "als-intelligibility-v1"  # User-facing alias (no leakage!)
        assert (
            client_v1._backend_model_name == "als-intelligibility-mtpa"
        )  # Backend name (internal)

        # Test v2 mapping
        client_v2 = intelligibility_client(
            session=mock_session, model_name="als-intelligibility-v2"
        )
        assert client_v2.model_name == "als-intelligibility-v2"  # User-facing alias (no leakage!)
        assert (
            client_v2._backend_model_name == "als-intelligibility-eals"
        )  # Backend name (internal)

    def test_intelligibility_client_initialises_with_intelligibility_builder(self, mock_session):
        """Test that intelligibility_client factory returns InferenceClient with IntelligibilityPayloadBuilder."""
        client = intelligibility_client(session=mock_session, model_name="als-intelligibility-v1")
        assert isinstance(client, InferenceClient)
        assert isinstance(client._payload_builder, IntelligibilityPayloadBuilder)
        assert client.model_name == "als-intelligibility-v1"
        assert client.score_key == "intelligibility_score"


class TestNaturalnessClient:

    def test_naturalness_client_valid_model(self, mock_session):
        """Test that naturalness_client returns InferenceClient for valid model."""
        client = naturalness_client(session=mock_session, model_name="als-naturalness-v1")
        assert isinstance(client, InferenceClient)
        assert client.model_name == "als-naturalness-v1"

    def test_naturalness_client_invalid_model_raises(self, mock_session):
        """Test that naturalness_client raises ValidationError for invalid model."""
        with pytest.raises(ValidationError, match="Invalid naturalness model"):
            naturalness_client(session=mock_session, model_name="invalid-model")

    def test_naturalness_client_sets_score_key(self, mock_session):
        """Test that naturalness_client sets score_key for aggregation."""
        client = naturalness_client(session=mock_session, model_name="als-naturalness-v1")
        assert client.score_key == "naturalness_score"

    def test_naturalness_client_maps_alias(self, mock_session):
        """Test that naturalness_client maps v1/v2 aliases to backend names."""
        # Test v1 mapping
        client_v1 = naturalness_client(session=mock_session, model_name="als-naturalness-v1")
        assert client_v1.model_name == "als-naturalness-v1"  # User-facing alias (no leakage!)
        assert client_v1._backend_model_name == "als-naturalness-mtpa"  # Backend name (internal)

        # Test v2 mapping
        client_v2 = naturalness_client(session=mock_session, model_name="als-naturalness-v2")
        assert client_v2.model_name == "als-naturalness-v2"  # User-facing alias (no leakage!)
        assert client_v2._backend_model_name == "als-naturalness-eals"  # Backend name (internal)
