# -*- coding: utf-8 -*-
"""
子 Agent 通信协议（SubAgent Protocol）

借鉴 Evolver 的 bridge.js sessions_spawn(JSON) 模式，
定义标准化的任务描述格式和结果格式。

设计目标：
- 任务和结果都有标准 JSON schema → 可跨进程传输
- 状态码枚举 → 主 Agent 可精确判断结果
- 可序列化为 JSON → 未来改用子进程/HTTP 通信时零修改

使用方式：
    from abyss.subagent.protocol import SubAgentProtocol, TaskStatus

    # 编码任务
    task_json = SubAgentProtocol.encode_task(sub_id="1", task="analyze code")

    # 解码结果
    result = SubAgentProtocol.decode_result(result_json)
    if result["status"] == TaskStatus.SUCCESS:
        print(result["output"])
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field, asdict
from enum import Enum
from typing import Any, Dict, List, Optional


class TaskStatus(str, Enum):
    """子任务执行状态。"""
    SUCCESS = "success"
    FAILED = "failed"
    TIMEOUT = "timeout"
    CANCELLED = "cancelled"


@dataclass
class TaskDefinition:
    """子任务定义——标准化的输入格式。"""
    task: str                              # 必填：自包含的任务描述
    sub_id: str = ""                       # 可选：子任务 ID，未指定时自动生成
    context: str = ""                      # 可选：附加上下文
    methodology: Optional[str] = None      # 可选：方法论纪律名
    metadata: Dict[str, Any] = field(default_factory=dict)  # 扩展元数据

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        # 始终保留 task 和 sub_id，其他字段空值可省略
        keep = {"task", "sub_id"}
        return {k: v for k, v in d.items() if k in keep or v}

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TaskDefinition":
        return cls(
            task=data.get("task", ""),
            sub_id=str(data.get("sub_id", "")),
            context=str(data.get("context", "")),
            methodology=data.get("methodology"),
            metadata=data.get("metadata", {}) or {},
        )


@dataclass
class TaskResult:
    """子任务结果——标准化的输出格式。"""
    status: TaskStatus = TaskStatus.SUCCESS
    sub_id: str = ""
    output: str = ""                       # 子 Agent 的最终回复文本
    error: str = ""                        # 失败时的错误信息
    duration_ms: int = 0                   # 执行耗时（毫秒）
    tool_rounds: int = 0                   # 工具调用轮数
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def is_success(self) -> bool:
        return self.status == TaskStatus.SUCCESS

    @property
    def is_failed(self) -> bool:
        return self.status in (TaskStatus.FAILED, TaskStatus.TIMEOUT, TaskStatus.CANCELLED)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "status": self.status.value,
            "sub_id": self.sub_id,
            "output": self.output,
            "error": self.error,
            "duration_ms": self.duration_ms,
            "tool_rounds": self.tool_rounds,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TaskResult":
        raw_status = data.get("status", "success")
        try:
            status = TaskStatus(raw_status)
        except ValueError:
            status = TaskStatus.FAILED
        return cls(
            status=status,
            sub_id=str(data.get("sub_id", "")),
            output=str(data.get("output", "")),
            error=str(data.get("error", "")),
            duration_ms=int(data.get("duration_ms", 0) or 0),
            tool_rounds=int(data.get("tool_rounds", 0) or 0),
            metadata=data.get("metadata", {}) or {},
        )

    @classmethod
    def success(cls, sub_id: str, output: str, **kwargs) -> "TaskResult":
        return cls(status=TaskStatus.SUCCESS, sub_id=sub_id, output=output, **kwargs)

    @classmethod
    def failure(cls, sub_id: str, error: str, **kwargs) -> "TaskResult":
        return cls(status=TaskStatus.FAILED, sub_id=sub_id, error=error, **kwargs)

    @classmethod
    def timeout(cls, sub_id: str, max_rounds: int, **kwargs) -> "TaskResult":
        return cls(
            status=TaskStatus.TIMEOUT,
            sub_id=sub_id,
            error=f"达到 {max_rounds} 轮上限仍未完成",
            **kwargs,
        )


class SubAgentProtocol:
    """子 Agent 通信协议：编码/解码任务和结果。

    借鉴 Evolver 的 sessions_spawn(JSON) 桥接模式：
    - 统一 JSON 格式
    - 无需正则解析（纯 JSON）
    - 参数和结果均有 schema 约束
    """

    VERSION = "1.0"

    @staticmethod
    def encode_task(task: TaskDefinition) -> str:
        """将任务定义编码为 JSON 字符串。"""
        payload = {
            "protocol": "subagent",
            "version": SubAgentProtocol.VERSION,
            "task": task.to_dict(),
        }
        if not payload["task"]["sub_id"]:
            import uuid
            payload["task"]["sub_id"] = str(uuid.uuid4())[:8]
        return json.dumps(payload, ensure_ascii=False)

    @staticmethod
    def decode_task(json_str: str) -> TaskDefinition:
        """从 JSON 字符串解码任务定义。"""
        payload = json.loads(json_str)
        task_data = payload.get("task", payload)
        return TaskDefinition.from_dict(task_data)

    @staticmethod
    def encode_result(result: TaskResult) -> str:
        """将任务结果编码为 JSON 字符串。"""
        payload = {
            "protocol": "subagent",
            "version": SubAgentProtocol.VERSION,
            "result": result.to_dict(),
        }
        return json.dumps(payload, ensure_ascii=False)

    @staticmethod
    def decode_result(json_str: str) -> TaskResult:
        """从 JSON 字符串解码任务结果。"""
        payload = json.loads(json_str)
        result_data = payload.get("result", payload)
        return TaskResult.from_dict(result_data)

    @staticmethod
    def build_summary(results: List[TaskResult]) -> str:
        """从结果列表生成汇总摘要。"""
        total = len(results)
        ok = sum(1 for r in results if r.is_success)
        failed = total - ok
        if failed == 0:
            return f"共 {total} 个子任务全部成功。"
        details = []
        for r in results:
            if r.is_failed:
                details.append(f"  [{r.status.value}] {r.sub_id}: {r.error[:80]}")
        header = f"共 {total} 个子任务，成功 {ok} 个，失败 {failed} 个。"
        if details:
            header += "\n" + "\n".join(details[:5])
            if len(details) > 5:
                header += f"\n  ... 还有 {len(details) - 5} 个失败项"
        return header
