"""Unit tests for SSHConnectionManager."""

from __future__ import annotations

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

import asyncssh
import pytest

from vmcluster_mcp.exec.ssh import SSHConnectionManager
from vmcluster_mcp.models import ExecResult


# Helper to build a mock asyncssh run result
def _mock_run_result(exit_status: int = 0, stdout: str = "output", stderr: str = "") -> MagicMock:
    r = MagicMock()
    r.exit_status = exit_status
    r.stdout = stdout
    r.stderr = stderr
    return r


# Helper to build a mock asyncssh connection
def _mock_conn(run_result=None, run_exc=None) -> AsyncMock:
    conn = AsyncMock()
    if run_exc is not None:
        conn.run = AsyncMock(side_effect=run_exc)
    else:
        conn.run = AsyncMock(return_value=run_result or _mock_run_result())
    conn.close = MagicMock()
    return conn


@pytest.mark.asyncio
async def test_execute_fresh_connection_returns_exec_result():
    """Test 1: execute() on a fresh manager opens a new connection and returns ExecResult."""
    manager = SSHConnectionManager()
    mock_conn = _mock_conn(run_result=_mock_run_result(exit_status=0, stdout="hello", stderr=""))

    with patch("asyncssh.connect", new_callable=AsyncMock, return_value=mock_conn):
        result = await manager.execute("10.0.0.1", 22, "root", "/key", "echo hello")

    assert isinstance(result, ExecResult)
    assert result.exit_code == 0
    assert result.stdout == "hello"
    assert result.stderr == ""
    assert result.timed_out is False
    assert result.duration_s >= 0


@pytest.mark.asyncio
async def test_execute_reuses_cached_connection():
    """Test 2: execute() called twice reuses the cached connection (one asyncssh.connect call)."""
    manager = SSHConnectionManager()
    mock_conn = _mock_conn(run_result=_mock_run_result())

    with patch("asyncssh.connect", new_callable=AsyncMock, return_value=mock_conn) as mock_connect:
        await manager.execute("10.0.0.1", 22, "root", "/key", "cmd1")
        await manager.execute("10.0.0.1", 22, "root", "/key", "cmd2")

    # Only one connection should have been opened
    assert mock_connect.call_count == 1
    # But run was called twice
    assert mock_conn.run.call_count == 2


@pytest.mark.asyncio
async def test_execute_different_hosts_open_separate_connections():
    """Test 3: execute() for different hosts opens separate connections."""
    manager = SSHConnectionManager()

    conn1 = _mock_conn(run_result=_mock_run_result(stdout="host1"))
    conn2 = _mock_conn(run_result=_mock_run_result(stdout="host2"))

    with patch("asyncssh.connect", new_callable=AsyncMock, side_effect=[conn1, conn2]) as mock_connect:
        r1 = await manager.execute("10.0.0.1", 22, "root", "/key", "hostname")
        r2 = await manager.execute("10.0.0.2", 22, "root", "/key", "hostname")

    assert mock_connect.call_count == 2
    assert r1.stdout == "host1"
    assert r2.stdout == "host2"


@pytest.mark.asyncio
async def test_execute_reconnects_on_stale_connection():
    """Test 4: When cached connection raises ConnectionLost, execute() retries with a fresh connection."""
    manager = SSHConnectionManager()

    # First connection: returns stale on run
    stale_conn = _mock_conn(run_exc=asyncssh.ConnectionLost("connection lost"))
    # Second connection: fresh, succeeds
    fresh_conn = _mock_conn(run_result=_mock_run_result(exit_status=0, stdout="reconnected"))

    with patch("asyncssh.connect", new_callable=AsyncMock, side_effect=[stale_conn, fresh_conn]) as mock_connect:
        result = await manager.execute("10.0.0.1", 22, "root", "/key", "echo ok")

    # Should have opened two connections (first stale, second fresh)
    assert mock_connect.call_count == 2
    assert result.exit_code == 0
    assert result.stdout == "reconnected"
    assert result.timed_out is False


@pytest.mark.asyncio
async def test_execute_returns_error_on_double_failure():
    """Test 5: When retry also fails, execute() returns ExecResult with error in stderr."""
    manager = SSHConnectionManager()

    exc = asyncssh.DisconnectError(asyncssh.DISC_CONNECTION_LOST, "connection lost")
    failing_conn = _mock_conn(run_exc=exc)

    with patch("asyncssh.connect", new_callable=AsyncMock, return_value=failing_conn):
        result = await manager.execute("10.0.0.1", 22, "root", "/key", "echo ok")

    assert result.exit_code == -1
    assert result.timed_out is False
    assert result.stderr  # some error message


@pytest.mark.asyncio
async def test_invalidate_removes_cached_connection():
    """Test 6: invalidate(ip) removes the cached connection; next execute() opens a new one."""
    manager = SSHConnectionManager()

    conn1 = _mock_conn(run_result=_mock_run_result(stdout="first"))
    conn2 = _mock_conn(run_result=_mock_run_result(stdout="second"))

    with patch("asyncssh.connect", new_callable=AsyncMock, side_effect=[conn1, conn2]) as mock_connect:
        await manager.execute("10.0.0.1", 22, "root", "/key", "cmd")
        assert mock_connect.call_count == 1

        await manager.invalidate("10.0.0.1")
        conn1.close.assert_called_once()

        r2 = await manager.execute("10.0.0.1", 22, "root", "/key", "cmd")
        assert mock_connect.call_count == 2
        assert r2.stdout == "second"


@pytest.mark.asyncio
async def test_close_all_closes_all_connections():
    """Test 7: close_all() closes all cached connections and clears the dict."""
    manager = SSHConnectionManager()

    conn1 = _mock_conn(run_result=_mock_run_result(stdout="h1"))
    conn2 = _mock_conn(run_result=_mock_run_result(stdout="h2"))

    with patch("asyncssh.connect", new_callable=AsyncMock, side_effect=[conn1, conn2]):
        await manager.execute("10.0.0.1", 22, "root", "/key", "cmd")
        await manager.execute("10.0.0.2", 22, "root", "/key", "cmd")

    await manager.close_all()

    conn1.close.assert_called_once()
    conn2.close.assert_called_once()
    assert len(manager._connections) == 0


@pytest.mark.asyncio
async def test_execute_timeout_returns_timed_out():
    """Test 8: Command timeout returns ExecResult with timed_out=True."""
    manager = SSHConnectionManager()
    # asyncssh.TimeoutError requires all ProcessError constructor args
    timeout_exc = asyncssh.TimeoutError(None, "sleep 999", None, None, None, None, "", "")
    mock_conn = _mock_conn(run_exc=timeout_exc)

    with patch("asyncssh.connect", new_callable=AsyncMock, return_value=mock_conn):
        result = await manager.execute("10.0.0.1", 22, "root", "/key", "sleep 999", timeout_s=1)

    assert result.timed_out is True
    assert result.exit_code == -1
    assert "timed out" in result.stderr.lower()


@pytest.mark.asyncio
async def test_execute_nonzero_exit_code():
    """Test 9: Non-zero exit code returns correct exit_code, stdout, stderr."""
    manager = SSHConnectionManager()
    mock_conn = _mock_conn(run_result=_mock_run_result(exit_status=1, stdout="", stderr="error msg"))

    with patch("asyncssh.connect", new_callable=AsyncMock, return_value=mock_conn):
        result = await manager.execute("10.0.0.1", 22, "root", "/key", "false")

    assert result.exit_code == 1
    assert result.stdout == ""
    assert result.stderr == "error msg"
    assert result.timed_out is False
