# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

"""Discussion adapter (round_robin) + multi-party handshake tests.

Covers:

* 5-way handshake: hub posts ``EV_CHANNEL_INVITE`` to every invitee,
  collects all-or-nothing acks, transitions to ``ACTIVE``.
* Round-robin speaker rotation: ``state.expected_next_speaker`` cycles
  through ``metadata.participants`` ``order`` after each accepted
  ``EV_TEXT``.
* ``validate_send`` rejects out-of-turn sends.
* Partial reject fails the channel (the handshake is all-or-nothing).
* ``Hub.hydrate()`` re-folds the WAL through ``DiscussionAdapter.fold``
  and recovers ``expected_next_speaker``.
* Hub auto-registers the adapter on ``Hub.open``.
* ``validate_create`` rejects unsupported ordering modes.

This suite uses ``ScriptedConfig`` so it runs offline and fast.
"""

import contextlib
from collections.abc import Awaitable, Callable

import pytest

from ag2 import Agent
from ag2.knowledge import DiskKnowledgeStore, MemoryKnowledgeStore
from ag2.network import (
    EV_CHANNEL_INVITE,
    EV_CHANNEL_INVITE_ACK,
    EV_CHANNEL_INVITE_REJECT,
    EV_CHANNEL_OPENED,
    EV_TEXT,
    Envelope,
    Hub,
    HubClient,
    LocalLink,
    Passport,
    Resume,
)
from ag2.network.adapters.discussion import (
    DISCUSSION_TYPE,
    ORDERING_ROUND_ROBIN,
    DiscussionAdapter,
    DiscussionState,
)
from ag2.network.channel import ChannelState
from ag2.network.client.agent_client import AgentClient
from ag2.network.errors import ProtocolError
from ag2.testing import TestConfig

from ._helpers import ScriptedConfig, wait_for_text_count


def _agent(name: str, *events: object) -> Agent:
    return Agent(name=name, config=TestConfig(*events))


def _scripted_agent(name: str, *replies: str) -> Agent:
    return Agent(name=name, config=ScriptedConfig(*replies))


def _make_rejecter(client: AgentClient) -> Callable[[Envelope], Awaitable[None]]:
    """Build a notify handler that explicitly rejects every invite.

    Used to drive the partial-reject path. Defined at module level
    (not inside a test) so the closure is created once at registration.
    """

    async def _reject(envelope: Envelope) -> None:
        if envelope.event_type != EV_CHANNEL_INVITE:
            return
        rejection = Envelope(
            channel_id=envelope.channel_id,
            sender_id=client.agent_id,
            audience=None,
            event_type=EV_CHANNEL_INVITE_REJECT,
            event_data={"reason": "not interested"},
            causation_id=envelope.envelope_id,
        )
        with contextlib.suppress(Exception):
            await client.send_envelope(rejection)

    return _reject


@pytest.mark.asyncio
async def test_default_discussion_adapter_registered_on_open() -> None:
    """``Hub.open`` auto-registers DiscussionAdapter for ``discussion@v1``."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0)
    adapter = hub._adapters.get((DISCUSSION_TYPE, 1))
    assert isinstance(adapter, DiscussionAdapter)
    await hub.close()


@pytest.mark.asyncio
async def test_discussion_validate_create_rejects_unsupported_ordering() -> None:
    """V1 only ships round_robin; dynamic / static raise at create time."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0)

    alice = await hub.register(_agent("alice"))
    bob = await hub.register(_agent("bob"))

    with pytest.raises(ProtocolError, match="ordering"):
        await hub.create_channel(
            creator_id=alice.agent_id,
            manifest_type=DISCUSSION_TYPE,
            participants=[bob.agent_id],
            knobs={"ordering": "dynamic"},
        )

    await hub.close()


@pytest.mark.asyncio
async def test_discussion_5_way_handshake_transitions_to_active() -> None:
    """5 invitees all auto-ack; channel activates with full participant list."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0)
    link = LocalLink(hub)

    names = ["alice", "bob", "carol", "dave", "erin"]
    clients: list[AgentClient] = []
    for name in names:
        hc = HubClient(link, hub=hub)
        client = await hc.register(_agent(name), Passport(name=name), Resume())
        clients.append(client)

    alice = clients[0]
    targets = [c.agent_id for c in clients[1:]]
    channel = await alice.open(
        type=DISCUSSION_TYPE,
        target=targets,
        knobs={"ordering": ORDERING_ROUND_ROBIN},
    )

    assert channel.state == ChannelState.ACTIVE
    assert channel.metadata.pending_acks == []
    assert channel.metadata.rejected_by == []
    assert {p.agent_id for p in channel.metadata.participants} == {c.agent_id for c in clients}

    wal = await hub.read_wal(channel.channel_id)
    invite_count = sum(1 for e in wal if e.event_type == EV_CHANNEL_INVITE)
    ack_count = sum(1 for e in wal if e.event_type == EV_CHANNEL_INVITE_ACK)
    assert invite_count == 4  # one invite per non-creator participant
    assert ack_count == 4
    assert any(e.event_type == EV_CHANNEL_OPENED for e in wal)

    state = hub._adapter_states[channel.channel_id]
    assert isinstance(state, DiscussionState)
    assert state.expected_next_speaker == alice.agent_id
    assert state.participant_order[0] == alice.agent_id

    for hc in [c._hub_client for c in clients]:
        await hc.close()
    await hub.close()


@pytest.mark.asyncio
async def test_discussion_round_robin_advances_through_participants() -> None:
    """Manual sends in turn order succeed; out-of-turn raises ProtocolError."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0)

    # Use plain TestConfig agents and skip the LLM auto-handler — we
    # only exercise the adapter by posting envelopes manually.
    alice = await hub.register(_agent("alice"), attach_plugin=False)
    bob = await hub.register(_agent("bob"), attach_plugin=False)
    carol = await hub.register(_agent("carol"), attach_plugin=False)

    # Auto-ack invites (handlers are off — install minimal ackers).
    for client in (bob, carol):
        client.on_envelope(_make_auto_acker(client))

    channel = await alice.open(
        type=DISCUSSION_TYPE,
        target=[bob.agent_id, carol.agent_id],
        knobs={"ordering": ORDERING_ROUND_ROBIN},
    )
    assert channel.state == ChannelState.ACTIVE

    # Round 1: alice → bob → carol.
    await channel.send("alice 1")
    state = hub._adapter_states[channel.channel_id]
    assert state.expected_next_speaker == bob.agent_id

    bob_envelope = Envelope(
        channel_id=channel.channel_id,
        sender_id=bob.agent_id,
        audience=None,
        event_type=EV_TEXT,
        event_data={"text": "bob 1"},
    )
    await hub.post_envelope(bob_envelope)
    state = hub._adapter_states[channel.channel_id]
    assert state.expected_next_speaker == carol.agent_id

    carol_envelope = Envelope(
        channel_id=channel.channel_id,
        sender_id=carol.agent_id,
        audience=None,
        event_type=EV_TEXT,
        event_data={"text": "carol 1"},
    )
    await hub.post_envelope(carol_envelope)
    state = hub._adapter_states[channel.channel_id]
    assert state.expected_next_speaker == alice.agent_id  # cycle back
    assert state.turn_count == 3
    assert state.last_speaker_id == carol.agent_id

    await hub.close()


@pytest.mark.asyncio
async def test_discussion_rejects_out_of_turn_send() -> None:
    """Sending out of round-robin order raises ProtocolError before WAL write."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0)

    alice = await hub.register(_agent("alice"))
    bob = await hub.register(_agent("bob"))
    carol = await hub.register(_agent("carol"))

    channel = await alice.open(
        type=DISCUSSION_TYPE,
        target=[bob.agent_id, carol.agent_id],
        knobs={"ordering": ORDERING_ROUND_ROBIN},
    )

    # Carol jumps the queue before alice has spoken.
    bad = Envelope(
        channel_id=channel.channel_id,
        sender_id=carol.agent_id,
        audience=None,
        event_type=EV_TEXT,
        event_data={"text": "I'm jumping the queue"},
    )
    with pytest.raises(ProtocolError, match="expects"):
        await hub.post_envelope(bad)

    await hub.close()


@pytest.mark.asyncio
async def test_discussion_partial_reject_fails_channel() -> None:
    """V1 all-or-nothing: any reject during handshake closes the channel."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0, invite_ack_timeout=2.0)

    alice = await hub.register(_agent("alice"))
    bob = await hub.register(_agent("bob"))
    # Carol rejects every invite.
    carol = await hub.register(_agent("carol"), attach_plugin=False)
    carol.on_envelope(_make_rejecter(carol))

    with pytest.raises(ProtocolError, match="rejected"):
        await alice.open(
            type=DISCUSSION_TYPE,
            target=[bob.agent_id, carol.agent_id],
            knobs={"ordering": ORDERING_ROUND_ROBIN},
        )

    await hub.close()


@pytest.mark.asyncio
async def test_discussion_hydrate_refolds_round_robin_state(tmp_path) -> None:
    """Re-opening hub mid-discussion recovers expected_next_speaker."""
    store = DiskKnowledgeStore(str(tmp_path))
    hub1 = await Hub.open(store, ttl_sweep_interval=0)

    alice = await hub1.register(_agent("alice"), attach_plugin=False)
    bob = await hub1.register(_agent("bob"), attach_plugin=False)
    carol = await hub1.register(_agent("carol"), attach_plugin=False)
    bob.on_envelope(_make_auto_acker(bob))
    carol.on_envelope(_make_auto_acker(carol))

    channel = await alice.open(
        type=DISCUSSION_TYPE,
        target=[bob.agent_id, carol.agent_id],
        knobs={"ordering": ORDERING_ROUND_ROBIN},
    )
    await channel.send("alice opens the floor")

    # Mid-discussion (bob is up next) → tear down the hub.
    await hub1.close()

    # Reopen against the same store.
    store2 = DiskKnowledgeStore(str(tmp_path))
    hub2 = await Hub.open(store2, ttl_sweep_interval=0)

    state = hub2._adapter_states[channel.channel_id]
    assert isinstance(state, DiscussionState)
    assert state.expected_next_speaker == bob.agent_id
    assert state.last_speaker_id == alice.agent_id
    assert state.turn_count == 1
    assert state.participant_order == [alice.agent_id, bob.agent_id, carol.agent_id]

    await hub2.close()


@pytest.mark.asyncio
async def test_discussion_llm_driven_round_robin_3_way() -> None:
    """End-to-end LLM-driven 3-way round-robin: each agent's handler fires
    only when the adapter's validate_send would accept their reply."""
    store = MemoryKnowledgeStore()
    hub = await Hub.open(store, ttl_sweep_interval=0)

    # alice's manual send opens turn 1; her LLM speaks again on her turn 2.
    alice = await hub.register(_scripted_agent("alice", "alice 2"))
    bob = await hub.register(_scripted_agent("bob", "bob 1"))
    carol = await hub.register(_scripted_agent("carol", "carol 1"))

    channel = await alice.open(
        type=DISCUSSION_TYPE,
        target=[bob.agent_id, carol.agent_id],
        knobs={"ordering": ORDERING_ROUND_ROBIN},
    )
    await channel.send("alice 1")

    # 4 EV_TEXT envelopes total: alice 1 / bob 1 / carol 1 / alice 2.
    # On alice's turn 3, bob has no script left → empty body → halt.
    wal = await wait_for_text_count(hub, channel.channel_id, expected=4)
    text_envelopes = [e for e in wal if e.event_type == EV_TEXT]
    assert [e.event_data["text"] for e in text_envelopes] == [
        "alice 1",
        "bob 1",
        "carol 1",
        "alice 2",
    ]
    assert [e.sender_id for e in text_envelopes] == [
        alice.agent_id,
        bob.agent_id,
        carol.agent_id,
        alice.agent_id,
    ]

    state = hub._adapter_states[channel.channel_id]
    assert state.expected_next_speaker == bob.agent_id
    assert state.turn_count == 4

    await hub.close()


def _make_auto_acker(client: AgentClient) -> Callable[[Envelope], Awaitable[None]]:
    """Minimal handler that auto-acks invites and ignores everything else.

    Used by tests that disable the default handler (which also runs
    LLM logic) but still need the multi-party handshake to complete.
    """

    async def _ack(envelope: Envelope) -> None:
        if envelope.event_type != EV_CHANNEL_INVITE:
            return
        ack = Envelope(
            channel_id=envelope.channel_id,
            sender_id=client.agent_id,
            audience=None,
            event_type=EV_CHANNEL_INVITE_ACK,
            event_data={"channel_id": envelope.channel_id},
            causation_id=envelope.envelope_id,
        )
        with contextlib.suppress(Exception):
            await client.send_envelope(ack)

    return _ack