"""Tests for health check, readiness, startup probes, and graceful shutdown.

测试 K8s 生命周期探针和优雅关闭机制。
"""

from __future__ import annotations

from typing import Any
from unittest.mock import AsyncMock, Mock, patch

import pytest
from fastapi.testclient import TestClient


@pytest.fixture
def mock_control_plane() -> Mock:
    """创建模拟的 Control Plane."""
    control_plane = Mock()
    control_plane.start = AsyncMock()
    control_plane.stop = AsyncMock()
    control_plane.list_engines = Mock(
        return_value=[
            {"engine_id": "engine-1", "model_id": "test-model", "state": "ready"},
            {"engine_id": "engine-2", "model_id": "test-model", "state": "ready"},
        ]
    )
    return control_plane


@pytest.fixture
def mock_adapter() -> Mock:
    """创建模拟的 Adapter."""
    adapter = Mock()
    adapter.initialize = AsyncMock()
    return adapter


@pytest.fixture
def app_with_mocks(mock_control_plane: Mock, mock_adapter: Mock) -> Any:
    """创建带有模拟依赖的应用."""
    with patch("sagellm_control.ControlPlaneManager", return_value=mock_control_plane):
        with patch("sagellm_gateway.server.SageLLMAdapter", return_value=mock_adapter):
            from sagellm_gateway.server import app

            yield app


class TestHealthProbes:
    """测试健康检查探针."""

    def test_health_endpoint_basic(self, app_with_mocks: Any) -> None:
        """测试基本健康检查端点."""
        client = TestClient(app_with_mocks)
        response = client.get("/health")
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "healthy"
        assert "state" in data
        assert "uptime_seconds" in data
        assert "active_requests" in data
        assert "sessions" in data

    def test_ready_endpoint_when_ready(
        self, app_with_mocks: Any, mock_control_plane: Mock, mock_adapter: Mock
    ) -> None:
        """测试就绪探针 - 服务就绪时."""
        from sagellm_gateway.server import AppState, _gateway_state

        # 确保状态为就绪
        _gateway_state.app_state = AppState.READY
        _gateway_state.control_plane = mock_control_plane
        _gateway_state.adapter = mock_adapter

        client = TestClient(app_with_mocks)

        response = client.get("/ready")
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "ready"
        assert "services" in data
        assert "control_plane" in data["services"]
        assert "adapter" in data["services"]
        assert data["services"]["control_plane"]["status"] == "up"
        assert data["services"]["adapter"]["status"] == "up"

    def test_ready_endpoint_control_plane_down(
        self, app_with_mocks: Any, mock_control_plane: Mock
    ) -> None:
        """测试就绪探针 - Control Plane 故障时."""
        from sagellm_gateway.server import AppState, _gateway_state

        # 确保状态为就绪
        _gateway_state.app_state = AppState.READY
        _gateway_state.control_plane = mock_control_plane

        # 模拟 Control Plane 故障
        mock_control_plane.list_engines.side_effect = Exception("Connection refused")

        client = TestClient(app_with_mocks)
        response = client.get("/ready")

        assert response.status_code == 503
        data = response.json()
        assert "not ready" in data["detail"].lower() or "not_ready" in str(data)

    def test_startup_endpoint_after_startup(self, app_with_mocks: Any) -> None:
        """测试启动探针 - 启动完成后."""
        from sagellm_gateway.server import AppState, _gateway_state

        # 确保状态为就绪（已完成启动）
        _gateway_state.app_state = AppState.READY

        client = TestClient(app_with_mocks)

        response = client.get("/startup")
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "started"
        assert "state" in data
        assert "uptime_seconds" in data

    def test_health_endpoint_during_operation(self, app_with_mocks: Any) -> None:
        """测试健康检查端点 - 运行期间."""
        client = TestClient(app_with_mocks)

        # 模拟一些请求
        for _ in range(3):
            response = client.get("/health")
            assert response.status_code == 200

    def test_ready_endpoint_reports_engine_count(
        self, app_with_mocks: Any, mock_control_plane: Mock
    ) -> None:
        """测试就绪探针报告引擎数量."""
        client = TestClient(app_with_mocks)

        # 等待服务就绪
        import time

        time.sleep(0.5)

        response = client.get("/ready")
        if response.status_code == 200:
            data = response.json()
            assert "services" in data
            assert "control_plane" in data["services"]
            assert "engine_count" in data["services"]["control_plane"]
            # 根据 mock 设置，应该有 2 个引擎
            assert data["services"]["control_plane"]["engine_count"] == 2


class TestGracefulShutdown:
    """测试优雅关闭机制."""

    @pytest.mark.asyncio
    async def test_reject_requests_during_shutdown(self, app_with_mocks: Any) -> None:
        """测试关闭期间拒绝新请求."""
        from sagellm_gateway.server import AppState, _gateway_state

        # 模拟正在关闭状态
        _gateway_state.app_state = AppState.SHUTTING_DOWN

        client = TestClient(app_with_mocks, raise_server_exceptions=False)
        response = client.get("/v1/models")

        assert response.status_code == 503
        assert "shutting down" in response.text.lower()

        # 恢复状态
        _gateway_state.app_state = AppState.READY

    @pytest.mark.asyncio
    async def test_active_request_tracking(self, app_with_mocks: Any) -> None:
        """测试活动请求追踪."""
        from sagellm_gateway.server import _gateway_state

        initial_count = _gateway_state.active_requests

        client = TestClient(app_with_mocks)

        # 发起请求
        response = client.get("/health")
        assert response.status_code == 200

        # 请求完成后，计数应该恢复
        assert _gateway_state.active_requests == initial_count


class TestRootEndpoint:
    """测试根端点."""

    def test_root_endpoint_includes_probe_paths(self, app_with_mocks: Any) -> None:
        """测试根端点包含探针路径."""
        client = TestClient(app_with_mocks)
        response = client.get("/")
        assert response.status_code == 200
        data = response.json()
        assert "endpoints" in data
        assert "health" in data["endpoints"]
        assert "ready" in data["endpoints"]
        assert "startup" in data["endpoints"]
        assert data["endpoints"]["health"] == "/health"
        assert data["endpoints"]["ready"] == "/ready"
        assert data["endpoints"]["startup"] == "/startup"


class TestProbeResilience:
    """测试探针的鲁棒性."""

    def test_health_probe_always_returns_200(self, app_with_mocks: Any) -> None:
        """测试健康探针总是返回 200（除非服务完全崩溃）."""
        client = TestClient(app_with_mocks)

        # 即使有错误，健康探针也应该返回 200
        for _ in range(5):
            response = client.get("/health")
            assert response.status_code == 200

    def test_ready_probe_handles_control_plane_exception(
        self, app_with_mocks: Any, mock_control_plane: Mock
    ) -> None:
        """测试就绪探针优雅处理 Control Plane 异常."""
        from sagellm_gateway.server import AppState, _gateway_state

        # 确保状态为就绪
        _gateway_state.app_state = AppState.READY
        _gateway_state.control_plane = mock_control_plane

        # 模拟 Control Plane 抛出异常
        mock_control_plane.list_engines.side_effect = RuntimeError("Database connection lost")

        client = TestClient(app_with_mocks)
        response = client.get("/ready")

        # 应该返回 503，而不是崩溃
        assert response.status_code == 503

    def test_startup_probe_reports_state(self, app_with_mocks: Any) -> None:
        """测试启动探针报告正确的状态."""
        import time

        client = TestClient(app_with_mocks)
        time.sleep(0.5)  # 等待启动完成

        response = client.get("/startup")
        assert response.status_code == 200
        data = response.json()
        assert data["state"] in ["ready", "started"]


class TestSessionIntegration:
    """测试会话管理与健康检查的集成."""

    def test_health_includes_session_stats(self, app_with_mocks: Any) -> None:
        """测试健康检查包含会话统计."""
        client = TestClient(app_with_mocks)

        # 创建一些会话
        client.post("/sessions", json={"title": "Test Session 1"})
        client.post("/sessions", json={"title": "Test Session 2"})

        # 检查健康状态
        response = client.get("/health")
        assert response.status_code == 200
        data = response.json()
        assert "sessions" in data
        assert "total_sessions" in data["sessions"]
        assert data["sessions"]["total_sessions"] >= 2


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
