"""Behavioral tests for retry/timeout logic in core.py.

Verifies that retry decorators actually retry on the right errors, give up
on non-retryable errors, and eventually time out -- without depending on
specific wait/stop numeric values.  All wait/stop times are patched to be
very short (0.01s / 0.1s) so tests finish quickly.

Targets
-------
1. _retry_on_unavailable  -- retries UNAVAILABLE, gives up on others
2. _ensure_workspace       -- retries UNAVAILABLE via same retry logic
3. _wait_for_trialbatch    -- retries RuntimeError, gives up on ValueError
4. StudyWrapper.wait       -- retries when is_trial_finished returns False
"""

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import grpc
import pytest
from tenacity import retry, retry_if_exception, stop_after_delay, wait_fixed

from aiauto.grpc_helper import AIAutoRpcError, _should_retry_rpc_error


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
class FakeRpcError(grpc.RpcError):
    """Fake grpc.RpcError with code() and details() for testing."""

    def __init__(self, status_code, details_msg=""):
        self._code = status_code
        self._details = details_msg

    def code(self):
        return self._code

    def details(self):
        return self._details


def _short_wait(*_args, **_kwargs):
    """Return a very short wait_fixed so retries happen instantly."""
    return wait_fixed(0.01)


def _short_stop(*_args, **_kwargs):
    """Return a very short stop_after_delay so timeout triggers quickly."""
    return stop_after_delay(0.1)


def _patch_retry_timing():
    """Context manager stack for patching wait/stop to short values.

    Returns nested context managers compatible with Python 3.8.
    """
    return patch("aiauto.core.wait_fixed", _short_wait)


def _patch_retry_timing_stop():
    return patch("aiauto.core.stop_after_delay", _short_stop)


def _make_study_wrapper(core_mod):
    """Build a minimal StudyWrapper with mocked internals."""
    sw = object.__new__(core_mod.StudyWrapper)
    sw._controller = MagicMock()
    sw._storage = MagicMock()
    sw.study_name = "test-study"
    sw._study = MagicMock()
    sw._last_trialbatch_name = "tb-test"
    return sw


# ---------------------------------------------------------------------------
# 1. _retry_on_unavailable behavioral tests
# ---------------------------------------------------------------------------
class TestRetryOnUnavailableBehavior:
    """_retry_on_unavailable: retry on UNAVAILABLE, fail fast on others."""

    def test_retries_unavailable_then_succeeds(self):
        """mock: 2x UNAVAILABLE then success -> final success, 3 calls."""
        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                from aiauto.core import _retry_on_unavailable

                call_count = 0

                @_retry_on_unavailable()
                def _rpc_call():
                    nonlocal call_count
                    call_count += 1
                    if call_count <= 2:
                        raise AIAutoRpcError("unavailable", "service starting")
                    return "ok"

                result = _rpc_call()
                assert result == "ok"
                assert call_count == 3

    def test_timeout_on_persistent_unavailable(self):
        """mock keeps returning UNAVAILABLE -> short timeout triggers error."""
        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                from aiauto.core import _retry_on_unavailable

                @_retry_on_unavailable()
                def _rpc_call():
                    raise AIAutoRpcError("unavailable", "still down")

                with pytest.raises(AIAutoRpcError, match="still down"):
                    _rpc_call()

    def test_non_retryable_error_fails_immediately(self):
        """INVALID_ARGUMENT (non-retryable) -> no retry, immediate error."""
        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                from aiauto.core import _retry_on_unavailable

                call_count = 0

                @_retry_on_unavailable()
                def _rpc_call():
                    nonlocal call_count
                    call_count += 1
                    raise AIAutoRpcError("invalid_argument", "bad request")

                with pytest.raises(AIAutoRpcError, match="bad request"):
                    _rpc_call()

                assert call_count == 1, "Non-retryable error should not trigger retry"


# ---------------------------------------------------------------------------
# 2. _ensure_workspace retry behavioral tests
#
# _ensure_workspace is a closure inside AIAutoController.__init__ that uses
# the same retry pattern (retry_if_exception(_should_retry_rpc_error) +
# wait_fixed + stop_after_delay).  We verify the *pattern* behaves correctly
# by constructing an equivalent decorated function with short timeouts.
# ---------------------------------------------------------------------------
class TestEnsureWorkspaceRetryBehavior:
    """_ensure_workspace: retries UNAVAILABLE from EnsureWorkspace gRPC call."""

    def test_ensure_workspace_retries_then_succeeds(self):
        """1x UNAVAILABLE then success -> 2 calls total."""
        call_count = 0
        success_response = SimpleNamespace(
            journal_grpc_storage_proxy_host_external="https://host/grpc/user1"
        )

        @retry(
            stop=stop_after_delay(0.1),
            wait=wait_fixed(0.01),
            retry=retry_if_exception(_should_retry_rpc_error),
            reraise=True,
        )
        def _ensure_workspace():
            nonlocal call_count
            call_count += 1
            if call_count <= 1:
                raise AIAutoRpcError("unavailable", "workspace not ready")
            return success_response

        result = _ensure_workspace()
        assert result.journal_grpc_storage_proxy_host_external == "https://host/grpc/user1"
        assert call_count == 2

    def test_ensure_workspace_timeout_on_persistent_unavailable(self):
        """Persistent UNAVAILABLE -> timeout error."""

        @retry(
            stop=stop_after_delay(0.1),
            wait=wait_fixed(0.01),
            retry=retry_if_exception(_should_retry_rpc_error),
            reraise=True,
        )
        def _ensure_workspace():
            raise AIAutoRpcError("unavailable", "never ready")

        with pytest.raises(AIAutoRpcError, match="never ready"):
            _ensure_workspace()


# ---------------------------------------------------------------------------
# 3. _wait_for_trialbatch behavioral tests
# ---------------------------------------------------------------------------
class TestWaitForTrialbatchBehavior:
    """_wait_for_trialbatch: retries RuntimeError, fails fast on ValueError."""

    def test_poll_retries_running_then_succeeds(self):
        """get_status returns "Running" 2x then completed -> retry succeeds."""
        call_count = 0

        def make_status(phase, completed):
            return {
                "trialbatches": {
                    "tb-test": {
                        "phase": phase,
                        "count_completed": completed,
                    }
                }
            }

        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                import aiauto.core as core_mod

                sw = _make_study_wrapper(core_mod)

                def mock_get_status(**kwargs):
                    nonlocal call_count
                    call_count += 1
                    if call_count <= 2:
                        return make_status("Running", 0)
                    return make_status("Succeeded", 3)

                sw.get_status = mock_get_status

                sw._wait_for_trialbatch(
                    trialbatch_name="tb-test",
                    n_trials=3,
                    wait_option=core_mod.WaitOption.WAIT_ALL_TRIALS,
                    timeout=300,
                )
                assert call_count == 3

    def test_poll_fails_immediately_on_failed_phase(self):
        """get_status returns "Failed" phase -> ValueError immediately, no retry."""
        call_count = 0

        def make_failed_status():
            return {
                "trialbatches": {
                    "tb-test": {
                        "phase": "Failed",
                        "conditions": [
                            {
                                "type": "DEGRADED",
                                "status": "True",
                                "message": "OOM killed",
                            }
                        ],
                        "count_completed": 0,
                    }
                }
            }

        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                import aiauto.core as core_mod

                sw = _make_study_wrapper(core_mod)

                def mock_get_status(**kwargs):
                    nonlocal call_count
                    call_count += 1
                    return make_failed_status()

                sw.get_status = mock_get_status

                with pytest.raises(ValueError, match="OOM killed"):
                    sw._wait_for_trialbatch(
                        trialbatch_name="tb-test",
                        n_trials=3,
                        wait_option=core_mod.WaitOption.WAIT_ALL_TRIALS,
                        timeout=300,
                    )

                assert call_count == 1, "Failed phase should not trigger retry"


# ---------------------------------------------------------------------------
# 4. StudyWrapper.wait behavioral tests
# ---------------------------------------------------------------------------
class TestStudyWrapperWaitBehavior:
    """StudyWrapper.wait: retries when is_trial_finished returns False."""

    def test_wait_retries_then_succeeds(self):
        """is_trial_finished returns False 2x then True -> wait returns True."""
        call_count = 0

        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                import aiauto.core as core_mod

                sw = _make_study_wrapper(core_mod)

                def mock_is_trial_finished(trial_id, trialbatch_name=None):
                    nonlocal call_count
                    call_count += 1
                    return call_count >= 3

                sw.is_trial_finished = mock_is_trial_finished

                result = sw.wait(trial_identifier=1, trialbatch_name="tb-test")
                assert result is True
                assert call_count == 3

    def test_wait_returns_false_on_timeout(self):
        """is_trial_finished always False -> timeout -> returns False."""
        with _patch_retry_timing():
            with _patch_retry_timing_stop():
                import aiauto.core as core_mod

                sw = _make_study_wrapper(core_mod)
                sw.is_trial_finished = MagicMock(return_value=False)

                result = sw.wait(trial_identifier=1, trialbatch_name="tb-test")
                assert result is False
                assert sw.is_trial_finished.call_count >= 2
