import unittest

from turbo_agent_core.schema.events import (
	ContentActionDeltaEvent,
	ContentActionDeltaPayload,
	ContentActionEndEvent,
	ContentActionEndPayload,
	ContentActionResultDeltaEvent,
	ContentActionResultDeltaPayload,
	ContentActionResultEndEvent,
	ContentActionResultEndPayload,
	ContentActionResultStartEvent,
	ContentActionResultStartPayload,
	ContentActionStartEvent,
	ContentActionStartPayload,
	ContentAgentResultDeltaEvent,
	ContentAgentResultDeltaPayload,
	ContentAgentResultEndEvent,
	ContentAgentResultEndPayload,
	ContentAgentResultStartEvent,
	ContentAgentResultStartPayload,
	ContentReasoningDeltaEvent,
	ContentReasoningDeltaPayload,
	ContentReasoningEndEvent,
	ContentReasoningEndPayload,
	ContentReasoningStartEvent,
	ContentReasoningStartPayload,
	ContentTextDeltaEvent,
	ContentTextDeltaPayload,
	ContentTextEndEvent,
	ContentTextEndPayload,
	ContentTextStartEvent,
	ContentTextStartPayload,
	ControlRollbackEvent,
	ControlRollbackPayload,
	ExecutorMetadata,
	RunLifecycleCompletedEvent,
	RunLifecycleCompletedPayload,
	RunLifecycleCreatedEvent,
	RunLifecycleCreatedPayload,
	UserInfo,
)
from turbo_agent_core.schema.states import Message
from turbo_agent_core.utils.aggregators import EventTreeAggregator


class TestEventTreeAggregatorV2(unittest.TestCase):
	BASE_TS = 1700000000000

	def setUp(self) -> None:
		self.user = UserInfo(id="user-1")
		self.executor = ExecutorMetadata(id="agent-root", name="Root Agent")

	def test_root_trace_always_uses_event_trace_id(self):
		agg = EventTreeAggregator(root_trace_id=None)
		agg.on_event(
			RunLifecycleCreatedEvent(
				trace_id="trace-root",
				trace_path=["trace-child"],
				run_id="run-child-1",
				timestamp=self.BASE_TS + 10,
				executor_id="AGENT:root@proj",
				executor_metadata=self.executor,
				user_metadata=self.user,
				payload=RunLifecycleCreatedPayload(input_data={"hello": "world"}),
			)
		)

		self.assertEqual(agg.root_trace_id, "trace-root")
		self.assertIn("trace-root", agg.nodes)
		self.assertIn("trace-child", agg.nodes)
		self.assertEqual(agg.nodes["trace-child"].parent_trace_id, "trace-root")
		self.assertEqual(agg.nodes["trace-root"].children, ["trace-child"])

	def test_tree_level_rollback_removes_child_nodes(self):
		agg = EventTreeAggregator(root_trace_id=None)
		child_created = RunLifecycleCreatedEvent(
			trace_id="trace-root",
			trace_path=["trace-child"],
			run_id="run-child-1",
			timestamp=self.BASE_TS + 50,
			executor_id="TOOL:search@proj",
			executor_metadata=self.executor,
			user_metadata=self.user,
			payload=RunLifecycleCreatedPayload(input_data={"query": "hello"}),
		)
		agg.on_event(child_created)

		agg.on_event(
			ControlRollbackEvent(
				trace_id="trace-root",
				run_id="run-root-rollback",
				timestamp=self.BASE_TS + 100,
				executor_id="AGENT:root@proj",
				payload=ControlRollbackPayload(target_timestamp=self.BASE_TS + 20, reason="回滚到子节点创建前"),
			)
		)

		self.assertEqual(list(agg.nodes.keys()), ["trace-root"])
		self.assertEqual(agg.nodes["trace-root"].children, [])
		self.assertEqual(agg.nodes["trace-root"].control["rollback"]["reason"], "回滚到子节点创建前")

	def test_snapshot_and_persistent_state_have_clear_boundaries(self):
		agg = EventTreeAggregator(root_trace_id="trace-root")
		events = [
			RunLifecycleCreatedEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 10,
				executor_id="AGENT:writer@proj",
				executor_metadata=self.executor,
				user_metadata=self.user,
				payload=RunLifecycleCreatedPayload(input_data={"conversation_id": "conv-1"}),
			),
			ContentTextStartEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 20,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				payload=ContentTextStartPayload(format="plain"),
			),
			ContentTextDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 30,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				payload=ContentTextDeltaPayload(delta="你好，世界"),
			),
			ContentTextEndEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 40,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				payload=ContentTextEndPayload(full_text="你好，世界"),
			),
			ContentReasoningStartEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 50,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				payload=ContentReasoningStartPayload(format="plain"),
			),
			ContentReasoningDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 60,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				payload=ContentReasoningDeltaPayload(delta="先分析问题"),
			),
			ContentReasoningEndEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 70,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				payload=ContentReasoningEndPayload(full_text="先分析问题"),
			),
			ContentActionStartEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 80,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionStartPayload(name="TOOL:search@proj/tool@1.0.0", call_type="TOOL", intent="检索"),
			),
			ContentActionDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 90,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionDeltaPayload(part="args", delta="he", key_path=["query"]),
			),
			ContentActionDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 100,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionDeltaPayload(part="args", delta="llo", key_path=["query"]),
			),
			ContentActionEndEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 110,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionEndPayload(arguments={"query": "hello"}, intent="检索 hello"),
			),
			ContentActionResultStartEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 120,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionResultStartPayload(record_id="record-1", status="success", mode="json"),
			),
			ContentActionResultDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 130,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionResultDeltaPayload(record_id="record-1", part="title", delta="检索结果"),
			),
			ContentActionResultDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 140,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionResultDeltaPayload(record_id="record-1", part="summary", delta="命中 1 条"),
			),
			ContentActionResultDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 150,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionResultDeltaPayload(record_id="record-1", part="reasoning", delta="命中最佳答案"),
			),
			ContentActionResultDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 160,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionResultDeltaPayload(record_id="record-1", part="output", delta="片段", key_path=["body"]),
			),
			ContentActionResultEndEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 170,
				executor_id="AGENT:writer@proj",
				react_id="react-1",
				action_id="action-1",
				payload=ContentActionResultEndPayload(record_id="record-1", full_result={"body": "最终结果"}, status="success", mode="json"),
			),
			ContentAgentResultStartEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 180,
				executor_id="AGENT:writer@proj",
				payload=ContentAgentResultStartPayload(mode="json"),
			),
			ContentAgentResultDeltaEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 190,
				executor_id="AGENT:writer@proj",
				payload=ContentAgentResultDeltaPayload(part="summary", delta="最终摘要"),
			),
			ContentAgentResultEndEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 200,
				executor_id="AGENT:writer@proj",
				payload=ContentAgentResultEndPayload(full_result={"answer": "最终答案"}, full_reasoning="最终推理", status="success", mode="json"),
			),
			RunLifecycleCompletedEvent(
				trace_id="trace-root",
				run_id="run-root-1",
				timestamp=self.BASE_TS + 210,
				executor_id="AGENT:writer@proj",
				payload=RunLifecycleCompletedPayload(output={"audit": "审计结果"}, usage={"tokenCost": 10}),
			),
		]

		for event in events:
			agg.on_event(event)

		snapshot = agg.snapshot_root()
		self.assertEqual(snapshot.lifecycle_output, {"audit": "审计结果"})
		self.assertEqual(snapshot.output, {"answer": "最终答案"})
		self.assertEqual(snapshot.reasoning_by_react["react-1"], "先分析问题")
		self.assertEqual(snapshot.actions["action-1"].result_title, "检索结果")
		self.assertEqual(snapshot.actions["action-1"].result_summary, "命中 1 条")
		self.assertEqual(snapshot.actions["action-1"].result_reasoning, "命中最佳答案")
		self.assertEqual(snapshot.actions["action-1"].raw_arguments, {"query": "hello"})

		persistent_full = agg.nodes["trace-root"].to_state(include_nested=True)
		self.assertIsInstance(persistent_full, Message)
		self.assertEqual(persistent_full.content, "你好，世界")
		self.assertEqual(persistent_full.reasoning_content, "最终推理")
		self.assertEqual(persistent_full.reactRounds[0].actions[0].title, "检索结果")
		self.assertEqual(persistent_full.reactRounds[0].actions[0].summary, "命中 1 条")
		self.assertEqual(persistent_full.reactRounds[0].actions[0].thought, "命中最佳答案")
		self.assertEqual(persistent_full.reactRounds[0].actions[0].records[0].related_message_action_ids, ["action-1"])

		persistent_flat = agg.nodes["trace-root"].to_state(include_nested=False)
		self.assertEqual(persistent_flat.reactRounds, [])
		self.assertEqual(persistent_flat.actions, [])


if __name__ == "__main__":
	unittest.main()