"""Tests to increase coverage for core.py missed lines.

Covers:
- AIAutoController.__init__ (mock gRPC + EnsureWorkspace)
- get_storage / get_artifact_store / get_artifact_tmp_dir / upload_artifact
- create_study
- StudyWrapper.__init__ / get_study / optimize / get_status / is_trial_finished / wait / __repr__
- _wait_for_trialbatch
"""

import sys
from unittest.mock import MagicMock, patch

import grpc
import pytest

from aiauto.core import (
    AIAutoController,
    StudyWrapper,
    WaitOption,
)
from aiauto.grpc_helper import AIAutoRpcError

# ── Helpers ──


def _make_controller_stub(tus_url="http://localhost:1080/files/"):
    """Create a minimal AIAutoController with mocked internals (skip __init__)."""
    with patch.object(AIAutoController, "__new__", lambda cls, *a, **kw: object.__new__(cls)):
        with patch.object(AIAutoController, "__init__", lambda self, *a, **kw: None):
            ctrl = object.__new__(AIAutoController)
            ctrl.stub = MagicMock()
            ctrl._metadata = (("authorization", "Bearer test"),)
            ctrl.token = "test"
            ctrl.tus_upload_url = tus_url
            ctrl.storage = MagicMock(name="mock_storage")
            ctrl._artifact_store = None
            ctrl.tmp_dir = "/tmp/ai_auto_test"
            ctrl.dashboard_url = "https://dashboard.example.com"
            ctrl._base_url = "http://localhost:8080"
            ctrl._channel = MagicMock()
            ctrl.storage_host_internal = "internal:50051"
            return ctrl


def _make_study_wrapper(
    study_name="test-study",
    controller=None,
    direction="minimize",
    directions=None,
    last_trialbatch=None,
):
    """Create a StudyWrapper with mocked controller."""
    ctrl = controller or _make_controller_stub()
    sw = StudyWrapper(
        study_name=study_name,
        storage=ctrl.storage,
        controller=ctrl,
        direction=direction,
        directions=directions,
    )
    if last_trialbatch:
        sw._last_trialbatch_name = last_trialbatch
    return sw


def _make_fake_rpc_error(code=grpc.StatusCode.UNAVAILABLE, details="service down"):
    """Create a fake gRPC RpcError."""

    class FakeRpcError(grpc.RpcError):
        def code(self):
            return code

        def details(self):
            return details

    return FakeRpcError()


@pytest.fixture()
def mock_proto():
    """Inject mock proto modules into sys.modules for lazy import coverage.

    core.py uses `from .v1 import aiauto_pb2, aiauto_pb2_grpc`
    inside methods (__init__, create_study, etc.). This fixture pre-populates
    sys.modules so the lazy import resolves to our mock objects.
    """
    mock_pb2 = MagicMock()
    mock_pb2_grpc = MagicMock()

    modules_to_inject = {
        "aiauto.v1": MagicMock(aiauto_pb2=mock_pb2, aiauto_pb2_grpc=mock_pb2_grpc),
        "aiauto.v1.aiauto_pb2": mock_pb2,
        "aiauto.v1.aiauto_pb2_grpc": mock_pb2_grpc,
    }

    saved = {}
    for mod_name, mod_mock in modules_to_inject.items():
        saved[mod_name] = sys.modules.get(mod_name)
        sys.modules[mod_name] = mod_mock

    yield mock_pb2, mock_pb2_grpc

    for mod_name, original in saved.items():
        if original is None:
            sys.modules.pop(mod_name, None)
        else:
            sys.modules[mod_name] = original


# ── AIAutoController.__init__ ──


class TestAIAutoControllerInit:
    """Test AIAutoController.__init__ with mocked gRPC."""

    @pytest.fixture(autouse=True)
    def clear_instances(self):
        yield
        AIAutoController._instances.clear()

    @patch("aiauto.core.AIAutoController._create_grpc_channel")
    @patch("aiauto.core._parse_grpc_url", return_value=("host", "50051", "user123"))
    @patch("aiauto.core.PathPrefixGrpcStorageProxy")
    def test_init_with_path_prefix_storage(self, mock_proxy, mock_parse, mock_channel, mock_proto):
        """__init__ uses PathPrefixGrpcStorageProxy when user_id is present."""
        _mock_pb2, mock_pb2_grpc = mock_proto
        mock_channel.return_value = MagicMock(name="channel")

        mock_stub = MagicMock()
        mock_pb2_grpc.AIAutoServiceStub.return_value = mock_stub

        mock_response = MagicMock()
        mock_response.journal_grpc_storage_proxy_host_external = "https://host:443/user123"
        mock_response.journal_grpc_storage_proxy_host_internal = "internal:50051"
        mock_response.dashboard_url = "https://dash.example.com"
        mock_response.tus_upload_url = "https://tus.example.com/files/"
        mock_stub.EnsureWorkspace.return_value = mock_response

        ctrl = AIAutoController.__new__(AIAutoController, "test-token-init")
        ctrl.__init__("test-token-init")

        assert ctrl.token == "test-token-init"
        mock_proxy.assert_called_once_with(host="host", port="50051", user_id="user123")
        assert ctrl.dashboard_url == "https://dash.example.com"
        assert ctrl.tus_upload_url == "https://tus.example.com/files/"

    @patch("aiauto.core.AIAutoController._create_grpc_channel")
    @patch("aiauto.core._parse_grpc_url", return_value=("host", "50051", ""))
    def test_init_without_user_id_falls_back_to_grpc_storage_proxy(
        self, mock_parse, mock_channel, mock_proto
    ):
        """__init__ falls back to optuna GrpcStorageProxy when no user_id."""
        _mock_pb2, mock_pb2_grpc = mock_proto
        mock_channel.return_value = MagicMock(name="channel")

        mock_stub = MagicMock()
        mock_pb2_grpc.AIAutoServiceStub.return_value = mock_stub

        mock_response = MagicMock()
        mock_response.journal_grpc_storage_proxy_host_external = "host:50051"
        mock_response.journal_grpc_storage_proxy_host_internal = "internal:50051"
        mock_response.dashboard_url = ""
        mock_response.tus_upload_url = ""
        mock_stub.EnsureWorkspace.return_value = mock_response

        with patch("aiauto.core.optuna.storages.GrpcStorageProxy") as mock_grpc_proxy:
            ctrl = AIAutoController.__new__(AIAutoController, "token-no-uid")
            ctrl.__init__("token-no-uid")

        mock_grpc_proxy.assert_called_once_with(host="host", port=50051)

    @patch("aiauto.core.AIAutoController._create_grpc_channel")
    def test_init_already_initialized_returns_early(self, mock_channel):
        """__init__ with same token returns early if already initialized."""
        mock_channel.return_value = MagicMock()

        ctrl = AIAutoController.__new__(AIAutoController, "token-early")
        ctrl.token = "token-early"

        # Should return early without calling EnsureWorkspace
        ctrl.__init__("token-early")

    @patch("aiauto.core.AIAutoController._create_grpc_channel")
    def test_init_rpc_error_raises_aiauto_rpc_error(self, mock_channel, mock_proto):
        """__init__ wraps AIAutoRpcError into RuntimeError with hint.

        _ensure_workspace() is inside try/except. With reraise=True, the last
        AIAutoRpcError propagates from tenacity, then except AIAutoRpcError catches it
        and wraps into RuntimeError("Failed to initialize workspace: ...").
        """
        _mock_pb2, mock_pb2_grpc = mock_proto
        mock_channel.return_value = MagicMock(name="channel")

        mock_stub = MagicMock()
        mock_pb2_grpc.AIAutoServiceStub.return_value = mock_stub
        mock_stub.EnsureWorkspace.side_effect = _make_fake_rpc_error(
            grpc.StatusCode.UNAUTHENTICATED, "bad token"
        )

        ctrl = AIAutoController.__new__(AIAutoController, "bad-token")
        with pytest.raises(RuntimeError, match="Failed to initialize workspace"):
            ctrl.__init__("bad-token")

    @patch("aiauto.core.AIAutoController._create_grpc_channel")
    def test_init_generic_exception_propagates_directly(self, mock_channel, mock_proto):
        """__init__ wraps non-RPC exceptions into RuntimeError with default hint.

        _ensure_workspace() is inside try/except. Non-retryable exceptions propagate
        from tenacity, then except Exception catches them and wraps into
        RuntimeError("Failed to initialize workspace: ...").
        """
        _mock_pb2, mock_pb2_grpc = mock_proto
        mock_channel.return_value = MagicMock(name="channel")

        mock_stub = MagicMock()
        mock_pb2_grpc.AIAutoServiceStub.return_value = mock_stub
        mock_stub.EnsureWorkspace.side_effect = ConnectionError("network failure")

        ctrl = AIAutoController.__new__(AIAutoController, "conn-fail")
        with pytest.raises(RuntimeError, match="Failed to initialize workspace"):
            ctrl.__init__("conn-fail")

    @patch("aiauto.core.AIAutoController._create_grpc_channel")
    def test_init_no_storage_host_retries_then_raises(self, mock_channel, mock_proto):
        """__init__ raises RuntimeError("OptunaWorkspace creation timed out") on empty host.

        Empty host triggers AIAutoRpcError("unavailable", ...) inside _ensure_workspace.
        This is retryable, so tenacity retries until stop condition. Without reraise=True,
        timeout raises RetryError. except RetryError wraps it into
        RuntimeError("OptunaWorkspace creation timed out after 300 seconds").
        """
        _mock_pb2, mock_pb2_grpc = mock_proto
        mock_channel.return_value = MagicMock(name="channel")

        mock_stub = MagicMock()
        mock_pb2_grpc.AIAutoServiceStub.return_value = mock_stub

        mock_response = MagicMock()
        mock_response.journal_grpc_storage_proxy_host_external = ""
        mock_stub.EnsureWorkspace.return_value = mock_response

        ctrl = AIAutoController.__new__(AIAutoController, "empty-host")
        # The empty host triggers AIAutoRpcError("unavailable", ...) which retries.
        # To avoid waiting 300s, we patch stop_after_delay to stop immediately.
        with patch("aiauto.core.stop_after_delay", return_value=MagicMock(return_value=True)):
            with pytest.raises(RuntimeError, match="OptunaWorkspace creation timed out"):
                ctrl.__init__("empty-host")


# ── get_storage / get_artifact_store / get_artifact_tmp_dir / upload_artifact ──


class TestControllerAccessors:
    """Test get_storage, get_artifact_store, get_artifact_tmp_dir, upload_artifact."""

    def test_get_storage_returns_storage(self):
        ctrl = _make_controller_stub()
        assert ctrl.get_storage() is ctrl.storage

    def test_get_artifact_tmp_dir_returns_tmp_dir(self):
        ctrl = _make_controller_stub()
        assert ctrl.get_artifact_tmp_dir() == "/tmp/ai_auto_test"

    @patch("os.path.isdir", return_value=True)
    @patch("aiauto.core.optuna.artifacts.FileSystemArtifactStore")
    def test_get_artifact_store_uses_artifacts_dir_when_exists(self, mock_fs, mock_isdir):
        """get_artifact_store uses /artifacts when it exists."""
        ctrl = _make_controller_stub()
        ctrl._artifact_store = None

        store = ctrl.get_artifact_store()

        mock_fs.assert_called_once_with("/artifacts")
        assert store == mock_fs.return_value

    @patch("os.path.isdir", side_effect=lambda p: p != "/artifacts")
    @patch("os.makedirs")
    @patch("aiauto.core.optuna.artifacts.FileSystemArtifactStore")
    def test_get_artifact_store_fallback_local_exists(self, mock_fs, mock_makedirs, mock_isdir):
        """get_artifact_store falls back to ./artifacts when /artifacts doesn't exist but ./artifacts does."""
        ctrl = _make_controller_stub()
        ctrl._artifact_store = None

        ctrl.get_artifact_store()

        mock_fs.assert_called_once_with("./artifacts")
        mock_makedirs.assert_not_called()

    @patch("os.path.isdir", return_value=False)
    @patch("os.makedirs")
    @patch("aiauto.core.optuna.artifacts.FileSystemArtifactStore")
    def test_get_artifact_store_fallback_local_makedirs(self, mock_fs, mock_makedirs, mock_isdir):
        """get_artifact_store creates ./artifacts via makedirs when neither /artifacts nor ./artifacts exist."""
        ctrl = _make_controller_stub()
        ctrl._artifact_store = None

        ctrl.get_artifact_store()

        mock_makedirs.assert_called_once_with("./artifacts", exist_ok=True)
        mock_fs.assert_called_once_with("./artifacts")

    def test_get_artifact_store_returns_cached(self):
        """get_artifact_store returns cached instance on subsequent calls."""
        ctrl = _make_controller_stub()
        fake_store = MagicMock(name="cached_store")
        ctrl._artifact_store = fake_store

        assert ctrl.get_artifact_store() is fake_store

    @patch("aiauto.core.optuna.artifacts.upload_artifact", return_value="artifact-id-123")
    def test_upload_artifact(self, mock_upload):
        """upload_artifact uploads and sets user_attr."""
        ctrl = _make_controller_stub()
        ctrl._artifact_store = MagicMock(name="store")

        mock_trial = MagicMock()

        result = ctrl.upload_artifact(mock_trial, "/path/to/model.pt")

        assert result == "artifact-id-123"
        mock_upload.assert_called_once_with(
            artifact_store=ctrl._artifact_store,
            storage=ctrl.storage,
            study_or_trial=mock_trial,
            file_path="/path/to/model.pt",
        )
        mock_trial.set_user_attr.assert_called_once_with("artifact_id", "artifact-id-123")


# ── create_study ──


class TestCreateStudy:
    """Test AIAutoController.create_study."""

    def test_create_study_with_direction(self, mock_proto):
        """create_study with single direction returns StudyWrapper."""
        _mock_pb2, _ = mock_proto
        ctrl = _make_controller_stub()

        mock_response = MagicMock()
        mock_response.study_name = "my-study"
        ctrl.stub.CreateStudy.return_value = mock_response

        result = ctrl.create_study("my-study", direction="minimize")

        assert isinstance(result, StudyWrapper)
        assert result.study_name == "my-study"

    def test_create_study_with_directions(self, mock_proto):
        """create_study with multiple directions returns StudyWrapper."""
        _mock_pb2, _ = mock_proto
        ctrl = _make_controller_stub()

        mock_response = MagicMock()
        mock_response.study_name = "multi-study"
        ctrl.stub.CreateStudy.return_value = mock_response

        result = ctrl.create_study(
            "multi-study", direction=None, directions=["minimize", "maximize"]
        )

        assert isinstance(result, StudyWrapper)
        assert result.study_name == "multi-study"

    def test_create_study_no_direction_raises(self, mock_proto):
        """create_study with no direction and no directions raises ValueError."""
        ctrl = _make_controller_stub()

        with pytest.raises(ValueError, match="Either 'direction' or 'directions'"):
            ctrl.create_study("bad-study", direction=None, directions=None)

    def test_create_study_both_direction_and_directions_raises(self, mock_proto):
        """create_study with both direction and directions raises ValueError."""
        ctrl = _make_controller_stub()

        with pytest.raises(ValueError, match="Cannot specify both"):
            ctrl.create_study(
                "bad-study", direction="minimize", directions=["minimize", "maximize"]
            )

    def test_create_study_grpc_error_reraises(self, mock_proto):
        """create_study re-raises AIAutoRpcError from gRPC."""
        ctrl = _make_controller_stub()
        ctrl.stub.CreateStudy.side_effect = _make_fake_rpc_error(
            grpc.StatusCode.ALREADY_EXISTS, "study exists"
        )

        with pytest.raises(AIAutoRpcError):
            ctrl.create_study("dup-study")

    def test_create_study_generic_exception_wraps_runtime(self, mock_proto):
        """create_study wraps non-gRPC exceptions in RuntimeError."""
        mock_pb2, _ = mock_proto
        ctrl = _make_controller_stub()

        # Make StudySpec raise a TypeError to trigger the generic except branch
        mock_pb2.StudySpec.side_effect = TypeError("bad arg")

        with pytest.raises(RuntimeError, match="Failed to create study"):
            ctrl.create_study("err-study")

    def test_create_study_uses_response_study_name(self, mock_proto):
        """create_study uses study_name from response if available."""
        ctrl = _make_controller_stub()

        mock_response = MagicMock()
        mock_response.study_name = "server-assigned-name"
        ctrl.stub.CreateStudy.return_value = mock_response

        result = ctrl.create_study("my-study", direction="minimize")

        assert result.study_name == "server-assigned-name"

    def test_create_study_falls_back_to_input_name(self, mock_proto):
        """create_study uses input study_name when response.study_name is empty."""
        ctrl = _make_controller_stub()

        mock_response = MagicMock()
        mock_response.study_name = ""  # empty from server
        ctrl.stub.CreateStudy.return_value = mock_response

        result = ctrl.create_study("input-name", direction="minimize")

        assert result.study_name == "input-name"


# ── StudyWrapper.__init__ / __repr__ ──


class TestStudyWrapperInit:
    """Test StudyWrapper initialization and repr."""

    def test_init_sets_all_fields(self):
        sw = _make_study_wrapper(study_name="test-sw", direction="minimize")
        assert sw.study_name == "test-sw"
        assert sw._direction == "minimize"
        assert sw._directions is None
        assert sw._study is None
        assert sw._last_trialbatch_name is None

    def test_init_with_directions(self):
        sw = _make_study_wrapper(
            study_name="multi", direction=None, directions=["minimize", "maximize"]
        )
        assert sw._directions == ["minimize", "maximize"]
        assert sw._direction is None

    def test_repr(self):
        sw = _make_study_wrapper(study_name="repr-test")
        r = repr(sw)
        assert "StudyWrapper" in r
        assert "repr-test" in r


# ── StudyWrapper.get_study ──


class TestStudyWrapperGetStudy:
    """Test StudyWrapper.get_study with mocked optuna."""

    @patch("aiauto.core.optuna.load_study")
    def test_get_study_loads_and_wraps_ask(self, mock_load):
        """get_study loads study and wraps ask() to add trialbatch_name."""
        mock_study = MagicMock()
        mock_load.return_value = mock_study

        sw = _make_study_wrapper()
        result = sw.get_study()

        assert result is mock_study
        mock_load.assert_called_once_with(study_name="test-study", storage=sw._storage)
        assert sw._study is mock_study

    @patch("aiauto.core.optuna.load_study")
    def test_get_study_cached_on_second_call(self, mock_load):
        """get_study returns cached study on second call."""
        mock_study = MagicMock()
        mock_load.return_value = mock_study

        sw = _make_study_wrapper()
        first = sw.get_study()
        second = sw.get_study()

        assert first is second
        assert mock_load.call_count == 1

    @patch("aiauto.core.optuna.load_study")
    def test_get_study_wrapped_ask_sets_trialbatch_name(self, mock_load):
        """Wrapped ask() sets trialbatch_name='ask_tell_local' on trial."""
        mock_study = MagicMock()
        mock_trial = MagicMock()
        mock_study.ask.return_value = mock_trial
        mock_load.return_value = mock_study

        sw = _make_study_wrapper()
        study = sw.get_study()

        trial = study.ask()
        assert trial is mock_trial
        mock_trial.set_user_attr.assert_called_with("trialbatch_name", "ask_tell_local")

    @patch("aiauto.core.optuna.load_study")
    def test_get_study_wrapped_ask_swallows_set_user_attr_error(self, mock_load):
        """Wrapped ask() does not raise even if set_user_attr fails."""
        mock_study = MagicMock()
        mock_trial = MagicMock()
        mock_trial.set_user_attr.side_effect = RuntimeError("storage error")
        mock_study.ask.return_value = mock_trial
        mock_load.return_value = mock_study

        sw = _make_study_wrapper()
        study = sw.get_study()

        trial = study.ask()
        assert trial is mock_trial

    @patch("aiauto.core.optuna.load_study", side_effect=Exception("connection refused"))
    def test_get_study_generic_error_raises_runtime_with_dashboard(self, mock_load):
        """get_study raises RuntimeError("Study gRPC storage not ready") on persistent error.

        _load_study has @retry with no retry= filter, so any exception triggers retry.
        Without reraise=True, timeout raises RetryError. except RetryError wraps it into
        RuntimeError("Study gRPC storage not ready after 30 seconds").
        """
        sw = _make_study_wrapper()
        sw._controller.dashboard_url = "https://dash.example.com"

        with patch("aiauto.core.stop_after_delay", return_value=MagicMock(return_value=True)):
            with pytest.raises(RuntimeError, match="Study gRPC storage not ready"):
                sw.get_study()

    @patch("aiauto.core.optuna.load_study", side_effect=Exception("connection refused"))
    def test_get_study_generic_error_no_dashboard(self, mock_load):
        """get_study raises RuntimeError("Study gRPC storage not ready") without dashboard.

        Same as with_dashboard variant: _load_study retries any exception, timeout raises
        RetryError, except RetryError wraps into RuntimeError("Study gRPC storage not ready").
        """
        sw = _make_study_wrapper()
        sw._controller.dashboard_url = ""

        with patch("aiauto.core.stop_after_delay", return_value=MagicMock(return_value=True)):
            with pytest.raises(RuntimeError, match="Study gRPC storage not ready"):
                sw.get_study()

    @patch("aiauto.core.optuna.load_study")
    def test_get_study_post_load_exception_with_dashboard(self, mock_load):
        """get_study except Exception (line 655-659): load succeeds but post-processing fails.

        _load_study succeeds, but accessing study.ask raises a non-RetryError exception.
        This hits the 'except Exception' branch (not 'except RetryError').
        """
        # study.ask property raises TypeError (not RetryError)
        mock_study = MagicMock()
        type(mock_study).ask = property(
            lambda self: (_ for _ in ()).throw(TypeError("bad property"))
        )
        mock_load.return_value = mock_study

        sw = _make_study_wrapper()
        sw._controller.dashboard_url = "https://dash.example.com"

        with pytest.raises(RuntimeError, match="Failed to get study.*dashboard"):
            sw.get_study()

    @patch("aiauto.core.optuna.load_study")
    def test_get_study_post_load_exception_no_dashboard(self, mock_load):
        """get_study except Exception (line 655-659): no dashboard URL in error message."""
        mock_study = MagicMock()
        type(mock_study).ask = property(
            lambda self: (_ for _ in ()).throw(TypeError("bad property"))
        )
        mock_load.return_value = mock_study

        sw = _make_study_wrapper()
        sw._controller.dashboard_url = ""

        with pytest.raises(RuntimeError, match="Failed to get study"):
            sw.get_study()


# ── StudyWrapper.get_status ──


class TestStudyWrapperGetStatus:
    """Test StudyWrapper.get_status with mocked gRPC."""

    def test_get_status_basic(self, mock_proto):
        """get_status returns converted protobuf response."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        mock_condition = MagicMock()
        mock_condition.type = "READY"
        mock_condition.status = "True"
        mock_condition.reason = ""
        mock_condition.message = ""
        mock_condition.last_transition_time = "2024-01-01T00:00:00Z"

        mock_tb_status = MagicMock()
        mock_tb_status.trialbatch_name = "tb-001"
        mock_tb_status.phase = "Running"
        mock_tb_status.conditions = [mock_condition]
        mock_tb_status.count_active = 2
        mock_tb_status.count_succeeded = 3
        mock_tb_status.count_pruned = 0
        mock_tb_status.count_failed = 1
        mock_tb_status.count_total = 6
        mock_tb_status.count_completed = 4
        mock_tb_status.completed_trials = {}

        mock_response = MagicMock()
        mock_response.study_name = "test-study"
        mock_response.trialbatches = {"tb-001": mock_tb_status}
        mock_response.dashboard_url = "https://dash.example.com"
        mock_response.updated_at = "2024-01-01T00:00:00Z"
        ctrl.stub.GetStatus.return_value = mock_response

        result = sw.get_status()

        assert result["study_name"] == "test-study"
        assert "tb-001" in result["trialbatches"]
        tb = result["trialbatches"]["tb-001"]
        assert tb["phase"] == "Running"
        assert tb["count_completed"] == 4

    def test_get_status_with_completed_trials(self, mock_proto):
        """get_status includes completed_trials details."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        mock_trial_info = MagicMock()
        mock_trial_info.trial_number = 5
        mock_trial_info.state = "COMPLETE"
        mock_trial_info.pod_name = "pod-abc"
        mock_trial_info.value = 0.95
        mock_trial_info.user_attrs = {"custom_key": "custom_val"}
        mock_trial_info.start_time = "2024-01-01T00:00:00Z"
        mock_trial_info.end_time = "2024-01-01T00:05:00Z"
        mock_trial_info.value_special = ""

        mock_tb_status = MagicMock()
        mock_tb_status.trialbatch_name = "tb-002"
        mock_tb_status.phase = "Completed"
        mock_tb_status.conditions = []
        mock_tb_status.count_active = 0
        mock_tb_status.count_succeeded = 5
        mock_tb_status.count_pruned = 0
        mock_tb_status.count_failed = 0
        mock_tb_status.count_total = 5
        mock_tb_status.count_completed = 5
        mock_tb_status.completed_trials = {"pod-abc": mock_trial_info}

        mock_response = MagicMock()
        mock_response.study_name = "test-study"
        mock_response.trialbatches = {"tb-002": mock_tb_status}
        mock_response.dashboard_url = ""
        mock_response.updated_at = ""
        ctrl.stub.GetStatus.return_value = mock_response

        result = sw.get_status(include_trials=True)

        tb = result["trialbatches"]["tb-002"]
        assert "completed_trials" in tb
        assert "pod-abc" in tb["completed_trials"]
        trial = tb["completed_trials"]["pod-abc"]
        assert trial["trialNumber"] == 5
        assert trial["state"] == "COMPLETE"

    def test_get_status_with_trialbatch_filter(self, mock_proto):
        """get_status passes trialbatch_name to request."""
        mock_pb2, _ = mock_proto
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        mock_response = MagicMock()
        mock_response.study_name = "test-study"
        mock_response.trialbatches = {}
        mock_response.dashboard_url = ""
        mock_response.updated_at = ""
        ctrl.stub.GetStatus.return_value = mock_response

        mock_request = MagicMock()
        mock_pb2.GetStatusRequest.return_value = mock_request

        sw.get_status(trialbatch_name="tb-filter")

        assert mock_request.trialbatch_name == "tb-filter"

    def test_get_status_grpc_error_reraises(self, mock_proto):
        """get_status re-raises AIAutoRpcError."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        ctrl.stub.GetStatus.side_effect = _make_fake_rpc_error(
            grpc.StatusCode.PERMISSION_DENIED, "forbidden"
        )

        with pytest.raises(AIAutoRpcError):
            sw.get_status()

    def test_get_status_generic_exception_wraps_runtime(self, mock_proto):
        """get_status wraps non-gRPC exceptions in RuntimeError."""
        mock_pb2, _ = mock_proto
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        mock_pb2.GetStatusRequest.side_effect = TypeError("bad proto")

        with pytest.raises(RuntimeError, match="Failed to get status"):
            sw.get_status()


# ── StudyWrapper.optimize ──


class TestStudyWrapperOptimize:
    """Test StudyWrapper.optimize with mocked gRPC."""

    def _setup_optimize(self, ctrl, trialbatch_name="tb-opt-001"):
        """Set up mock for successful optimize call."""
        mock_response = MagicMock()
        mock_response.trialbatch_name = trialbatch_name
        ctrl.stub.Optimize.return_value = mock_response
        return mock_response

    def test_optimize_basic_no_wait(self, mock_proto):
        """optimize with WAIT_NO returns trialbatch_name immediately."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=5,
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"
        assert sw._last_trialbatch_name == "tb-opt-001"

    def test_optimize_string_wait_option(self, mock_proto):
        """optimize converts string wait_option to WaitOption enum."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=5,
            wait_option="WAIT_NO",
        )

        assert result == "tb-opt-001"

    def test_optimize_invalid_string_wait_option_raises(self):
        """optimize with invalid string wait_option raises ValueError."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="Invalid wait_option"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=5,
                wait_option="INVALID_OPTION",
            )

    def test_optimize_cpu_with_custom_dev_shm_raises(self):
        """optimize raises when CPU mode but custom dev_shm_size."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="can only be used with GPU"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                use_gpu=False,
                dev_shm_size="4Gi",
            )

    def test_optimize_dev_shm_exceeds_max_raises(self):
        """optimize raises when dev_shm_size exceeds 4Gi."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="exceeds max size"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                use_gpu=True,
                dev_shm_size="8Gi",
            )

    def test_optimize_tmp_cache_exceeds_max_raises(self):
        """optimize raises when tmp_cache_size exceeds 4Gi."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="exceeds max size"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                tmp_cache_size="10Gi",
            )

    def test_optimize_gpu_model_string(self, mock_proto):
        """optimize with gpu_model as string creates correct mapping."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        with patch("aiauto.core._fetch_available_gpu_models", return_value={"gpu_4090"}):
            result = sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=5,
                use_gpu=True,
                gpu_model="gpu_4090",
                wait_option=WaitOption.WAIT_NO,
            )

        assert result == "tb-opt-001"

    def test_optimize_gpu_model_dict(self, mock_proto):
        """optimize with gpu_model as dict validates and passes."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        with patch(
            "aiauto.core._fetch_available_gpu_models", return_value={"gpu_3090", "gpu_4090"}
        ):
            result = sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=10,
                use_gpu=True,
                gpu_model={"gpu_3090": 5, "gpu_4090": 4},
                wait_option=WaitOption.WAIT_NO,
            )

        assert result == "tb-opt-001"

    def test_optimize_gpu_model_invalid_type_raises(self):
        """optimize raises when gpu_model is not str or dict."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="gpu_model must be str or dict"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=5,
                gpu_model=42,
                wait_option=WaitOption.WAIT_NO,
            )

    def test_optimize_gpu_model_no_available_raises(self):
        """optimize raises when no GPU models available in cluster."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with patch("aiauto.core._fetch_available_gpu_models", return_value=set()):
            with pytest.raises(ValueError, match="No GPU models available"):
                sw.optimize(
                    objective=lambda trial: 1.0,
                    n_trials=5,
                    gpu_model="gpu_4090",
                    wait_option=WaitOption.WAIT_NO,
                )

    def test_optimize_gpu_model_invalid_name_raises(self):
        """optimize raises for invalid GPU model name."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with patch("aiauto.core._fetch_available_gpu_models", return_value={"gpu_3090"}):
            with pytest.raises(ValueError, match="Invalid GPU model"):
                sw.optimize(
                    objective=lambda trial: 1.0,
                    n_trials=5,
                    gpu_model="gpu_9090",
                    wait_option=WaitOption.WAIT_NO,
                )

    def test_optimize_gpu_model_negative_value_raises(self):
        """optimize raises when gpu_model dict has non-positive values."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with patch("aiauto.core._fetch_available_gpu_models", return_value={"gpu_3090"}):
            with pytest.raises(ValueError, match="must be positive"):
                sw.optimize(
                    objective=lambda trial: 1.0,
                    n_trials=5,
                    gpu_model={"gpu_3090": -1},
                    wait_option=WaitOption.WAIT_NO,
                )

    def test_optimize_gpu_model_exceeds_n_trials_raises(self):
        """optimize raises when GPU allocations exceed n_trials."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with patch(
            "aiauto.core._fetch_available_gpu_models", return_value={"gpu_3090", "gpu_4090"}
        ):
            with pytest.raises(ValueError, match="exceeds n_trials"):
                sw.optimize(
                    objective=lambda trial: 1.0,
                    n_trials=10,
                    gpu_model={"gpu_3090": 6, "gpu_4090": 6},
                    wait_option=WaitOption.WAIT_NO,
                )

    def test_optimize_image_pull_both_methods_raises(self):
        """optimize raises when both registry and docker_config_json provided."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="Cannot use both"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                image_pull_registry="registry.example.com",
                image_pull_username="user",
                image_pull_password="pass",
                image_pull_docker_config_json={"auths": {}},
                wait_option=WaitOption.WAIT_NO,
            )

    def test_optimize_image_pull_registry_incomplete_raises(self):
        """optimize raises when registry method is incomplete."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="requires all 3 fields"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                image_pull_registry="registry.example.com",
                image_pull_username="user",
                wait_option=WaitOption.WAIT_NO,
            )

    def test_optimize_grpc_error_reraises(self, mock_proto):
        """optimize re-raises AIAutoRpcError from gRPC."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        ctrl.stub.Optimize.side_effect = _make_fake_rpc_error(
            grpc.StatusCode.INTERNAL, "server error"
        )

        with pytest.raises(AIAutoRpcError):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                wait_option=WaitOption.WAIT_NO,
            )

    def test_optimize_generic_exception_wraps_runtime(self, mock_proto):
        """optimize wraps non-gRPC exceptions in RuntimeError."""
        mock_pb2, _ = mock_proto
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        mock_pb2.ObjectivePayload.side_effect = TypeError("bad serialize")

        with pytest.raises(RuntimeError, match="Failed to start optimization"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                wait_option=WaitOption.WAIT_NO,
            )

    def test_optimize_with_registry_credentials(self, mock_proto):
        """optimize passes image pull registry credentials."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=1,
            image_pull_registry="registry.example.com",
            image_pull_username="user",
            image_pull_password="pass123",
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"

    def test_optimize_with_docker_config_json(self, mock_proto):
        """optimize passes docker config json credentials."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=1,
            image_pull_docker_config_json={"auths": {"reg": {"auth": "abc"}}},
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"

    def test_optimize_default_resources_cpu(self, mock_proto):
        """optimize sets default CPU resources when none specified."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=1,
            use_gpu=False,
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"

    def test_optimize_default_resources_gpu(self, mock_proto):
        """optimize sets default GPU resources when use_gpu=True."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=1,
            use_gpu=True,
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"

    def test_optimize_custom_resources_and_image(self, mock_proto):
        """optimize with user-provided resources/limits/runtime_image skips default branches.

        Covers branch partials 800->806 (resources_requests not None),
        806->809 (resources_limits not None), 809->816 (runtime_image not None/empty).
        """
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=1,
            use_gpu=False,
            resources_requests={"cpu": "4", "memory": "8Gi"},
            resources_limits={"cpu": "8", "memory": "16Gi"},
            runtime_image="custom-image:v1",
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"

    def test_optimize_empty_dev_shm_empty_tmp_cache_gpu(self, mock_proto):
        """optimize with empty dev_shm_size and empty tmp_cache_size skips both checks.

        Covers branch partials 751->765 (dev_shm_size falsy skip)
        and 765->780 (tmp_cache_size falsy skip).
        use_gpu=True to bypass the dev_shm_size != "500Mi" CPU guard at line 744.
        """
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl)

        result = sw.optimize(
            objective=lambda trial: 1.0,
            n_trials=1,
            use_gpu=True,
            dev_shm_size="",
            tmp_cache_size="",
            wait_option=WaitOption.WAIT_NO,
        )

        assert result == "tb-opt-001"

    def test_optimize_invalid_dev_shm_format_raises(self):
        """optimize raises on invalid dev_shm_size format."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="dev_shm_size"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                use_gpu=True,
                dev_shm_size="badformat",
            )

    def test_optimize_invalid_tmp_cache_format_raises(self):
        """optimize raises on invalid tmp_cache_size format."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)

        with pytest.raises(ValueError, match="tmp_cache_size"):
            sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=1,
                tmp_cache_size="xyz",
            )

    def test_optimize_with_wait_calls_wait_for_trialbatch(self, mock_proto):
        """optimize with WAIT_ALL calls _wait_for_trialbatch."""
        ctrl = _make_controller_stub()
        sw = _make_study_wrapper(controller=ctrl)
        self._setup_optimize(ctrl, trialbatch_name="tb-wait")

        with patch.object(sw, "_wait_for_trialbatch") as mock_wait:
            result = sw.optimize(
                objective=lambda trial: 1.0,
                n_trials=5,
                wait_option=WaitOption.WAIT_ALL_TRIALS,
                wait_timeout=60,
            )

        assert result == "tb-wait"
        mock_wait.assert_called_once_with("tb-wait", 5, WaitOption.WAIT_ALL_TRIALS, 60)


# ── StudyWrapper._wait_for_trialbatch ──


class TestWaitForTrialbatch:
    """Test StudyWrapper._wait_for_trialbatch."""

    def _make_status_response(self, tb_name, phase="Running", completed=0, conditions=None):
        """Build a status dict matching get_status output."""
        return {
            "study_name": "test-study",
            "trialbatches": {
                tb_name: {
                    "trialbatch_name": tb_name,
                    "phase": phase,
                    "conditions": conditions or [],
                    "count_active": 1,
                    "count_succeeded": completed,
                    "count_pruned": 0,
                    "count_failed": 0,
                    "count_total": 5,
                    "count_completed": completed,
                }
            },
            "dashboard_url": "",
            "updated_at": "",
        }

    def test_wait_all_trials_completed(self):
        """_wait_for_trialbatch returns when all trials completed."""
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(return_value=self._make_status_response("tb-1", completed=5))

        sw._wait_for_trialbatch("tb-1", 5, WaitOption.WAIT_ALL_TRIALS, 30)

    def test_wait_atleast_one_completed(self):
        """_wait_for_trialbatch returns when at least one trial completed."""
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(return_value=self._make_status_response("tb-1", completed=1))

        sw._wait_for_trialbatch("tb-1", 5, WaitOption.WAIT_ATLEAST_ONE_TRIAL, 30)

    def test_wait_trialbatch_not_found_raises(self):
        """_wait_for_trialbatch raises RuntimeError("Timeout after...") when tb not found.

        Without reraise=True, timeout raises RetryError. except RetryError wraps it
        into RuntimeError("Timeout after N seconds...").
        """
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(
            return_value={
                "study_name": "test",
                "trialbatches": {},
                "dashboard_url": "",
                "updated_at": "",
            }
        )

        with pytest.raises(RuntimeError, match="Timeout after"):
            sw._wait_for_trialbatch("tb-missing", 5, WaitOption.WAIT_ALL_TRIALS, 0)

    def test_wait_trialbatch_failed_degraded_condition(self):
        """_wait_for_trialbatch raises ValueError immediately on Failed phase with DEGRADED."""
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(
            return_value=self._make_status_response(
                "tb-fail",
                phase="Failed",
                conditions=[
                    {
                        "type": "DEGRADED",
                        "status": "True",
                        "reason": "Error",
                        "message": "OOM killed",
                        "last_transition_time": "",
                    }
                ],
            )
        )

        with pytest.raises(ValueError, match="OOM killed"):
            sw._wait_for_trialbatch("tb-fail", 5, WaitOption.WAIT_ALL_TRIALS, 30)

    def test_wait_trialbatch_failed_non_degraded_condition(self):
        """_wait_for_trialbatch extracts message from non-DEGRADED condition."""
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(
            return_value=self._make_status_response(
                "tb-fail2",
                phase="Failed",
                conditions=[
                    {
                        "type": "SOME_OTHER",
                        "status": "True",
                        "reason": "Error",
                        "message": "image pull failed",
                        "last_transition_time": "",
                    }
                ],
            )
        )

        with pytest.raises(ValueError, match="image pull failed"):
            sw._wait_for_trialbatch("tb-fail2", 5, WaitOption.WAIT_ALL_TRIALS, 30)

    def test_wait_trialbatch_failed_non_degraded_skip_then_match(self):
        """_wait_for_trialbatch: second for-loop skips first condition, matches second.

        Covers branch partial 918->917: first condition has status=False (skip),
        second condition has status=True with message (match).
        """
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(
            return_value=self._make_status_response(
                "tb-fail-skip",
                phase="Failed",
                conditions=[
                    {
                        "type": "READY",
                        "status": "False",
                        "reason": "NotReady",
                        "message": "",
                        "last_transition_time": "",
                    },
                    {
                        "type": "AVAILABLE",
                        "status": "True",
                        "reason": "Error",
                        "message": "pod evicted",
                        "last_transition_time": "",
                    },
                ],
            )
        )

        with pytest.raises(ValueError, match="pod evicted"):
            sw._wait_for_trialbatch("tb-fail-skip", 5, WaitOption.WAIT_ALL_TRIALS, 30)

    def test_wait_trialbatch_failed_no_conditions(self):
        """_wait_for_trialbatch raises generic message when no condition has message."""
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(
            return_value=self._make_status_response(
                "tb-fail3",
                phase="Failed",
                conditions=[],
            )
        )

        with pytest.raises(ValueError, match="failed"):
            sw._wait_for_trialbatch("tb-fail3", 5, WaitOption.WAIT_ALL_TRIALS, 30)

    def test_wait_timeout_all_trials(self):
        """_wait_for_trialbatch raises RuntimeError("Timeout after...") for WAIT_ALL_TRIALS.

        Without reraise=True, timeout raises RetryError. except RetryError wraps it
        into RuntimeError("Timeout after N seconds...") with dashboard hint.
        """
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(return_value=self._make_status_response("tb-slow", completed=0))

        with pytest.raises(RuntimeError, match="Timeout after"):
            sw._wait_for_trialbatch("tb-slow", 5, WaitOption.WAIT_ALL_TRIALS, 0)

    def test_wait_timeout_atleast_one(self):
        """_wait_for_trialbatch raises RuntimeError("Timeout after...") for WAIT_ATLEAST_ONE.

        Without reraise=True, timeout raises RetryError. except RetryError wraps it
        into RuntimeError("Timeout after N seconds...") with dashboard hint.
        """
        sw = _make_study_wrapper()
        sw.get_status = MagicMock(return_value=self._make_status_response("tb-slow", completed=0))

        with pytest.raises(RuntimeError, match="Timeout after"):
            sw._wait_for_trialbatch("tb-slow", 5, WaitOption.WAIT_ATLEAST_ONE_TRIAL, 0)


# ── StudyWrapper.is_trial_finished ──


class TestIsTrialFinished:
    """Test StudyWrapper.is_trial_finished."""

    def test_trial_number_without_trialbatch_raises(self):
        """is_trial_finished with int and no trialbatch raises ValueError."""
        sw = _make_study_wrapper()
        sw._last_trialbatch_name = None

        with pytest.raises(ValueError, match="trialbatch_name is required"):
            sw.is_trial_finished(5)

    def test_no_trialbatch_name_returns_false(self):
        """is_trial_finished with str identifier and no trialbatch returns False."""
        sw = _make_study_wrapper()
        sw._last_trialbatch_name = None

        result = sw.is_trial_finished("pod-abc")
        assert result is False

    def test_trialbatch_not_found_returns_false(self):
        """is_trial_finished returns False when trialbatch not in status."""
        sw = _make_study_wrapper(last_trialbatch="tb-old")
        sw.get_status = MagicMock(
            return_value={
                "study_name": "test",
                "trialbatches": {},
                "dashboard_url": "",
                "updated_at": "",
            }
        )

        assert sw.is_trial_finished("pod-xyz") is False

    def test_trial_found_by_number(self):
        """is_trial_finished returns True when trial number matches."""
        sw = _make_study_wrapper(last_trialbatch="tb-1")
        sw.get_status = MagicMock(
            return_value={
                "study_name": "test",
                "trialbatches": {
                    "tb-1": {
                        "completed_trials": {
                            "pod-a": {"trialNumber": 5, "state": "COMPLETE"},
                            "pod-b": {"trialNumber": 3, "state": "COMPLETE"},
                        }
                    }
                },
                "dashboard_url": "",
                "updated_at": "",
            }
        )

        assert sw.is_trial_finished(5) is True
        assert sw.is_trial_finished(99) is False

    def test_trial_found_by_pod_name(self):
        """is_trial_finished returns True when pod name matches."""
        sw = _make_study_wrapper(last_trialbatch="tb-1")
        sw.get_status = MagicMock(
            return_value={
                "study_name": "test",
                "trialbatches": {
                    "tb-1": {
                        "completed_trials": {
                            "pod-abc": {"trialNumber": 1, "state": "COMPLETE"},
                        }
                    }
                },
                "dashboard_url": "",
                "updated_at": "",
            }
        )

        assert sw.is_trial_finished("pod-abc") is True
        assert sw.is_trial_finished("pod-xyz") is False

    def test_trial_finished_uses_provided_trialbatch(self):
        """is_trial_finished uses explicitly provided trialbatch_name."""
        sw = _make_study_wrapper(last_trialbatch="tb-old")
        sw.get_status = MagicMock(
            return_value={
                "study_name": "test",
                "trialbatches": {
                    "tb-new": {
                        "completed_trials": {
                            "pod-x": {"trialNumber": 1, "state": "COMPLETE"},
                        }
                    }
                },
                "dashboard_url": "",
                "updated_at": "",
            }
        )

        assert sw.is_trial_finished("pod-x", trialbatch_name="tb-new") is True

    def test_trial_finished_exception_returns_false(self):
        """is_trial_finished returns False on exception."""
        sw = _make_study_wrapper(last_trialbatch="tb-1")
        sw.get_status = MagicMock(side_effect=RuntimeError("network error"))

        assert sw.is_trial_finished("pod-abc") is False


# ── StudyWrapper.wait ──


class TestStudyWrapperWait:
    """Test StudyWrapper.wait method."""

    def test_wait_int_without_trialbatch_raises(self):
        """wait with int and no trialbatch raises ValueError."""
        sw = _make_study_wrapper()
        sw._last_trialbatch_name = None

        with pytest.raises(ValueError, match="trialbatch_name is required"):
            sw.wait(5)

    def test_wait_str_without_trialbatch_raises(self):
        """wait with str and no trialbatch raises RuntimeError."""
        sw = _make_study_wrapper()
        sw._last_trialbatch_name = None

        with pytest.raises(RuntimeError, match="No TrialBatch tracked"):
            sw.wait("pod-abc")

    def test_wait_returns_true_when_finished(self):
        """wait returns True when trial finishes within timeout."""
        sw = _make_study_wrapper(last_trialbatch="tb-1")
        sw.is_trial_finished = MagicMock(return_value=True)

        result = sw.wait("pod-abc", timeout=5)
        assert result is True

    def test_wait_returns_false_on_timeout(self):
        """wait returns False when trial does not finish within timeout."""
        sw = _make_study_wrapper(last_trialbatch="tb-1")
        sw.is_trial_finished = MagicMock(return_value=False)

        result = sw.wait("pod-abc", timeout=0)
        assert result is False

    def test_wait_uses_provided_trialbatch(self):
        """wait uses explicitly provided trialbatch_name."""
        sw = _make_study_wrapper(last_trialbatch="tb-old")
        sw.is_trial_finished = MagicMock(return_value=True)

        result = sw.wait(5, trialbatch_name="tb-explicit", timeout=5)
        assert result is True
        sw.is_trial_finished.assert_called_with(5, trialbatch_name="tb-explicit")
