import json
from collections.abc import Callable
from unittest.mock import patch

import pytest
from src.worker.clients.step_functions import (
    StaleTaskToken,
    send_task_failure,
    send_task_heartbeat,
    send_task_success,
)
from src.worker.connectors.aws import stepfunctions


class TestSendTaskSuccess:
    def test_serializes_output_as_json(self) -> None:
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            send_task_success(task_token="abc123", output={"valid": True, "findings": []})

        mock_sfn.send_task_success.assert_called_once_with(
            taskToken="abc123",
            output=json.dumps({"valid": True, "findings": []}),
        )


class TestSendTaskFailure:
    def test_passes_error_and_cause(self) -> None:
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            send_task_failure(
                task_token="abc123",
                error="ValidationError",
                cause="stereo reference not found",
            )

        mock_sfn.send_task_failure.assert_called_once_with(
            taskToken="abc123",
            error="ValidationError",
            cause="stereo reference not found",
        )


_STALE_TOKEN_EXCEPTIONS = (
    pytest.param(stepfunctions.exceptions.TaskDoesNotExist, id="task_does_not_exist"),
    pytest.param(stepfunctions.exceptions.TaskTimedOut, id="task_timed_out"),
)


def _raise(exc_class: type[Exception]) -> Callable[..., None]:
    def _side_effect(*_: object, **__: object) -> None:
        raise exc_class({"Error": {"Code": exc_class.__name__, "Message": "stale"}}, "SendTask")

    return _side_effect


class TestStaleTokenTranslation:
    @pytest.mark.parametrize("exc_class", _STALE_TOKEN_EXCEPTIONS)
    def test_send_task_success_translates_to_stale(self, exc_class: type[Exception]) -> None:
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            mock_sfn.send_task_success.side_effect = _raise(exc_class)
            with pytest.raises(StaleTaskToken):
                send_task_success(task_token="abc", output={})

    @pytest.mark.parametrize("exc_class", _STALE_TOKEN_EXCEPTIONS)
    def test_send_task_failure_translates_to_stale(self, exc_class: type[Exception]) -> None:
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            mock_sfn.send_task_failure.side_effect = _raise(exc_class)
            with pytest.raises(StaleTaskToken):
                send_task_failure(task_token="abc", error="x", cause="y")

    @pytest.mark.parametrize("exc_class", _STALE_TOKEN_EXCEPTIONS)
    def test_send_task_heartbeat_translates_to_stale(self, exc_class: type[Exception]) -> None:
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            mock_sfn.send_task_heartbeat.side_effect = _raise(exc_class)
            with pytest.raises(StaleTaskToken):
                send_task_heartbeat(task_token="abc")

    def test_invalid_token_propagates_unwrapped(self) -> None:
        # InvalidToken is structural (malformed token), not stale (lifecycle).
        # It signals a bug or config drift and should surface, not be silently
        # treated as designed recovery.
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            mock_sfn.send_task_heartbeat.side_effect = _raise(stepfunctions.exceptions.InvalidToken)
            with pytest.raises(stepfunctions.exceptions.InvalidToken):
                send_task_heartbeat(task_token="abc")

    def test_non_stale_exception_propagates(self) -> None:
        with patch("src.worker.clients.step_functions.stepfunctions") as mock_sfn:
            mock_sfn.send_task_heartbeat.side_effect = RuntimeError("throttled")
            with pytest.raises(RuntimeError, match="throttled"):
                send_task_heartbeat(task_token="abc")