import pytest
from app.services.presence import (
    add_agent_connection,
    remove_agent_connection,
    get_agent_connection,
    get_connected_agents,
    set_agent_status,
    get_available_agents,
    connected_agents
)

class MockWebSocket:
    def __init__(self, client_id: str = "mock"):
        self.client_id = client_id
        self.sent_messages = []
        self.is_closed = False

    async def send_text(self, text: str):
        if self.is_closed:
            raise RuntimeError("WebSocket is closed")
        self.sent_messages.append(text)

    async def close(self, code: int = 1000, reason: str = ""):
        self.is_closed = True

@pytest.fixture(autouse=True)
def reset_agent_pool():
    connected_agents.clear()
    yield
    connected_agents.clear()

def test_add_two_agents_simultaneously():
    ws1 = MockWebSocket("ws1")
    ws2 = MockWebSocket("ws2")

    conn1 = add_agent_connection("agent_alpha", ws1)
    conn2 = add_agent_connection("agent_beta", ws2)

    assert conn1["agentId"] == "agent_alpha"
    assert conn1["status"] == "AVAILABLE"
    assert conn2["agentId"] == "agent_beta"
    assert conn2["status"] == "AVAILABLE"

    all_agents = get_connected_agents()
    assert len(all_agents) == 2
    assert "agent_alpha" in all_agents
    assert "agent_beta" in all_agents

def test_retrieve_agent_independently():
    ws1 = MockWebSocket("ws1")
    ws2 = MockWebSocket("ws2")

    add_agent_connection("agent_alpha", ws1)
    add_agent_connection("agent_beta", ws2)

    alpha_info = get_agent_connection("agent_alpha")
    assert alpha_info is not None
    assert alpha_info["websocket"] == ws1

    beta_info = get_agent_connection("agent_beta")
    assert beta_info is not None
    assert beta_info["websocket"] == ws2

    non_existent = get_agent_connection("agent_unknown")
    assert non_existent is None

def test_disconnect_one_agent_does_not_affect_others():
    ws1 = MockWebSocket("ws1")
    ws2 = MockWebSocket("ws2")

    add_agent_connection("agent_alpha", ws1)
    add_agent_connection("agent_beta", ws2)

    # Disconnect agent_alpha
    remove_agent_connection("agent_alpha", ws1)

    assert get_agent_connection("agent_alpha") is None
    assert get_agent_connection("agent_beta") is not None
    assert len(get_connected_agents()) == 1

def test_same_agent_reconnect_prevents_duplicates():
    ws1 = MockWebSocket("ws1")
    ws2 = MockWebSocket("ws2_reconnected")

    add_agent_connection("agent_alpha", ws1)
    assert len(get_connected_agents()) == 1
    assert get_agent_connection("agent_alpha")["websocket"] == ws1

    # Same agent reconnects with a new websocket
    add_agent_connection("agent_alpha", ws2)

    all_agents = get_connected_agents()
    assert len(all_agents) == 1
    assert get_agent_connection("agent_alpha")["websocket"] == ws2

    # Attempting to disconnect old stale socket should NOT remove the active replacement
    remove_agent_connection("agent_alpha", ws1)
    assert get_agent_connection("agent_alpha") is not None
    assert get_agent_connection("agent_alpha")["websocket"] == ws2

def test_agent_presence_status_transitions():
    ws = MockWebSocket("ws1")
    add_agent_connection("agent_alpha", ws, status_val="AVAILABLE")

    assert get_agent_connection("agent_alpha")["status"] == "AVAILABLE"

    # Transition to BUSY
    ok_busy = set_agent_status("agent_alpha", "BUSY")
    assert ok_busy is True
    assert get_agent_connection("agent_alpha")["status"] == "BUSY"

    # Transition to OFFLINE
    ok_offline = set_agent_status("agent_alpha", "OFFLINE")
    assert ok_offline is True
    assert get_agent_connection("agent_alpha")["status"] == "OFFLINE"

    # Transition with invalid status returns False
    ok_invalid = set_agent_status("agent_alpha", "UNKNOWN_STATUS")
    assert ok_invalid is False

def test_get_available_agents_filtering():
    ws1 = MockWebSocket("ws1")
    ws2 = MockWebSocket("ws2")
    ws3 = MockWebSocket("ws3")

    add_agent_connection("agent_alpha", ws1, status_val="AVAILABLE")
    add_agent_connection("agent_beta", ws2, status_val="BUSY")
    add_agent_connection("agent_gamma", ws3, status_val="AVAILABLE")

    avail = get_available_agents()
    assert len(avail) == 2
    assert "agent_alpha" in avail
    assert "agent_gamma" in avail
    assert "agent_beta" not in avail