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

import asyncio
from contextlib import AsyncExitStack, ExitStack
from datetime import UTC, datetime
from uuid import uuid4

import pytest
from dirty_equals import HasAttributes, IsInstance

from ag2 import Agent, Context, TaskInject
from ag2.context import StreamId
from ag2.events import (
    BaseEvent,
    HaltEvent,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    TaskCancelled,
    TaskCompleted,
    TaskExpired,
    TaskFailed,
    TaskProgress,
    TaskStarted,
    ToolCallEvent,
    Usage,
    UsageEvent,
)
from ag2.exceptions import ConfigNotProvidedError
from ag2.history import MemoryStorage
from ag2.stream import MemoryStream
from ag2.task import Task, TaskSpec, TaskState
from ag2.testing import TestConfig, TrackingConfig
from ag2.tools.subagents import DeliveryError

TASK_EVENTS = (TaskStarted, TaskProgress, TaskCompleted, TaskFailed, TaskCancelled, TaskExpired)


def task_events(events: list[BaseEvent]) -> list[BaseEvent]:
    return [event for event in events if isinstance(event, TASK_EVENTS)]


async def fail() -> str:
    """Fail the model turn."""
    raise ValueError("work failed")


class SlowTool:
    def __init__(self) -> None:
        self.entered = asyncio.Event()
        self.cleaned_up = asyncio.Event()

    async def wait(self) -> str:
        """Wait until the invocation is cancelled."""
        self.entered.set()
        try:
            await asyncio.Event().wait()
        finally:
            self.cleaned_up.set()
        return "unreachable"


class FailingCleanup:
    """A user observer whose teardown raises."""

    def __init__(self, error: BaseException) -> None:
        self.error = error

    def register(self, stack: ExitStack | AsyncExitStack, context: Context) -> None:
        stack.callback(self.close)

    def close(self) -> None:
        raise self.error


class TerminalProbe:
    """Tries to message a delegated worker once its terminal is observed."""

    def __init__(self) -> None:
        self.outcomes: list[str] = []

    async def check(self, event: TaskCompleted) -> None:
        try:
            await event.origin.send(HaltEvent(reason="already finished", source="test"))
        except DeliveryError:
            self.outcomes.append("retired")
        else:
            self.outcomes.append("still addressable")


class StallingStorage(MemoryStorage):
    """A store whose first history read after the model reply hangs, like a stalled round-trip."""

    def __init__(self) -> None:
        super().__init__()
        self.replied = False
        self.stalled = asyncio.Event()

    async def save_event(self, event: BaseEvent, context: Context) -> None:
        if isinstance(event, ModelResponse):
            self.replied = True
        await super().save_event(event, context)

    async def get_history(self, stream_id: StreamId) -> list[BaseEvent]:
        if self.replied and not self.stalled.is_set():
            self.stalled.set()
            await asyncio.Event().wait()
        return list(await super().get_history(stream_id))


class UsageWritesFail(MemoryStorage):
    """A store that cannot record a delegated worker's usage."""

    async def save_event(self, event: BaseEvent, context: Context) -> None:
        if isinstance(event, UsageEvent):
            raise ConnectionError("usage store unavailable")
        await super().save_event(event, context)


class BeforeFailure(BaseEvent):
    """Mark the last model event before an execution failure."""


class StallingRetentionStorage(MemoryStorage):
    """Lose source usage, then stall its retention after the model fails."""

    def __init__(self, source_id: StreamId) -> None:
        super().__init__()
        self.source_id = source_id
        self.armed = False
        self.stalled = asyncio.Event()

    async def save_event(self, event: BaseEvent, context: Context) -> None:
        if context.stream.id == self.source_id:
            if isinstance(event, UsageEvent):
                return
            if isinstance(event, BeforeFailure):
                self.armed = True
        await super().save_event(event, context)

    async def get_history(self, stream_id: StreamId) -> list[BaseEvent]:
        if stream_id == self.source_id and self.armed and not self.stalled.is_set():
            self.stalled.set()
            await asyncio.Event().wait()
        return list(await super().get_history(stream_id))


class LosesThenFailsUsage(MemoryStorage):
    """A store that silently loses the first usage record, then cannot record usage at all."""

    def __init__(self) -> None:
        super().__init__()
        self.lost = False

    async def save_event(self, event: BaseEvent, context: Context) -> None:
        if isinstance(event, UsageEvent):
            if not self.lost:
                self.lost = True
                return
            raise ConnectionError("usage store unavailable")
        await super().save_event(event, context)


class Abort(BaseException):
    """An interrupt-style error that is not an ``Exception``."""


def abort() -> str:
    """Abort the model turn."""
    raise Abort


async def expire_on_start(event: TaskStarted, ctx: Context) -> None:
    await ctx.send(TaskExpired(task_id=event.task_id, agent_name=event.agent_name, objective=event.objective))


def fail_on_start(event: TaskStarted) -> None:
    raise ValueError("startup observer failed")


async def report_progress(task: TaskInject) -> str:
    """Report progress on the caller's task."""
    if task is not None:
        await task.progress({"from_tool": True})
    return "ok"


async def index_documents(ctx: Context, tool: SlowTool) -> None:
    async with Task(owner_id="indexer", spec=TaskSpec(title="index"), context=ctx):
        await tool.wait()


def delegating_lead(worker: Agent, *replies: str) -> Agent:
    return Agent(
        "lead",
        config=TestConfig(
            ToolCallEvent(id="call-1", name="task_worker", arguments='{"objective": "summarise the report"}'),
            *replies,
        ),
        tools=[worker.as_tool(description="Summarise documents")],
    )


@pytest.mark.asyncio
@pytest.mark.parametrize("entry", ["ask", "run"])
async def test_untagged_calls_publish_no_task_events(entry: str) -> None:
    agent = Agent("analyst", config=TestConfig("done"))

    if entry == "ask":
        reply = await agent.ask("analyse the data")
    else:
        async with agent.run("analyse the data") as run:
            reply = await run.result()

    assert task_events(await reply.history.get_events()) == []


@pytest.mark.asyncio
async def test_tagged_ask_reports_its_task_without_leaking_the_prompt() -> None:
    tracking = TrackingConfig(TestConfig("done"))
    agent = Agent("analyst", config=tracking)

    reply = await agent.ask("confidential: customer list", capability="analysis")

    started, completed = task_events(await reply.history.get_events())
    assert started == IsInstance(TaskStarted) & HasAttributes(
        agent_name="analyst",
        objective="analyst",
        spec=TaskSpec(title="analyst", capability="analysis"),
        expires_at=None,
        parent_tool_call_id=None,
    )
    assert completed == IsInstance(TaskCompleted) & HasAttributes(
        task_id=started.task_id,
        result="done",
        task_stream=reply.context.stream.id,
    )
    [call] = tracking.calls
    assert task_events(list(call.messages)) == []


@pytest.mark.asyncio
async def test_ttl_publishes_a_deadline() -> None:
    agent = Agent("analyst", config=TestConfig("done"))

    reply = await agent.ask("analyse", ttl_seconds=60)

    [started, _] = task_events(await reply.history.get_events())
    assert started == HasAttributes(spec=TaskSpec(title="analyst"))
    assert datetime.fromisoformat(started.expires_at) > datetime.now(UTC)


@pytest.mark.asyncio
async def test_expiry_reported_by_an_observer_is_final_and_does_not_stop_the_run() -> None:
    stream = MemoryStream()
    stream.where(TaskStarted).subscribe(expire_on_start, sync_to_thread=False)
    agent = Agent("analyst", config=TestConfig("done"))

    reply = await agent.ask("analyse", stream=stream, ttl_seconds=1)

    assert reply.body == "done"
    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskExpired]


@pytest.mark.asyncio
async def test_run_publishes_the_terminal_after_the_block_exits() -> None:
    agent = Agent("analyst", config=TestConfig("done"))

    async with agent.run("analyse", capability="analysis") as run:
        reply = await run.result()
        assert [type(event) for event in task_events(await reply.history.get_events())] == [TaskStarted]

    assert [type(event) for event in task_events(await reply.history.get_events())] == [TaskStarted, TaskCompleted]


@pytest.mark.asyncio
async def test_run_left_without_driving_publishes_nothing() -> None:
    stream = MemoryStream()
    agent = Agent("analyst", config=TestConfig("unused"))

    async with agent.run("analyse", capability="analysis", stream=stream):
        pass

    assert task_events(await stream.history.get_events()) == []


@pytest.mark.asyncio
async def test_run_started_and_abandoned_is_cancelled_before_calling_the_model() -> None:
    stream = MemoryStream()
    tracking = TrackingConfig(TestConfig("unused"))
    agent = Agent("analyst", config=tracking)

    async with agent.run("analyse", capability="analysis", stream=stream) as run:
        run.start()

    assert tracking.calls == []
    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskCancelled]


@pytest.mark.asyncio
@pytest.mark.parametrize("entry", ["ask", "run"])
async def test_failing_tool_fails_the_task_and_raises(entry: str) -> None:
    stream = MemoryStream()
    agent = Agent("analyst", config=TestConfig(ToolCallEvent(name="fail")), tools=[fail])

    with pytest.raises(ValueError, match="work failed"):
        if entry == "ask":
            await agent.ask("analyse", capability="analysis", stream=stream)
        else:
            async with agent.run("analyse", capability="analysis", stream=stream) as run:
                await run.result()

    started, failed = task_events(await stream.history.get_events())
    assert started == IsInstance(TaskStarted)
    assert failed == IsInstance(TaskFailed) & HasAttributes(task_id=started.task_id, error=IsInstance(ValueError))


@pytest.mark.asyncio
async def test_missing_model_config_fails_the_task() -> None:
    stream = MemoryStream()

    with pytest.raises(ConfigNotProvidedError):
        await Agent("analyst").ask("analyse", capability="analysis", stream=stream)

    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskFailed]


@pytest.mark.asyncio
async def test_cancelled_ask_is_reported_after_tool_cleanup() -> None:
    stream = MemoryStream()
    tool = SlowTool()
    agent = Agent("analyst", config=TestConfig(ToolCallEvent(name="wait")), tools=[tool.wait])

    driver = asyncio.create_task(agent.ask("analyse", capability="analysis", stream=stream))
    await asyncio.wait_for(tool.entered.wait(), 1)
    driver.cancel()

    with pytest.raises(asyncio.CancelledError):
        await driver
    assert tool.cleaned_up.is_set()
    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskCancelled]


@pytest.mark.asyncio
async def test_run_left_while_a_tool_runs_is_cancelled_after_tool_cleanup() -> None:
    stream = MemoryStream()
    tool = SlowTool()
    agent = Agent("analyst", config=TestConfig(ToolCallEvent(name="wait")), tools=[tool.wait])

    async with agent.run("analyse", capability="analysis", stream=stream) as run:
        run.start()
        await asyncio.wait_for(tool.entered.wait(), 1)

    assert tool.cleaned_up.is_set()
    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskCancelled]


@pytest.mark.asyncio
async def test_cancelled_standalone_task_is_reported_as_cancelled() -> None:
    ctx = Context(MemoryStream())
    tool = SlowTool()

    driver = asyncio.create_task(index_documents(ctx, tool))
    await asyncio.wait_for(tool.entered.wait(), 1)
    driver.cancel()

    with pytest.raises(asyncio.CancelledError):
        await driver
    assert [type(event) for event in task_events(await ctx.stream.history.get_events())] == [
        TaskStarted,
        TaskCancelled,
    ]


@pytest.mark.asyncio
@pytest.mark.parametrize("entry", ["reply.ask", "reply.run", "resume"])
async def test_each_tagged_continuation_is_a_new_task(entry: str) -> None:
    agent = Agent("analyst", config=TestConfig("first", "second", "third"))
    first = await agent.ask("first", capability="analysis")

    if entry == "reply.ask":
        second = await first.ask("second", capability="followup")
    elif entry == "reply.run":
        async with first.run("second", capability="followup") as run:
            second = await run.result()
    else:
        second = await agent.resume(
            *await first.history.get_events(),
            ModelRequest.ensure_request(["second"]),
            stream=first.context.stream,
            capability="followup",
        )
    await second.ask("third")

    starts = [event for event in await second.history.get_events() if isinstance(event, TaskStarted)]
    assert [event.spec.capability for event in starts] == ["analysis", "followup"]
    assert starts[0].task_id != starts[1].task_id


@pytest.mark.asyncio
async def test_tagged_run_keeps_the_callers_task_visible_to_tools() -> None:
    ctx = Context(MemoryStream())
    agent = Agent(
        "analyst",
        config=TestConfig(ToolCallEvent(name="report_progress"), "done"),
        tools=[report_progress],
    )

    async with Task(owner_id="pipeline", spec=TaskSpec(title="pipeline"), context=ctx) as task:
        await agent.ask("analyse", stream=ctx.stream, dependencies=ctx.dependencies, capability="analysis")

    assert task.metadata.progress == {"from_tool": True}
    assert task.state == TaskState.COMPLETED


@pytest.mark.asyncio
async def test_delegation_reports_the_workers_task_to_the_parent() -> None:
    worker = Agent("worker", config=TestConfig(ModelResponse(ModelMessage("summary"), usage=Usage(prompt_tokens=3))))
    lead = delegating_lead(worker, "done")

    reply = await lead.ask("handle the report")

    started, completed = task_events(await reply.history.get_events())
    assert started == IsInstance(TaskStarted) & HasAttributes(
        agent_name="worker",
        objective="summarise the report",
        parent_tool_call_id="call-1",
    )
    assert completed == IsInstance(TaskCompleted) & HasAttributes(
        task_id=started.task_id,
        result="summary",
        usage=HasAttributes(prompt_tokens=3),
    )
    assert (await reply.usage()).total.prompt_tokens == 3


@pytest.mark.asyncio
async def test_tagged_parent_terminal_includes_the_workers_usage() -> None:
    worker = Agent("worker", config=TestConfig(ModelResponse(ModelMessage("summary"), usage=Usage(prompt_tokens=3))))
    lead = delegating_lead(worker, "done")

    reply = await lead.ask("handle the report", capability="triage")

    completed = [event for event in await reply.history.get_events() if isinstance(event, TaskCompleted)]
    assert completed == [
        HasAttributes(agent_name="worker", usage=HasAttributes(prompt_tokens=3)),
        HasAttributes(agent_name="lead", usage=HasAttributes(prompt_tokens=3)),
    ]


@pytest.mark.asyncio
async def test_failed_delegation_is_reported_and_the_parent_continues() -> None:
    worker = Agent("worker", config=TestConfig(ToolCallEvent(name="fail")), tools=[fail])
    lead = delegating_lead(worker, "recovered")

    reply = await lead.ask("handle the report")

    assert reply.body == "recovered"
    started, failed = task_events(await reply.history.get_events())
    assert failed == IsInstance(TaskFailed) & HasAttributes(task_id=started.task_id, error=IsInstance(ValueError))


@pytest.mark.asyncio
async def test_finished_worker_cannot_be_messaged_from_its_terminal() -> None:
    stream = MemoryStream()
    probe = TerminalProbe()
    stream.where(TaskCompleted).subscribe(probe.check, sync_to_thread=False)
    lead = delegating_lead(Agent("worker", config=TestConfig("summary")), "done")

    await lead.ask("handle the report", stream=stream)

    assert probe.outcomes == ["retired"]


@pytest.mark.asyncio
@pytest.mark.parametrize("entry", ["ask", "run", "reply.ask", "reply.run"])
async def test_observer_cleanup_failure_is_what_the_caller_and_observers_see(entry: str) -> None:
    error = ValueError("cleanup failed")
    stream = MemoryStream()
    agent = Agent("analyst", config=TestConfig("initial", "done"))
    options = {"capability": "analysis", "observers": [FailingCleanup(error)]}

    with pytest.raises(ValueError) as raised:
        if entry == "ask":
            await agent.ask("analyse", stream=stream, **options)
        elif entry == "run":
            async with agent.run("analyse", stream=stream, **options) as run:
                await run.result()
        else:
            previous = await agent.ask("initial", stream=stream)
            if entry == "reply.ask":
                await previous.ask("analyse", **options)
            else:
                async with previous.run("analyse", **options) as run:
                    await run.result()

    assert raised.value is error
    started, failed = task_events(await stream.history.get_events())
    assert failed == IsInstance(TaskFailed) & HasAttributes(task_id=started.task_id, error=error)


@pytest.mark.asyncio
async def test_failing_start_observer_fails_the_task_and_frees_the_stream() -> None:
    stream = MemoryStream()
    subscription = stream.where(TaskStarted).subscribe(fail_on_start, sync_to_thread=False)
    agent = Agent("analyst", config=TestConfig("done"))

    with pytest.raises(ValueError, match="startup observer failed"):
        await agent.ask("analyse", capability="analysis", stream=stream)
    stream.unsubscribe(subscription)

    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskFailed]
    assert (await asyncio.wait_for(agent.ask("next", stream=stream), 1)).body == "done"


@pytest.mark.asyncio
@pytest.mark.timeout(5)
async def test_cancellation_while_the_store_reads_usage_is_still_reported() -> None:
    storage = StallingStorage()
    stream = MemoryStream(storage=storage)
    agent = Agent("analyst", config=TestConfig("done"))

    driver = asyncio.create_task(agent.ask("analyse", capability="analysis", stream=stream))
    await asyncio.wait_for(storage.stalled.wait(), 1)
    driver.cancel()

    with pytest.raises(asyncio.CancelledError):
        await driver
    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskCancelled]


@pytest.mark.asyncio
@pytest.mark.timeout(5)
async def test_cancellation_during_failed_worker_usage_retention_overrides_the_failure() -> None:
    source_id = uuid4()
    storage = StallingRetentionStorage(source_id)
    stream = MemoryStream(storage=storage)
    worker = Agent(
        "worker",
        config=TestConfig(UsageEvent(Usage(total_tokens=30)), BeforeFailure(), ValueError("model failed")),
    )
    lead = Agent(
        "lead",
        config=TestConfig(ToolCallEvent(name="task_worker", arguments='{"objective": "work"}'), "done"),
        tools=[worker.as_tool(description="Work", stream=MemoryStream(id=source_id, storage=storage))],
    )

    driver = asyncio.create_task(lead.ask("work", stream=stream, capability="lead"))
    await asyncio.wait_for(storage.stalled.wait(), 1)
    driver.cancel()

    with pytest.raises(asyncio.CancelledError):
        await driver
    assert task_events(await stream.history.get_events()) == [
        IsInstance(TaskStarted) & HasAttributes(agent_name="lead"),
        IsInstance(TaskStarted) & HasAttributes(agent_name="worker"),
        IsInstance(TaskCancelled) & HasAttributes(agent_name="worker"),
        IsInstance(TaskCancelled) & HasAttributes(agent_name="lead"),
    ]


@pytest.mark.asyncio
@pytest.mark.timeout(5)
async def test_usage_retention_storage_failure_preserves_worker_cancellation() -> None:
    stream = MemoryStream()
    worker = Agent("worker", config=TestConfig(UsageEvent(Usage(total_tokens=30)), asyncio.CancelledError()))
    lead = Agent(
        "lead",
        config=TestConfig(ToolCallEvent(name="task_worker", arguments='{"objective": "work"}'), "done"),
        tools=[worker.as_tool(description="Work", stream=MemoryStream(storage=LosesThenFailsUsage()))],
    )

    with pytest.raises(asyncio.CancelledError):
        await lead.ask("work", stream=stream, capability="lead")

    assert task_events(await stream.history.get_events()) == [
        IsInstance(TaskStarted) & HasAttributes(agent_name="lead"),
        IsInstance(TaskStarted) & HasAttributes(agent_name="worker"),
        IsInstance(TaskCancelled) & HasAttributes(agent_name="worker"),
        IsInstance(TaskCancelled) & HasAttributes(agent_name="lead"),
    ]


@pytest.mark.asyncio
@pytest.mark.timeout(5)
@pytest.mark.parametrize("storage", [UsageWritesFail, LosesThenFailsUsage])
async def test_unrecordable_worker_usage_fails_the_delegation_and_the_parent_continues(
    storage: type[MemoryStorage],
) -> None:
    stream = MemoryStream(storage=storage())
    worker = Agent("worker", config=TestConfig(ModelResponse(ModelMessage("summary"), usage=Usage(prompt_tokens=3))))
    lead = delegating_lead(worker, "recovered")

    reply = await asyncio.wait_for(lead.ask("handle the report", stream=stream), 1)

    assert reply.body == "recovered"
    started, failed = task_events(await stream.history.get_events())
    assert failed == IsInstance(TaskFailed) & HasAttributes(task_id=started.task_id, error=IsInstance(ConnectionError))


@pytest.mark.asyncio
@pytest.mark.timeout(5)
async def test_interrupt_raised_by_a_tool_fails_the_task_and_frees_the_stream() -> None:
    stream = MemoryStream()
    agent = Agent("analyst", config=TestConfig(ToolCallEvent(name="abort")), tools=[abort])

    with pytest.raises(Abort):
        async with agent.run("analyse", capability="analysis", stream=stream) as run:
            await run.result()

    assert [type(event) for event in task_events(await stream.history.get_events())] == [TaskStarted, TaskFailed]
    assert (
        await asyncio.wait_for(Agent("next", config=TestConfig("done")).ask("next", stream=stream), 1)
    ).body == "done"