import pytest
from unittest.mock import patch, AsyncMock, MagicMock
import asyncio

from dracula.retry import is_retryable, retry, async_retry
from dracula.exceptions import (
    ChatException,
    InvalidAPIKeyException,
    ValidationException,
)


class DummyException(Exception):
    """Dummy exception class used for testing purposes."""

    pass


# ==========================================
# 1. Tests for is_retryable() function
# ==========================================


def test_is_retryable_non_retryable_exceptions():
    """API Key and Validation errors should NEVER be retried."""
    assert is_retryable(InvalidAPIKeyException("Bad key")) is False
    assert is_retryable(ValidationException("Bad input")) is False


def test_is_retryable_status_codes():
    """Exceptions containing specific HTTP status codes should be retried."""
    assert is_retryable(Exception("Error 429: Too Many Requests")) is True
    assert is_retryable(Exception("Server returned 503 unavailable")) is True
    assert is_retryable(Exception("Gateway Timeout 504")) is True


def test_is_retryable_network_errors():
    """Generic network connection errors should be retried."""
    assert is_retryable(ConnectionError("Connection reset by peer")) is True
    assert is_retryable(TimeoutError("Read timeout")) is True


def test_is_retryable_unhandled_exceptions():
    """Random unhandled exceptions should NOT be retried."""
    assert is_retryable(ValueError("Some random math error")) is False
    assert is_retryable(KeyError("Missing dictionary key")) is False


# ==========================================
# 2. Tests for synchronous retry() function
# ==========================================


@patch("time.sleep")
def test_retry_success_first_try(mock_sleep):
    """If the function succeeds on the first try, sleep should not be called."""
    mock_func = MagicMock(return_value="Success!")

    result = retry(mock_func, max_retries=3)

    assert result == "Success!"
    assert mock_func.call_count == 1
    mock_sleep.assert_not_called()


@patch("time.sleep")
def test_retry_success_after_failures(mock_sleep):
    """Function should succeed on the 3rd try after 2 failures, with correct delays."""
    mock_func = MagicMock(
        side_effect=[
            Exception("503 Service Unavailable"),
            Exception("503 Service Unavailable"),
            "Success!",
        ]
    )

    result = retry(mock_func, max_retries=3, retry_delay=1.0)

    assert result == "Success!"
    assert mock_func.call_count == 3
    assert mock_sleep.call_count == 2

    mock_sleep.assert_any_call(1.0)
    mock_sleep.assert_any_call(2.0)


@patch("time.sleep")
def test_retry_exhausted(mock_sleep):
    """If all retries are exhausted, it should raise a ChatException."""
    mock_func = MagicMock(side_effect=Exception("500 Internal Error"))

    with pytest.raises(ChatException) as exc_info:
        retry(mock_func, max_retries=2, retry_delay=1.0)

    assert "Failed after 2 retries" in str(exc_info.value)
    assert mock_func.call_count == 3  # Initial try + 2 retries
    assert mock_sleep.call_count == 2


def test_retry_non_retryable_fast_fail():
    """If an API Key error occurs, it should fail immediately without retrying."""
    mock_func = MagicMock(side_effect=InvalidAPIKeyException("Invalid Key"))

    with pytest.raises(InvalidAPIKeyException):
        retry(mock_func, max_retries=5)

    assert mock_func.call_count == 1  # Only called once, no retries triggered


# ==========================================
# 3. Tests for async_retry() function
# ==========================================


@pytest.mark.asyncio
@patch("asyncio.sleep", new_callable=AsyncMock)
async def test_async_retry_success(mock_sleep):
    """If the async function succeeds on the first try."""

    async def mock_coroutine():
        return "Async Success!"

    result = await async_retry(mock_coroutine)

    assert result == "Async Success!"
    mock_sleep.assert_not_called()


@pytest.mark.asyncio
@patch("asyncio.sleep", new_callable=AsyncMock)
async def test_async_retry_fails_then_succeeds(mock_sleep):
    """If the async function fails once, then succeeds on the second try."""
    call_count = 0

    async def mock_coroutine():
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise Exception("Rate limit 429")
        return "Recovered!"

    result = await async_retry(mock_coroutine, max_retries=2, retry_delay=1.0)

    assert result == "Recovered!"
    assert call_count == 2
    mock_sleep.assert_awaited_once_with(1.0)


@pytest.mark.asyncio
@patch("asyncio.sleep", new_callable=AsyncMock)
async def test_async_retry_exhausted(mock_sleep):
    """If the async function exhausts all retries, it should raise a ChatException."""

    async def mock_coroutine():
        raise Exception("502 Bad Gateway")

    with pytest.raises(ChatException):
        await async_retry(mock_coroutine, max_retries=1)
