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

import asyncio
import logging
from collections.abc import Iterable
from contextlib import suppress
from uuid import UUID, uuid4

import pytest

from ag2.acp.sessions import (
    SessionBusyError,
    SessionConfig,
    SessionLimitError,
    SessionStore,
    UnknownSessionError,
)
from ag2.context import StreamId
from ag2.events import BaseEvent, ModelMessage
from ag2.history import MemoryStorage


class _SizedStorage(MemoryStorage):
    """A ``Storage`` that defines ``__len__``, so it is falsy until something is written."""

    def __init__(self) -> None:
        super().__init__()
        self.stream_ids: set[StreamId] = set()

    def __len__(self) -> int:
        return len(self.stream_ids)

    async def set_history(self, stream_id: StreamId, events: Iterable[BaseEvent]) -> None:
        self.stream_ids.add(stream_id)
        await super().set_history(stream_id, events)


@pytest.mark.asyncio
class TestIdentity:
    async def test_each_session_gets_a_distinct_id_and_stream(self) -> None:
        store = SessionStore()

        first = await store.create()
        second = await store.create()

        assert first.session_id != second.session_id
        assert first.stream_id != second.stream_id

    async def test_the_session_id_is_its_stream_id(self) -> None:
        """The id a Client holds is the key its history is stored under — nothing else is needed to find it."""
        store = SessionStore()

        session = await store.create()

        assert UUID(session.session_id) == session.stream_id
        assert session.session_id == session.stream_id.hex

    async def test_get_returns_the_same_session(self) -> None:
        store = SessionStore()

        created = await store.create()

        assert await store.get(created.session_id) is created

    async def test_unknown_id_raises(self) -> None:
        store = SessionStore()

        with pytest.raises(UnknownSessionError):
            await store.get("never-issued")

    async def test_client_context_is_captured_verbatim(self) -> None:
        store = SessionStore()

        session = await store.create(cwd="/work", additional_directories=["/extra"], meta={"ag2.space": {"room": "!r"}})

        assert session.cwd == "/work"
        assert session.additional_directories == ["/extra"]
        assert session.meta == {"ag2.space": {"room": "!r"}}


@pytest.mark.asyncio
class TestHistoryIsolation:
    async def test_two_sessions_never_share_history(self) -> None:
        store = SessionStore()
        first = await store.create()
        second = await store.create()

        await store.stream(first).history.replace([ModelMessage("only for first")])

        assert [e.content for e in await store.stream(first).history.get_events()] == ["only for first"]
        assert list(await store.stream(second).history.get_events()) == []

    async def test_history_accumulates_across_turns_of_one_session(self) -> None:
        store = SessionStore()
        session = await store.create()

        await store.stream(session).history.replace([ModelMessage("turn one")])
        # A *fresh* stream object per turn still reads the session's history back.
        later = store.stream(session)

        assert [e.content for e in await later.history.get_events()] == ["turn one"]

    async def test_a_session_keeps_one_stream_for_its_lifetime(self) -> None:
        """A stream carries an inbox and background tasks, not just history.

        Handing out a fresh object per turn would strand anything a background
        task delivered after its own turn finished.
        """
        store = SessionStore()
        session = await store.create()

        assert store.stream(session) is store.stream(session)

    async def test_a_message_enqueued_during_one_turn_survives_to_the_next(self) -> None:
        store = SessionStore()
        session = await store.create()

        store.stream(session).enqueue("late result")

        assert store.stream(session).pending_messages


@pytest.mark.asyncio
class TestSuppliedStorage:
    """A supplied backend holds the history even while it is empty, and so possibly falsy."""

    async def test_an_empty_backend_is_used(self) -> None:
        storage = _SizedStorage()
        assert not storage
        store = SessionStore(storage=storage)
        session = await store.create()

        await store.stream(session).history.replace([ModelMessage("kept")])

        assert [e.content for e in await storage.get_history(session.stream_id)] == ["kept"]

    async def test_an_empty_backend_flows_through_from_config(self) -> None:
        storage = _SizedStorage()
        store = SessionStore.from_config(SessionConfig(storage=storage))
        session = await store.create()

        await store.stream(session).history.replace([ModelMessage("kept")])

        assert [e.content for e in await storage.get_history(session.stream_id)] == ["kept"]


@pytest.mark.asyncio
class TestQueueing:
    async def test_concurrent_prompts_on_one_session_run_one_at_a_time(self) -> None:
        store = SessionStore()
        session = await store.create()
        order: list[str] = []

        async def prompt(label: str) -> None:
            async with session.turn():
                order.append(f"start {label}")
                await asyncio.sleep(0.01)
                order.append(f"end {label}")

        await asyncio.gather(prompt("a"), prompt("b"))

        # Interleaving would give "start a", "start b", "end a", "end b".
        assert order in (
            ["start a", "end a", "start b", "end b"],
            ["start b", "end b", "start a", "end a"],
        )

    async def test_the_running_turn_does_not_count_against_the_queue(self) -> None:
        store = SessionStore.from_config(SessionConfig(max_queued=1))
        session = await store.create()
        release = asyncio.Event()

        async def blocked() -> None:
            async with session.turn():
                await release.wait()

        running = asyncio.create_task(blocked())
        await asyncio.sleep(0)  # let it take the lock

        assert session.queued == 0  # the holder is running, not queued

        queued = asyncio.create_task(blocked())
        await asyncio.sleep(0)

        assert session.queued == 1

        release.set()
        await asyncio.gather(running, queued)

    async def test_queue_overflow_is_rejected(self) -> None:
        store = SessionStore.from_config(SessionConfig(max_queued=2))
        session = await store.create()
        release = asyncio.Event()

        async def blocked() -> None:
            async with session.turn():
                await release.wait()

        running = asyncio.create_task(blocked())
        await asyncio.sleep(0)  # let it take the lock
        waiters = [asyncio.create_task(blocked()) for _ in range(2)]
        await asyncio.sleep(0)  # let both queue behind it

        with pytest.raises(SessionBusyError):
            async with session.turn():
                pass  # pragma: no cover - the guard fires before the body

        release.set()
        await asyncio.gather(running, *waiters)

    async def test_queue_drains_and_the_session_is_reusable(self) -> None:
        store = SessionStore.from_config(SessionConfig(max_queued=1))
        session = await store.create()

        async with session.turn():
            pass

        assert session.queued == 0
        async with session.turn():
            pass


@pytest.mark.asyncio
class TestCancellation:
    async def test_cancel_stops_the_running_turn(self) -> None:
        store = SessionStore()
        session = await store.create()
        started = asyncio.Event()

        async def never_finishes() -> None:
            async with session.turn():
                started.set()
                await asyncio.Event().wait()

        session.turn_task = asyncio.create_task(never_finishes())
        await started.wait()

        await session.cancel()

        with pytest.raises(asyncio.CancelledError):
            await session.turn_task

    async def test_cancel_also_drops_prompts_queued_behind(self) -> None:
        store = SessionStore()
        session = await store.create()
        started = asyncio.Event()
        reached_agent = False

        async def running() -> None:
            async with session.turn():
                started.set()
                await asyncio.Event().wait()

        async def queued() -> None:
            nonlocal reached_agent
            async with session.turn():
                reached_agent = True

        session.turn_task = asyncio.create_task(running())
        await started.wait()
        waiter = asyncio.create_task(queued())
        await asyncio.sleep(0)

        await session.cancel()

        with pytest.raises(asyncio.CancelledError):
            await session.turn_task
        with pytest.raises(asyncio.CancelledError):
            await waiter
        assert reached_agent is False

    async def test_cancel_leaves_other_sessions_untouched(self) -> None:
        store = SessionStore()
        cancelled = await store.create()
        other = await store.create()
        finished = False

        async def work() -> None:
            nonlocal finished
            async with other.turn():
                await asyncio.sleep(0.01)
                finished = True

        task = asyncio.create_task(work())
        await asyncio.sleep(0)

        await cancelled.cancel()
        await task

        assert finished is True

    async def test_events_already_emitted_survive_a_cancel(self) -> None:
        store = SessionStore()
        session = await store.create()
        await store.stream(session).history.replace([ModelMessage("already streamed")])

        await session.cancel()

        assert [e.content for e in await store.stream(session).history.get_events()] == ["already streamed"]


@pytest.mark.asyncio
class TestEviction:
    async def test_lru_overflow_drops_the_oldest_session(self) -> None:
        store = SessionStore(max_sessions=2)

        first = await store.create()
        await store.create()
        await store.create()

        assert len(store) == 2
        with pytest.raises(UnknownSessionError):
            await store.get(first.session_id)

    async def test_touching_a_session_saves_it_from_eviction(self) -> None:
        store = SessionStore(max_sessions=2)
        first = await store.create()
        second = await store.create()

        await store.get(first.session_id)  # first is now most-recently used
        await store.create()

        assert await store.get(first.session_id) is first
        with pytest.raises(UnknownSessionError):
            await store.get(second.session_id)

    async def test_idle_sessions_expire(self) -> None:
        now = 0.0
        store = SessionStore(ttl=10.0, clock=lambda: now)
        session = await store.create()

        now = 11.0

        with pytest.raises(UnknownSessionError):
            await store.get(session.session_id)

    async def test_evicted_history_is_deleted(self) -> None:
        store = SessionStore(max_sessions=1)
        first = await store.create()
        await store.stream(first).history.replace([ModelMessage("gone")])

        await store.create()  # evicts first

        assert list(await store.stream(first).history.get_events()) == []

    async def test_a_session_under_recovery_is_not_evicted(self) -> None:
        """Holding the turn lock without a prompt — a repair, a load replay — is still work.

        Evicting mid-way would drop a session the Client is about to use, and
        by default its history with it.
        """
        store = SessionStore(max_sessions=1)
        session = await store.create()

        async with session.recovery():
            with pytest.raises(SessionLimitError):
                await store.create()

        assert await store.get(session.session_id) is session

    async def test_from_config_takes_a_clock(self) -> None:
        now = 0.0
        store = SessionStore.from_config(SessionConfig(ttl=10.0), clock=lambda: now)
        session = await store.create()

        now = 11.0

        with pytest.raises(UnknownSessionError):
            await store.get(session.session_id)


@pytest.mark.asyncio
class TestRetention:
    """``retain_history`` separates letting a session go from deleting its conversation."""

    async def test_retained_history_survives_aclose(self) -> None:
        store = SessionStore(retain_history=True)
        session = await store.create()
        await store.stream(session).history.replace([ModelMessage("kept")])

        await store.aclose()

        assert len(store) == 0
        assert [e.content for e in await store.stream(session).history.get_events()] == ["kept"]

    async def test_retained_history_survives_lru_eviction(self) -> None:
        store = SessionStore(max_sessions=1, retain_history=True)
        first = await store.create()
        await store.stream(first).history.replace([ModelMessage("kept")])

        await store.create()  # evicts first

        with pytest.raises(UnknownSessionError):
            await store.get(first.session_id)
        assert [e.content for e in await store.stream(first).history.get_events()] == ["kept"]

    async def test_retained_history_survives_ttl_expiry(self) -> None:
        now = 0.0
        store = SessionStore(ttl=10.0, retain_history=True, clock=lambda: now)
        session = await store.create()
        await store.stream(session).history.replace([ModelMessage("kept")])

        now = 11.0
        with pytest.raises(UnknownSessionError):
            await store.get(session.session_id)

        assert [e.content for e in await store.stream(session).history.get_events()] == ["kept"]

    async def test_close_always_drops_history_even_when_retained(self) -> None:
        """Retention is for sessions that were let go; a close means the conversation is over."""
        store = SessionStore(retain_history=True)
        session = await store.create()
        await store.stream(session).history.replace([ModelMessage("gone")])

        await store.close(session.session_id)

        assert list(await store.stream(session).history.get_events()) == []

    async def test_retention_flows_through_from_config(self) -> None:
        store = SessionStore.from_config(SessionConfig(retain_history=True, storage=MemoryStorage()))

        assert store.retains_history is True
        assert SessionStore.from_config(SessionConfig()).retains_history is False


class TestRetentionConfigWarns:
    """In-memory retention only grows the process; say so rather than let it look durable."""

    def test_retaining_in_the_default_storage_warns(self, caplog: pytest.LogCaptureFixture) -> None:
        with caplog.at_level(logging.WARNING, logger="ag2.acp.sessions"):
            SessionConfig(retain_history=True)

        assert any("in-memory" in record.message for record in caplog.records)

    def test_retaining_in_a_supplied_storage_is_quiet(self, caplog: pytest.LogCaptureFixture) -> None:
        with caplog.at_level(logging.WARNING, logger="ag2.acp.sessions"):
            SessionConfig(retain_history=True, storage=MemoryStorage())

        assert caplog.records == []


@pytest.mark.asyncio
class TestAdopt:
    """``get_or_adopt`` is the ``session/load`` seam: an id the store never issued, registered over its history."""

    async def test_adopt_registers_under_the_stream_id_hex(self) -> None:
        store = SessionStore()
        stream_id = uuid4()

        session, adopted = await store.get_or_adopt(stream_id, cwd="/work")

        assert adopted is True
        assert (session.session_id, session.stream_id, session.cwd) == (stream_id.hex, stream_id, "/work")
        assert await store.get(stream_id.hex) is session

    async def test_adopting_a_live_id_returns_the_existing_session_untouched(self) -> None:
        store = SessionStore()
        live = await store.create(cwd="/original")

        session, adopted = await store.get_or_adopt(live.stream_id, cwd="/other")

        assert session is live
        assert adopted is False
        assert live.cwd == "/original"  # context seeds an adopted session only

    async def test_an_adopted_session_reads_the_history_written_under_its_stream_id(self) -> None:
        """Two stores, one storage: what the first wrote, the second finds by id alone."""
        storage = MemoryStorage()
        first_store = SessionStore(storage=storage, retain_history=True)
        original = await first_store.create()
        await first_store.stream(original).history.replace([ModelMessage("from before")])
        await first_store.aclose()

        second_store = SessionStore(storage=storage)
        adopted, _ = await second_store.get_or_adopt(original.stream_id)

        assert [e.content for e in await second_store.stream(adopted).history.get_events()] == ["from before"]

    async def test_adopt_respects_the_cap(self) -> None:
        store = SessionStore(max_sessions=1)
        busy = await store.create()
        started = asyncio.Event()

        async def hold() -> None:
            async with busy.turn():
                started.set()
                await asyncio.Event().wait()

        task = asyncio.create_task(hold())
        busy.turn_task = task
        await started.wait()

        with pytest.raises(SessionLimitError):
            await store.get_or_adopt(uuid4())
        task.cancel()

    async def test_adopt_evicts_an_idle_session_to_make_room(self) -> None:
        store = SessionStore(max_sessions=1)
        idle = await store.create()

        adopted, _ = await store.get_or_adopt(uuid4())

        assert await store.get(adopted.session_id) is adopted
        with pytest.raises(UnknownSessionError):
            await store.get(idle.session_id)

    async def test_an_adopted_session_starts_idle(self) -> None:
        store = SessionStore()

        session, _ = await store.get_or_adopt(uuid4())

        assert session.is_idle

    async def test_forget_unregisters_without_touching_history(self) -> None:
        """Undoing a refused adoption must leave the conversation as it was found."""
        store = SessionStore()
        session, _ = await store.get_or_adopt(uuid4())
        await store.stream(session).history.replace([ModelMessage("not ours to delete")])

        await store.forget(session.session_id)

        with pytest.raises(UnknownSessionError):
            await store.get(session.session_id)
        assert [e.content for e in await store.stream(session).history.get_events()] == ["not ours to delete"]
        await store.forget("never-issued")  # nothing to undo is not an error


@pytest.mark.asyncio
class TestClose:
    async def test_close_drops_the_session_and_its_history(self) -> None:
        store = SessionStore()
        session = await store.create()
        await store.stream(session).history.replace([ModelMessage("gone")])

        await store.close(session.session_id)

        with pytest.raises(UnknownSessionError):
            await store.get(session.session_id)
        assert list(await store.stream(session).history.get_events()) == []

    async def test_close_cancels_live_work_and_waits_for_it(self) -> None:
        """Unlike eviction, an explicit close *does* stop work in progress.

        It also waits for the turn to unwind before deleting history, so the
        agent's last writes cannot land in a store that has already been purged.
        """
        store = SessionStore()
        session = await store.create()
        started = asyncio.Event()

        async def running() -> None:
            async with session.turn():
                started.set()
                await asyncio.Event().wait()

        task = asyncio.create_task(running())
        session.turn_task = task
        await started.wait()

        await store.close(session.session_id)

        assert task.done()
        assert task.cancelled()

    async def test_closing_an_unknown_session_raises(self) -> None:
        store = SessionStore()

        with pytest.raises(UnknownSessionError):
            await store.close("never-issued")

    async def test_aclose_drops_every_session(self) -> None:
        store = SessionStore()
        first = await store.create()
        second = await store.create()

        await store.aclose()

        assert len(store) == 0
        for session in (first, second):
            with pytest.raises(UnknownSessionError):
                await store.get(session.session_id)


def test_invalid_bounds_are_rejected() -> None:
    with pytest.raises(ValueError):
        SessionStore(max_sessions=0)
    with pytest.raises(ValueError):
        SessionStore(ttl=0)
    with pytest.raises(ValueError):
        SessionStore(max_queued=0)


@pytest.mark.asyncio
class TestEvictionSparesLiveWork:
    """An eviction policy is a memory bound, not a licence to kill running turns."""

    @staticmethod
    async def _busy(session: object) -> "asyncio.Task[None]":
        started = asyncio.Event()

        async def hold() -> None:
            async with session.turn():  # type: ignore[attr-defined]
                started.set()
                await asyncio.Event().wait()

        task = asyncio.create_task(hold())
        session.turn_task = task  # type: ignore[attr-defined]
        await started.wait()
        return task

    async def test_pressure_on_the_cap_does_not_cancel_a_running_turn(self) -> None:
        """Admission fails instead — see :class:`TestTheCapIsAHardBound`."""
        store = SessionStore(max_sessions=1)
        busy = await store.create()
        task = await self._busy(busy)

        with pytest.raises(SessionLimitError):
            await store.create()
        await asyncio.sleep(0.01)

        assert not task.done()
        task.cancel()

    async def test_pressure_on_the_cap_does_not_delete_a_running_session_s_history(self) -> None:
        store = SessionStore(max_sessions=1)
        busy = await store.create()
        await store.stream(busy).history.replace([ModelMessage("mid-turn work")])
        task = await self._busy(busy)

        with pytest.raises(SessionLimitError):
            await store.create()
        await asyncio.sleep(0.01)

        assert [e.content for e in await store.stream(busy).history.get_events()] == ["mid-turn work"]
        task.cancel()

    async def test_an_idle_session_is_still_evicted(self) -> None:
        store = SessionStore(max_sessions=1)
        idle = await store.create()

        await store.create()

        with pytest.raises(UnknownSessionError):
            await store.get(idle.session_id)

    async def test_the_idle_victim_is_chosen_over_the_busy_one(self) -> None:
        store = SessionStore(max_sessions=2)
        busy = await store.create()
        idle = await store.create()
        task = await self._busy(busy)

        await store.create()

        assert await store.get(busy.session_id) is busy
        with pytest.raises(UnknownSessionError):
            await store.get(idle.session_id)
        task.cancel()

    async def test_ttl_does_not_expire_a_session_mid_turn(self) -> None:
        now = 0.0
        store = SessionStore(ttl=10.0, clock=lambda: now)
        busy = await store.create()
        task = await self._busy(busy)

        now = 11.0
        await store.create()  # any access runs the expiry sweep

        assert await store.get(busy.session_id) is busy
        assert not task.done()
        task.cancel()


@pytest.mark.asyncio
class TestEveryIssuedIdIsUsable:
    """A session id handed to a Client must still be there when it is used."""

    async def test_room_is_made_before_the_session_exists(self) -> None:
        """A session that gets created is never a candidate for its own eviction.

        Admission now runs *before* the session is built, so the id handed back
        cannot already have been evicted to make room for itself.
        """
        store = SessionStore(max_sessions=2)
        idle = await store.create()

        fresh = await store.create()

        assert await store.get(fresh.session_id) is fresh
        assert await store.get(idle.session_id) is idle

    async def test_the_cap_still_holds_when_an_idle_victim_exists(self) -> None:
        store = SessionStore(max_sessions=1)
        old = await store.create()

        fresh = await store.create()

        assert await store.get(fresh.session_id) is fresh
        with pytest.raises(UnknownSessionError):
            await store.get(old.session_id)


@pytest.mark.asyncio
class TestTtlMeasuresIdleTime:
    """Expiry is about being unused, not about when the last prompt started."""

    async def test_a_turn_longer_than_the_ttl_does_not_expire_its_session(self) -> None:
        now = 0.0
        store = SessionStore(ttl=10.0, clock=lambda: now)
        session = await store.create()
        await store.stream(session).history.replace([ModelMessage("slow work")])

        async with session.turn():
            now = 25.0  # the turn itself outlives the TTL

        assert await store.get(session.session_id) is session
        assert [e.content for e in await store.stream(session).history.get_events()] == ["slow work"]

    async def test_the_clock_restarts_when_a_turn_finishes(self) -> None:
        now = 0.0
        store = SessionStore(ttl=10.0, clock=lambda: now)
        session = await store.create()

        async with session.turn():
            now = 25.0

        now = 30.0  # only 5s idle since the turn ended
        assert await store.get(session.session_id) is session

        now = 45.0  # now genuinely idle past the TTL
        with pytest.raises(UnknownSessionError):
            await store.get(session.session_id)


@pytest.mark.asyncio
class TestTheCapIsAHardBound:
    """A cap that stretches while turns are in flight is not a cap at all."""

    @staticmethod
    async def _busy(store: SessionStore) -> "tuple[object, asyncio.Task[None]]":
        session = await store.create()
        started = asyncio.Event()

        async def hold() -> None:
            async with session.turn():  # type: ignore[attr-defined]
                started.set()
                await asyncio.Event().wait()

        task = asyncio.create_task(hold())
        session.turn_task = task  # type: ignore[attr-defined]
        await started.wait()
        return session, task

    async def test_admission_is_refused_when_nothing_can_be_evicted(self) -> None:
        store = SessionStore(max_sessions=1)
        _busy, task = await self._busy(store)

        with pytest.raises(SessionLimitError):
            await store.create()

        assert len(store) == 1
        task.cancel()

    async def test_a_client_cannot_grow_the_registry_by_staying_busy(self) -> None:
        store = SessionStore(max_sessions=2)
        tasks = []
        for _ in range(2):
            _session, task = await self._busy(store)
            tasks.append(task)

        for _ in range(5):
            with pytest.raises(SessionLimitError):
                await store.create()

        assert len(store) == 2
        for task in tasks:
            task.cancel()

    async def test_a_slot_frees_up_once_a_turn_finishes(self) -> None:
        store = SessionStore(max_sessions=1)
        _session, task = await self._busy(store)

        with pytest.raises(SessionLimitError):
            await store.create()

        task.cancel()
        with suppress(asyncio.CancelledError):
            await task

        fresh = await store.create()
        assert await store.get(fresh.session_id) is fresh
        assert len(store) == 1