"""Tests for MCP hot reload and auto-reconnect functionality."""

import time
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from starlette.testclient import TestClient

from atlas.modules.mcp_tools.client import MCPToolManager


class TestMCPAdminEndpoints:
    """Integration tests for MCP admin endpoints."""

    def test_mcp_status_endpoint_requires_admin(self):
        """Test that MCP status endpoint requires admin access."""
        from main import app
        client = TestClient(app)

        # Non-admin user should be denied
        r = client.get("/admin/mcp/status", headers={"X-User-Email": "user@example.com"})
        assert r.status_code in (302, 403)

    def test_mcp_status_endpoint_returns_data(self):
        """Test that MCP status endpoint returns expected data structure."""
        from main import app
        client = TestClient(app)

        # Admin user should get response
        r = client.get("/admin/mcp/status", headers={"X-User-Email": "admin@example.com"})
        assert r.status_code == 200

        data = r.json()
        assert "connected_servers" in data
        assert "configured_servers" in data
        assert "failed_servers" in data
        assert "auto_reconnect" in data
        assert "tool_counts" in data
        assert "prompt_counts" in data

        # Check auto_reconnect structure
        auto_reconnect = data["auto_reconnect"]
        assert "enabled" in auto_reconnect
        assert "base_interval" in auto_reconnect
        assert "max_interval" in auto_reconnect
        assert "backoff_multiplier" in auto_reconnect
        assert "running" in auto_reconnect

    def test_mcp_status_marks_failed_servers_not_connected(self):
        """Servers with recorded failures should not appear as connected."""
        from main import app
        client = TestClient(app)

        # Seed a fake failure in the MCP manager
        from atlas.infrastructure.app_factory import app_factory
        mcp = app_factory.get_mcp_manager()
        mcp._failed_servers["failing-server"] = {
            "last_attempt": time.time(),
            "attempt_count": 1,
            "error": "Simulated failure",
        }
        mcp.clients["failing-server"] = AsyncMock()
        mcp.available_tools["failing-server"] = {"tools": [], "config": {}}
        mcp.available_prompts["failing-server"] = {"prompts": [], "config": {}}

        r = client.get("/admin/mcp/status", headers={"X-User-Email": "admin@example.com"})
        assert r.status_code == 200

        data = r.json()
        assert "failing-server" not in data["connected_servers"]

    def test_mcp_reload_endpoint_requires_admin(self):
        """Test that MCP reload endpoint requires admin access."""
        from main import app
        client = TestClient(app)

        # Non-admin user should be denied
        r = client.post("/admin/mcp/reload", headers={"X-User-Email": "user@example.com"})
        assert r.status_code in (302, 403)

    def test_mcp_reconnect_endpoint_requires_admin(self):
        """Test that MCP reconnect endpoint requires admin access."""
        from main import app
        client = TestClient(app)

        # Non-admin user should be denied
        r = client.post("/admin/mcp/reconnect", headers={"X-User-Email": "user@example.com"})
        assert r.status_code in (302, 403)

    def test_mcp_reconnect_endpoint_returns_data(self):
        """Test that MCP reconnect endpoint returns expected data structure."""
        from main import app
        client = TestClient(app)

        # Admin user should get response
        r = client.post("/admin/mcp/reconnect", headers={"X-User-Email": "admin@example.com"})
        assert r.status_code == 200

        data = r.json()
        assert "message" in data
        assert "result" in data
        assert "current_servers" in data
        assert "failed_servers" in data
        assert "triggered_by" in data

    def test_admin_dashboard_includes_mcp_endpoints(self):
        """Test that admin dashboard lists MCP endpoints."""
        from main import app
        client = TestClient(app)

        r = client.get("/admin/", headers={"X-User-Email": "admin@example.com"})
        assert r.status_code == 200

        data = r.json()
        endpoints = data.get("available_endpoints", [])
        assert "/admin/mcp/reload" in endpoints
        assert "/admin/mcp/reconnect" in endpoints
        assert "/admin/mcp/status" in endpoints


class TestMCPFailedServerTracking:
    """Tests for tracking failed MCP server connections."""

    def test_record_server_failure_new_server(self):
        """Test recording first failure for a server."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager._failed_servers = {}

        manager._record_server_failure("test-server", "Connection refused")

        assert "test-server" in manager._failed_servers
        assert manager._failed_servers["test-server"]["attempt_count"] == 1
        assert manager._failed_servers["test-server"]["error"] == "Connection refused"
        assert "last_attempt" in manager._failed_servers["test-server"]

    def test_record_server_failure_existing_server(self):
        """Test recording additional failures for an already-failed server."""
        manager = MCPToolManager.__new__(MCPToolManager)
        initial_time = time.time() - 100
        manager._failed_servers = {
            "test-server": {
                "last_attempt": initial_time,
                "attempt_count": 2,
                "error": "Old error"
            }
        }

        manager._record_server_failure("test-server", "New error")

        assert manager._failed_servers["test-server"]["attempt_count"] == 3
        assert manager._failed_servers["test-server"]["error"] == "New error"
        assert manager._failed_servers["test-server"]["last_attempt"] > initial_time

    def test_clear_server_failure(self):
        """Test clearing failure tracking after successful connection."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager._failed_servers = {
            "test-server": {
                "last_attempt": time.time(),
                "attempt_count": 3,
                "error": "Some error"
            }
        }

        manager._clear_server_failure("test-server")

        assert "test-server" not in manager._failed_servers

    def test_clear_server_failure_nonexistent(self):
        """Test clearing a server that wasn't tracked (should not error)."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager._failed_servers = {}

        # Should not raise any exception
        manager._clear_server_failure("nonexistent-server")

        assert "nonexistent-server" not in manager._failed_servers

    def test_get_failed_servers(self):
        """Test getting failed servers info."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager._failed_servers = {
            "server1": {"attempt_count": 1, "error": "Error 1"},
            "server2": {"attempt_count": 3, "error": "Error 2"}
        }

        result = manager.get_failed_servers()

        assert result == manager._failed_servers
        # Verify it returns a copy, not the original dict
        assert result is not manager._failed_servers


class TestMCPBackoffCalculation:
    """Tests for exponential backoff calculation."""

    @patch('atlas.modules.mcp_tools.client.config_manager')
    def test_calculate_backoff_first_attempt(self, mock_config_manager):
        """Test backoff calculation for first retry attempt."""
        mock_settings = MagicMock()
        mock_settings.mcp_reconnect_interval = 60
        mock_settings.mcp_reconnect_max_interval = 300
        mock_settings.mcp_reconnect_backoff_multiplier = 2.0
        mock_config_manager.app_settings = mock_settings

        manager = MCPToolManager.__new__(MCPToolManager)

        delay = manager._calculate_backoff_delay(1)

        assert delay == 60  # Base interval for first attempt

    @patch('atlas.modules.mcp_tools.client.config_manager')
    def test_calculate_backoff_exponential(self, mock_config_manager):
        """Test exponential backoff for subsequent attempts."""
        mock_settings = MagicMock()
        mock_settings.mcp_reconnect_interval = 60
        mock_settings.mcp_reconnect_max_interval = 300
        mock_settings.mcp_reconnect_backoff_multiplier = 2.0
        mock_config_manager.app_settings = mock_settings

        manager = MCPToolManager.__new__(MCPToolManager)

        # Second attempt: 60 * 2^1 = 120
        assert manager._calculate_backoff_delay(2) == 120
        # Third attempt: 60 * 2^2 = 240
        assert manager._calculate_backoff_delay(3) == 240
        # Fourth attempt: 60 * 2^3 = 480, but capped at 300
        assert manager._calculate_backoff_delay(4) == 300

    @patch('atlas.modules.mcp_tools.client.config_manager')
    def test_calculate_backoff_max_cap(self, mock_config_manager):
        """Test that backoff is capped at max_interval."""
        mock_settings = MagicMock()
        mock_settings.mcp_reconnect_interval = 60
        mock_settings.mcp_reconnect_max_interval = 300
        mock_settings.mcp_reconnect_backoff_multiplier = 2.0
        mock_config_manager.app_settings = mock_settings

        manager = MCPToolManager.__new__(MCPToolManager)

        # Very high attempt count should still be capped
        delay = manager._calculate_backoff_delay(10)

        assert delay == 300


class TestMCPConfigReload:
    """Tests for MCP configuration hot-reload."""

    @patch('atlas.modules.mcp_tools.client.config_manager')
    def test_reload_config_updates_servers(self, mock_config_manager):
        """Test that reload_config updates server configuration."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"old-server": {"description": "Old"}}
        manager._failed_servers = {"old-server": {"attempt_count": 1}}

        # Mock new config
        mock_new_config = MagicMock()
        mock_server = MagicMock()
        mock_server.model_dump.return_value = {"description": "New"}
        mock_new_config.servers = {"new-server": mock_server}
        mock_config_manager.reload_mcp_config.return_value = mock_new_config

        result = manager.reload_config()

        assert "old-server" in result["removed"]
        assert "new-server" in result["added"]
        assert manager.servers_config == {"new-server": {"description": "New"}}
        # Old failed server tracking should be cleared
        assert "old-server" not in manager._failed_servers

    @patch('atlas.modules.mcp_tools.client.config_manager')
    def test_reload_config_preserves_unchanged(self, mock_config_manager):
        """Test that reload_config identifies unchanged servers."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"existing-server": {"description": "Existing"}}
        manager._failed_servers = {}

        # Mock config with same server
        mock_new_config = MagicMock()
        mock_server = MagicMock()
        mock_server.model_dump.return_value = {"description": "Updated"}
        mock_new_config.servers = {"existing-server": mock_server}
        mock_config_manager.reload_mcp_config.return_value = mock_new_config

        result = manager.reload_config()

        assert "existing-server" in result["unchanged"]
        assert result["added"] == []
        assert result["removed"] == []


@pytest.mark.asyncio
class TestMCPDiscoveryResilience:
    """Tests to ensure tool/prompt discovery tolerates removed servers."""

    async def test_discover_tools_skips_removed_server(self):
        """If a client exists without config, discovery should not crash."""
        manager = MCPToolManager.__new__(MCPToolManager)
        # Simulate one valid server and one removed server
        manager.servers_config = {"server-a": {"description": "A"}}
        manager._failed_servers = {}
        manager.clients = {
            "server-a": AsyncMock(),
            "removed-server": AsyncMock(),
        }

        async def fake_discover(server_name, client):  # noqa: ARG001
            if server_name == "server-a":
                return {"tools": [MagicMock(name="t1")], "config": {"description": "A"}}
            # Simulate that removed-server had a client but config was deleted
            raise RuntimeError("Simulated failure for removed-server")

        manager._discover_tools_for_server = fake_discover  # type: ignore[assignment]

        # Should complete without raising, and only include server-a in available_tools
        await manager.discover_tools()
        assert "server-a" in manager.available_tools
        assert "removed-server" not in manager.available_tools

    @pytest.mark.asyncio
    async def test_discover_prompts_skips_removed_server(self):
        """Prompt discovery should also skip servers missing from config."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"server-a": {"description": "A"}}
        manager._failed_servers = {}
        manager.clients = {
            "server-a": AsyncMock(),
            "removed-server": AsyncMock(),
        }

        async def fake_discover_prompts(server_name, client):  # noqa: ARG001
            if server_name == "server-a":
                return {"prompts": [MagicMock(name="p1")], "config": {"description": "A"}}
            raise RuntimeError("Simulated failure for removed-server")

        manager._discover_prompts_for_server = fake_discover_prompts  # type: ignore[assignment]

        await manager.discover_prompts()
        assert "server-a" in manager.available_prompts
        assert "removed-server" not in manager.available_prompts


@pytest.mark.asyncio
class TestMCPDiscoveryFailureTracking:
    """Tests for tracking discovery failures so admin status can reflect them."""

    async def test_tool_discovery_failure_records_failed_server(self):
        """Tool discovery exception should record server in _failed_servers."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"bad-server": {"description": "Bad"}}
        manager.clients = {"bad-server": AsyncMock()}
        manager._failed_servers = {}

        async def failing_discover(server_name, client):  # noqa: ARG001
            raise RuntimeError("Simulated tool discovery failure")

        manager._discover_tools_for_server = failing_discover  # type: ignore[assignment]

        await manager.discover_tools()

        assert "bad-server" in manager._failed_servers
        error = manager._failed_servers["bad-server"]["error"]
        assert "Simulated tool discovery failure" in error

    async def test_prompt_discovery_failure_records_failed_server(self):
        """Prompt discovery exception should record server in _failed_servers."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"bad-server": {"description": "Bad"}}
        manager.clients = {"bad-server": AsyncMock()}
        manager._failed_servers = {}

        async def failing_discover_prompts(server_name, client):  # noqa: ARG001
            raise RuntimeError("Simulated prompt discovery failure")

        manager._discover_prompts_for_server = failing_discover_prompts  # type: ignore[assignment]

        await manager.discover_prompts()

        assert "bad-server" in manager._failed_servers
        error = manager._failed_servers["bad-server"]["error"]
        assert "Simulated prompt discovery failure" in error


@pytest.mark.asyncio
class TestMCPReconnection:
    """Tests for MCP server reconnection functionality."""

    async def test_reconnect_skips_when_no_failed_servers(self):
        """Test that reconnect returns early when no servers have failed."""
        manager = MCPToolManager.__new__(MCPToolManager)
        manager._failed_servers = {}

        result = await manager.reconnect_failed_servers()

        assert result["attempted"] == []
        assert result["reconnected"] == []
        assert result["still_failed"] == []
        assert result["skipped_backoff"] == []

    @patch('atlas.modules.mcp_tools.client.config_manager')
    async def test_reconnect_respects_backoff(self, mock_config_manager):
        """When not forced, reconnect should skip servers still in backoff."""
        mock_settings = MagicMock()
        mock_settings.mcp_reconnect_interval = 60
        mock_settings.mcp_reconnect_max_interval = 300
        mock_settings.mcp_reconnect_backoff_multiplier = 2.0
        mock_config_manager.app_settings = mock_settings

        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"test-server": {"description": "Test"}}
        manager.clients = {}
        # Server failed just now, should be in backoff
        manager._failed_servers = {
            "test-server": {
                "last_attempt": time.time(),
                "attempt_count": 1,
                "error": "Connection refused"
            }
        }

        result = await manager.reconnect_failed_servers(force=False)

        assert result["attempted"] == []
        assert result["skipped_backoff"][0]["server"] == "test-server"

    @patch('atlas.modules.mcp_tools.client.config_manager')
    async def test_reconnect_attempts_after_backoff(self, mock_config_manager):
        """When backoff has elapsed, reconnect should attempt server."""
        mock_settings = MagicMock()
        mock_settings.mcp_reconnect_interval = 60
        mock_settings.mcp_reconnect_max_interval = 300
        mock_settings.mcp_reconnect_backoff_multiplier = 2.0
        mock_config_manager.app_settings = mock_settings

        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"test-server": {"description": "Test"}}
        manager.clients = {}
        # Server failed long ago, backoff period has passed
        manager._failed_servers = {
            "test-server": {
                "last_attempt": time.time() - 120,  # 2 minutes ago
                "attempt_count": 1,
                "error": "Connection refused"
            }
        }

        # Mock the initialization method to return None (still failing)
        manager._initialize_single_client = AsyncMock(return_value=None)

        result = await manager.reconnect_failed_servers(force=False)

        assert "test-server" in result["attempted"]
        assert "test-server" in result["still_failed"]
        manager._initialize_single_client.assert_called_once()

    @patch('atlas.modules.mcp_tools.client.config_manager')
    async def test_reconnect_force_ignores_backoff(self, mock_config_manager):
        """Forced reconnect should attempt even inside backoff window."""
        mock_settings = MagicMock()
        mock_settings.mcp_reconnect_interval = 60
        mock_settings.mcp_reconnect_max_interval = 300
        mock_settings.mcp_reconnect_backoff_multiplier = 2.0
        mock_config_manager.app_settings = mock_settings

        manager = MCPToolManager.__new__(MCPToolManager)
        manager.servers_config = {"test-server": {"description": "Test"}}
        manager.clients = {}
        # Server failed just now, so normally it would be in backoff
        manager._failed_servers = {
            "test-server": {
                "last_attempt": time.time(),
                "attempt_count": 1,
                "error": "Connection refused"
            }
        }

        # Mock the initialization method to return None (still failing)
        manager._initialize_single_client = AsyncMock(return_value=None)

        # With force=True, it should attempt despite backoff
        result = await manager.reconnect_failed_servers(force=True)

        assert "test-server" in result["attempted"]
        manager._initialize_single_client.assert_called_once()


class TestConfigManagerMCPReload:
    """Tests for ConfigManager MCP reload functionality."""

    @patch('atlas.modules.config.config_manager.ConfigManager._search_paths')
    @patch('atlas.modules.config.config_manager.ConfigManager._load_file_with_error_handling')
    @patch('atlas.modules.config.config_manager.ConfigManager._validate_mcp_compliance_levels')
    def test_reload_mcp_config_clears_cache(
        self, mock_validate, mock_load, mock_search
    ):
        """Test that reload_mcp_config clears the cached config."""
        from atlas.modules.config.config_manager import ConfigManager

        manager = ConfigManager()
        # Pre-populate cache
        manager._mcp_config = MagicMock()
        manager._tool_approvals_config = MagicMock()

        # Mock the config loading
        mock_search.return_value = []
        mock_load.return_value = {"test-server": {"description": "Test"}}

        manager.reload_mcp_config()

        # Cache should have been cleared and reloaded
        assert manager._tool_approvals_config is None or manager._tool_approvals_config != MagicMock()
