from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional

from turbo_agent_core.schema.events import BaseEvent
from turbo_agent_core.utils.aggregators import EventTreeAggregator


class BaseExecutionStore(ABC):
    """执行记录存储抽象接口（面向事件流）。

    设计目标：让 Store Worker/消费者可以根据 Core Event 协议，直接落库。

    说明：当前仓库尚未对外发布，不考虑兼容性，可直接按设计调整接口。
    """

    async def connect(self) -> None:
        """建立连接（可选）。"""

    async def disconnect(self) -> None:
        """断开连接（可选）。"""

    @abstractmethod
    async def save_trace_snapshot(self, aggregator: EventTreeAggregator, created_resources: Optional[Dict[str, set]] = None) -> None:
        """保存 Trace 聚合快照。

        由 DataStreamConsumer 定期调用（每 2s 或结束时）。
        Store 实现需负责将 aggregator 中的状态映射到业务表（Conversation/Message/Action）。
        """

    # --- Conversation ---

    @abstractmethod
    async def upsert_conversation_from_run_created(self, event: BaseEvent) -> str:
        """处理 run.lifecycle.created：创建/更新 Conversation，并返回 conversation_id。

        约定：默认使用 event.trace_id 作为 conversation_id，保证跨进程可寻址。
        """

    @abstractmethod
    async def update_conversation_from_run_lifecycle(self, event: BaseEvent) -> None:
        """处理 run.lifecycle.*：更新 Conversation 状态、错误、输出等。"""

    # --- Message（文本流） ---

    @abstractmethod
    async def upsert_text_message_start(
        self,
        *,
        trace_id: str,
        run_id: str,
        executor_id: Optional[str],
        role: str,
        mode: str,
        json_schema: Optional[Dict[str, Any]],
    ) -> str:
        """处理 content.text.start：创建/更新 Message 并返回 message_id。

        重要说明：
        - content.text.start 只表示"文本内容块开始"，不承担 run.lifecycle 的开始标记语义。
        - 会话/任务/工具调用的绑定关系，应以 run.lifecycle(status="created"/"resuming") 携带的 input_data 为准。
        """

    @abstractmethod
    async def append_text_message_delta(
        self,
        *,
        message_id: str,
        delta: str,
        key_path: Optional[list[Any]] = None,
    ) -> None:
        """处理 content.text.delta：增量追加消息内容。"""

    @abstractmethod
    async def finish_text_message(self, *, message_id: str, full_text: str) -> None:
        """处理 content.text.end：完成消息（status=finished）。"""

    # --- Action（工具/动作） ---

    @abstractmethod
    async def upsert_action_start(
        self,
        *,
        trace_id: str,
        run_id: str,
        message_id: str,
        action_id: str,
        index: int,
        name: str,
        is_client: bool,
        intent: Optional[str],
    ) -> str:
        """处理 content.action.start：创建/更新 Action 记录并返回 action_record_id。"""

    @abstractmethod
    async def append_action_delta(
        self,
        *,
        action_record_id: str,
        part: str,
        delta: str,
        key_path: Optional[list[Any]] = None,
    ) -> None:
        """处理 content.action.delta：增量更新 intent/args/thought 等字段。"""

    @abstractmethod
    async def finish_action(self, *, action_record_id: str, full_arguments: Dict[str, Any]) -> None:
        """处理 content.action.end：写入最终 arguments，并更新状态。"""

    # --- Action Result ---

    @abstractmethod
    async def start_action_result(self, *, action_record_id: str, status: str) -> None:
        """处理 content.action.result.start：标记动作结果块开始（动作执行态）。

        注意：这里的"执行态"仅指 Action/Tool 的局部状态，不等价于 run.lifecycle.running。
        """

    @abstractmethod
    async def append_action_result_delta(
        self,
        *,
        action_record_id: str,
        part: str,
        delta: str,
        key_path: Optional[list[Any]] = None,
    ) -> None:
        """处理 content.action.result.delta：增量更新 title/summary/output/error 等字段。"""

    @abstractmethod
    async def finish_action_result(
        self,
        *,
        action_record_id: str,
        status: str,
        full_result: Any,
        images: Optional[list[str]] = None,
        files: Optional[list[str]] = None,
    ) -> None:
        """处理 content.action.result.end：写入最终结果并更新状态（Succeed/Failed）。"""
