import json
import threading

from a2a_agent.config import AgentConfig
from a2a_agent.protocol import AgentCard
from a2a_agent import runtime
from a2a_agent import topics


class FakePublishInfo:
    def __init__(self) -> None:
        self.wait_calls = 0

    def wait_for_publish(self, timeout: int) -> None:
        self.wait_calls += 1


class FakeRuntimeClient:
    def __init__(self) -> None:
        self.subscriptions = {}
        self.publications: list[tuple[str, str, int, bool, FakePublishInfo]] = []

    def subscribe(self, topic: str, callback) -> None:
        self.subscriptions[topic] = callback

    def publish(
        self,
        topic: str,
        payload: str,
        *,
        qos: int,
        retain: bool,
    ) -> FakePublishInfo:
        info = FakePublishInfo()
        self.publications.append((topic, payload, qos, retain, info))
        return info

    def disconnect(self) -> None:
        return None

    def loop_stop(self) -> None:
        return None


class FakeConnectingClient(FakeRuntimeClient):
    def __init__(self) -> None:
        super().__init__()
        self.connection_events: list[str] = []
        self.on_connect_hooks = []

    def set_last_will(self, topic: str, payload: str, *, retain: bool) -> None:
        self.connection_events.append("last-will")

    def add_on_connect_hook(self, fn) -> None:
        self.connection_events.append("hook")
        self.on_connect_hooks.append(fn)

    def connect(self, host: str, port: int, *, tls: bool) -> None:
        self.connection_events.append("connect")

    def loop_start(self) -> None:
        self.connection_events.append("loop-start")

    def wait_connected(self) -> bool:
        self.connection_events.append("wait-connected")
        return True


def test_worker_reconnect_hook_republishes_online_status_and_card(monkeypatch) -> None:
    config = AgentConfig(
        room="team",
        device_name="worker-a",
        role="worker",
        token="worker-token",
    )
    fake = FakeRuntimeClient()
    captured: dict[str, object] = {}

    monkeypatch.setenv("A2A_SINGLE_INSTANCE", "0")

    def fake_connected_client(config, role, status_topic, on_connect_hook):
        captured["role"] = role
        captured["status_topic"] = status_topic
        captured["hook"] = on_connect_hook
        return fake

    monkeypatch.setattr(runtime, "load_config", lambda: config)
    monkeypatch.setattr(runtime, "_connected_client", fake_connected_client)
    monkeypatch.setattr(runtime, "_install_signal_handlers", lambda event: None)
    monkeypatch.setattr(runtime, "start_runtime_auto_update_monitor", lambda event: None)

    stop = threading.Event()
    stop.set()
    runtime.run_worker(stop)
    fake.publications.clear()

    hook = captured["hook"]
    assert callable(hook)
    stop.clear()
    hook(fake)

    assert captured["role"] == "worker"
    assert captured["status_topic"] == topics.agent_status_topic("team", "worker-a")
    assert [publication[0] for publication in fake.publications] == [
        topics.agent_status_topic("team", "worker-a"),
        topics.agent_card_topic("team", "worker-a"),
    ]
    status = json.loads(fake.publications[0][1])
    assert status["status"] == "online"
    assert AgentCard.from_json(fake.publications[1][1]).name == "worker-a"
    assert all(publication[2:4] == (1, True) for publication in fake.publications)
    assert all(publication[4].wait_calls == 0 for publication in fake.publications)

    fake.publications.clear()
    stop.set()
    hook(fake)
    assert fake.publications == []


def test_commander_reconnect_hook_republishes_online_status(monkeypatch) -> None:
    config = AgentConfig(
        room="team",
        device_name="commander",
        role="commander",
    )
    fake = FakeRuntimeClient()
    captured: dict[str, object] = {}

    def fake_connected_client(config, role, status_topic, on_connect_hook):
        captured["role"] = role
        captured["status_topic"] = status_topic
        captured["hook"] = on_connect_hook
        return fake

    monkeypatch.setattr(runtime, "load_config", lambda: config)
    monkeypatch.setattr(runtime, "_connected_client", fake_connected_client)
    monkeypatch.setattr(runtime, "_install_signal_handlers", lambda event: None)

    stop = threading.Event()
    stop.set()
    runtime.run_commander(stop)
    fake.publications.clear()

    hook = captured["hook"]
    assert callable(hook)
    stop.clear()
    hook(fake)

    assert captured["role"] == "commander"
    assert captured["status_topic"] == topics.agent_status_topic("team", "commander")
    assert len(fake.publications) == 1
    publication = fake.publications[0]
    assert publication[0] == topics.agent_status_topic("team", "commander")
    assert json.loads(publication[1])["status"] == "online"
    assert publication[2:4] == (1, True)
    assert publication[4].wait_calls == 0

    fake.publications.clear()
    stop.set()
    hook(fake)
    assert fake.publications == []


def test_connected_client_registers_hook_before_connect(monkeypatch) -> None:
    config = AgentConfig(
        room="team",
        device_name="worker-a",
        role="worker",
        token="worker-token",
    )
    fake = FakeConnectingClient()
    hook_clients = []
    status_topic = topics.agent_status_topic("team", "worker-a")
    monkeypatch.setattr(runtime, "MqttClient", lambda client_id: fake)

    client = runtime._connected_client(
        config,
        "worker",
        status_topic,
        lambda connected_client: hook_clients.append(connected_client),
    )

    assert client is fake
    assert fake.connection_events == [
        "last-will",
        "hook",
        "connect",
        "loop-start",
        "wait-connected",
    ]
    assert len(fake.on_connect_hooks) == 1
    fake.on_connect_hooks[0]()
    assert hook_clients == [fake]
