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

import asyncio
import json
from collections.abc import Sequence
from dataclasses import dataclass, field

import pytest

from ag2 import Agent, Context, MemoryStream, observer
from ag2.compact import CompactTrigger, TailWindowCompact
from ag2.events import (
    BaseEvent,
    ChildEvent,
    HaltEvent,
    HumanInputRequest,
    InertEvent,
    ModelMessage,
    ModelResponse,
    TaskCancelled,
    TaskCompleted,
    TaskStarted,
    ToolCallEvent,
    Usage,
    UsageEvent,
)
from ag2.events.task_events import TaskEvent
from ag2.knowledge import KnowledgePlugin, KnowledgeStore, MemoryKnowledgeStore
from ag2.observers import BaseObserver
from ag2.testing import TestConfig
from ag2.usage import UsageReport
from ag2.watch import EventWatch


@dataclass
class Journal:
    events: list[BaseEvent] = field(default_factory=list)

    async def record(self, event: BaseEvent) -> None:
        self.events.append(event)


class Diagnostic(InertEvent):
    text: str


@dataclass
class CompactionProbe:
    calls: int = 0

    async def compact(
        self, events: Sequence[BaseEvent], context: Context, store: KnowledgeStore | None
    ) -> list[BaseEvent]:
        self.calls += 1
        return list(events)


@pytest.mark.asyncio
async def test_custom_inert_event_does_not_trigger_ancestor_compaction() -> None:
    compaction = CompactionProbe()
    child = Agent("child", config=TestConfig(Diagnostic(text="diagnostic"), "done"))
    root = Agent(
        "root",
        config=TestConfig(ToolCallEvent(name="task_child", arguments='{"objective":"work"}'), "done"),
        tools=[child.as_tool(description="Work")],
        plugins=[
            KnowledgePlugin(
                store=MemoryKnowledgeStore(),
                compact=compaction,
                compact_trigger=CompactTrigger(max_events=7),
                expose_tool=False,
                write_event_log=False,
            )
        ],
    )

    reply = await root.ask("start")

    assert any(isinstance(event, Diagnostic) for event in await reply.history.get_events())
    assert compaction.calls == 0


def spending(name: str, tokens: int, *calls: ToolCallEvent, tools: list | None = None) -> Agent:
    return Agent(
        name,
        config=TestConfig(
            *calls,
            ModelResponse(ModelMessage(f"{name} done"), usage=Usage(total_tokens=tokens), model=f"{name}-model"),
        ),
        tools=tools or [],
    )


def three_levels() -> Agent:
    leaf = spending("leaf", 30)
    worker = spending(
        "worker",
        20,
        ToolCallEvent(name="task_leaf", arguments='{"objective": "dig"}'),
        tools=[leaf.as_tool(description="Dig")],
    )
    return spending(
        "root",
        10,
        ToolCallEvent(name="task_worker", arguments='{"objective": "organize"}'),
        tools=[worker.as_tool(description="Organize")],
    )


def reload(events: list[BaseEvent]) -> list[BaseEvent]:
    return [type(event).from_dict(json.loads(json.dumps(event.to_dict()))) for event in events]


def task_ids(events: list[BaseEvent]) -> dict[str, str]:
    return {event.agent_name: event.task_id for event in events if isinstance(event, TaskStarted)}


@pytest.mark.asyncio
async def test_stored_usage_names_its_spender_at_every_level() -> None:
    reply = await three_levels().ask("start")

    history = reload(list(await reply.history.get_events()))

    ids = task_ids(history)
    spent = sorted(
        (event.agent_name or "", event.task_id or "", list(event.path), event.usage.total_tokens)
        for event in history
        if isinstance(event, UsageEvent)
    )
    assert spent == sorted([
        ("", "", [], 10),
        ("worker", ids["worker"], [ids["worker"]], 20),
        ("leaf", ids["leaf"], [ids["worker"], ids["leaf"]], 30),
    ])
    assert UsageReport.from_events(history).total.total_tokens == 60
    assert UsageReport.from_events([*history, *reload(history)]).total.total_tokens == 60


@pytest.mark.asyncio
async def test_descendant_usage_keeps_its_spender_through_compaction() -> None:
    compaction = KnowledgePlugin(
        store=MemoryKnowledgeStore(),
        compact=TailWindowCompact(target=1),
        compact_trigger=CompactTrigger(max_events=1),
        expose_tool=False,
        write_event_log=False,
    )
    leaf = spending("leaf", 30)
    worker = Agent(
        "worker",
        config=TestConfig(ToolCallEvent(name="task_leaf", arguments='{"objective": "dig"}'), "worker done"),
        tools=[leaf.as_tool(description="Dig")],
        plugins=[compaction],
    )
    root = Agent(
        "root",
        config=TestConfig(ToolCallEvent(name="task_worker", arguments='{"objective": "organize"}'), "root done"),
        tools=[worker.as_tool(description="Organize")],
        plugins=[compaction],
    )

    reply = await root.ask("start")

    history = reload(list(await reply.history.get_events()))
    spent = [event for event in history if isinstance(event, UsageEvent)]
    assert [(event.agent_name, len(event.path), event.usage.total_tokens) for event in spent] == [("leaf", 2, 30)]
    assert UsageReport.from_events(history).total.total_tokens == 30


@pytest.mark.asyncio
async def test_typed_lifecycle_observer_on_root_sees_child_and_grandchild() -> None:
    journal = Journal()

    await three_levels().ask("start", observers=[observer(TaskStarted | TaskCompleted, callback=journal.record)])

    ids = task_ids(journal.events)
    assert [(type(event).__name__, event.agent_name, list(event.path)) for event in journal.events] == [
        ("TaskStarted", "worker", [ids["worker"]]),
        ("TaskStarted", "leaf", [ids["worker"], ids["leaf"]]),
        ("TaskCompleted", "leaf", [ids["worker"], ids["leaf"]]),
        ("TaskCompleted", "worker", [ids["worker"]]),
    ]


@pytest.mark.asyncio
async def test_lifecycle_events_are_not_wrapped_at_any_ancestor() -> None:
    reply = await three_levels().ask("start")

    history = list(await reply.history.get_events())
    assert not any(isinstance(event, ChildEvent) and isinstance(event.inner, TaskEvent) for event in history)
    assert {(type(event).__name__, event.agent_name) for event in history if isinstance(event, TaskEvent)} == {
        ("TaskStarted", "worker"),
        ("TaskStarted", "leaf"),
        ("TaskCompleted", "leaf"),
        ("TaskCompleted", "worker"),
    }


class LeafStopper:
    async def stop(self, event: TaskStarted) -> None:
        if event.agent_name == "leaf":
            await event.origin.send(HaltEvent(reason="not needed", source="root"))


@pytest.mark.asyncio
async def test_root_halts_grandchild_through_its_forwarded_lifecycle_event() -> None:
    journal = Journal()
    stopper = LeafStopper()

    reply = await three_levels().ask(
        "start",
        observers=[observer(TaskStarted, callback=stopper.stop), observer(TaskCancelled, callback=journal.record)],
    )

    assert reply.body == "root done"
    assert [(event.agent_name, len(event.path), event.reason) for event in journal.events] == [
        ("leaf", 2, "not needed")
    ]


class SpendByModel(BaseObserver):
    """Sums every ``UsageEvent`` it is shown, with no deduplication, like a cost observer."""

    def __init__(self) -> None:
        super().__init__("spend-by-model", watch=EventWatch(UsageEvent))
        self.tokens: dict[str | None, int] = {}

    async def process(self, events: list[BaseEvent], ctx: Context) -> None:
        for event in events:
            if isinstance(event, UsageEvent):
                self.tokens[event.model] = self.tokens.get(event.model, 0) + (event.usage.total_tokens or 0)


@pytest.mark.asyncio
async def test_usage_watch_on_root_counts_each_descendant_call_once() -> None:
    spend = SpendByModel()

    await three_levels().ask("start", observers=[spend])

    assert spend.tokens == {"root-model": 10, "worker-model": 20, "leaf-model": 30}


@dataclass
class Calls:
    count: int = 0

    async def probe(self) -> str:
        """Count a call."""
        self.count += 1
        return "probed"


@dataclass
class Hook:
    questions: list[str] = field(default_factory=list)

    async def answer(self, event: HumanInputRequest) -> str:
        self.questions.append(event.content)
        return "yes"


@pytest.mark.asyncio
async def test_forwarded_events_do_not_run_the_roots_tools_model_or_hitl() -> None:
    leaf_calls, root_calls, hook = Calls(), Calls(), Hook()
    leaf = Agent(
        "leaf",
        config=TestConfig(
            ToolCallEvent(name="probe", arguments="{}"),
            ModelResponse(ModelMessage("leaf done"), usage=Usage(total_tokens=30)),
        ),
        tools=[leaf_calls.probe],
    )
    worker = Agent(
        "worker",
        config=TestConfig(ToolCallEvent(name="task_leaf", arguments='{"objective": "dig"}'), "worker done"),
        tools=[leaf.as_tool(description="Dig")],
    )
    root = Agent(
        "root",
        config=TestConfig(ToolCallEvent(name="task_worker", arguments='{"objective": "organize"}'), "root done"),
        tools=[worker.as_tool(description="Organize"), root_calls.probe],
        hitl_hook=hook.answer,
    )

    reply = await root.ask("start")

    assert reply.body == "root done"
    assert (leaf_calls.count, root_calls.count, hook.questions) == (1, 0, [])
    history = list(await reply.history.get_events())
    assert len([event for event in history if isinstance(event, ModelResponse)]) == 2


@dataclass
class UsageHold:
    entered: asyncio.Event = field(default_factory=asyncio.Event)
    release: asyncio.Event = field(default_factory=asyncio.Event)
    context: Context | None = None

    async def hold(self, event: UsageEvent, ctx: Context) -> None:
        self.context = ctx
        self.entered.set()
        await self.release.wait()


@pytest.mark.asyncio
async def test_root_cancel_keeps_grandchild_usage_once_at_every_level() -> None:
    hold = UsageHold()
    leaf = Agent("leaf", config=TestConfig(UsageEvent(Usage(total_tokens=30)), "leaf done"))
    middle = Agent(
        "middle",
        config=TestConfig(ToolCallEvent(name="task_leaf", arguments='{"objective": "dig"}'), "middle done"),
        tools=[leaf.as_tool(description="Dig")],
        observers=[observer(UsageEvent, callback=hold.hold, interrupt=True)],
    )
    root = Agent(
        "root",
        config=TestConfig(ToolCallEvent(name="task_middle", arguments='{"objective": "organize"}'), "root done"),
        tools=[middle.as_tool(description="Organize")],
    )
    stream = MemoryStream()

    turn = asyncio.create_task(root.ask("start", stream=stream))
    await asyncio.wait_for(hold.entered.wait(), 2)
    turn.cancel()
    with pytest.raises(asyncio.CancelledError):
        await turn

    assert hold.context is not None
    for history, cancelled in (
        (list(await stream.history.get_events()), "middle"),
        (list(await hold.context.stream.history.get_events()), "leaf"),
    ):
        records = [index for index, event in enumerate(history) if isinstance(event, UsageEvent)]
        terminals = [
            index
            for index, event in enumerate(history)
            if isinstance(event, TaskCancelled) and event.agent_name == cancelled
        ]
        assert len(records) == 1, cancelled
        assert len(terminals) == 1, cancelled
        assert records[0] < terminals[0], cancelled
        assert UsageReport.from_events(history).total == Usage(total_tokens=30)