"""
Memory Trace Store - 内存存储实现

用于测试和简单场景，数据不持久化
"""

from typing import Dict, List, Optional

from agent.trace.models import Trace, Step


class MemoryTraceStore:
    """内存 Trace 存储"""

    def __init__(self):
        self._traces: Dict[str, Trace] = {}
        self._steps: Dict[str, Step] = {}
        self._trace_steps: Dict[str, List[str]] = {}  # trace_id -> [step_ids]

    async def create_trace(self, trace: Trace) -> str:
        self._traces[trace.trace_id] = trace
        self._trace_steps[trace.trace_id] = []
        return trace.trace_id

    async def get_trace(self, trace_id: str) -> Optional[Trace]:
        return self._traces.get(trace_id)

    async def update_trace(self, trace_id: str, **updates) -> None:
        trace = self._traces.get(trace_id)
        if trace:
            for key, value in updates.items():
                if hasattr(trace, key):
                    setattr(trace, key, value)

    async def list_traces(
        self,
        mode: Optional[str] = None,
        agent_type: Optional[str] = None,
        uid: Optional[str] = None,
        status: Optional[str] = None,
        limit: int = 50
    ) -> List[Trace]:
        traces = list(self._traces.values())

        # 过滤
        if mode:
            traces = [t for t in traces if t.mode == mode]
        if agent_type:
            traces = [t for t in traces if t.agent_type == agent_type]
        if uid:
            traces = [t for t in traces if t.uid == uid]
        if status:
            traces = [t for t in traces if t.status == status]

        # 排序（最新的在前）
        traces.sort(key=lambda t: t.created_at, reverse=True)

        return traces[:limit]

    async def add_step(self, step: Step) -> str:
        self._steps[step.step_id] = step

        # 添加到 trace 的 steps 列表
        if step.trace_id in self._trace_steps:
            self._trace_steps[step.trace_id].append(step.step_id)

        # 更新 trace 的 total_steps
        trace = self._traces.get(step.trace_id)
        if trace:
            trace.total_steps += 1

        return step.step_id

    async def get_step(self, step_id: str) -> Optional[Step]:
        return self._steps.get(step_id)

    async def get_trace_steps(self, trace_id: str) -> List[Step]:
        step_ids = self._trace_steps.get(trace_id, [])
        steps = [self._steps[sid] for sid in step_ids if sid in self._steps]
        steps.sort(key=lambda s: s.sequence)
        return steps

    async def get_step_children(self, step_id: str) -> List[Step]:
        children = []
        for step in self._steps.values():
            if step.parent_id == step_id:
                children.append(step)
        children.sort(key=lambda s: s.sequence)
        return children