"""Agent 通信协议定义 - 统一消息格式"""

from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any

from agentkit.core.exceptions import TaskCancelledError


class TaskStatus(str, Enum):
    """任务状态枚举"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    PARTIALLY_COMPLETED = "partially_completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    HANDOFF = "handoff"


class AgentStatus(str, Enum):
    """Agent 状态枚举"""
    ONLINE = "online"
    OFFLINE = "offline"
    BUSY = "busy"


@dataclass
class AgentCapability:
    """Agent 能力声明"""
    agent_name: str
    agent_type: str
    version: str
    supported_tasks: list[str]
    max_concurrency: int
    description: str
    input_schema: dict[str, Any] | None = None
    output_schema: dict[str, Any] | None = None

    def to_dict(self) -> dict:
        d = {
            "agent_name": self.agent_name,
            "agent_type": self.agent_type,
            "version": self.version,
            "supported_tasks": self.supported_tasks,
            "max_concurrency": self.max_concurrency,
            "description": self.description,
        }
        if self.input_schema is not None:
            d["input_schema"] = self.input_schema
        if self.output_schema is not None:
            d["output_schema"] = self.output_schema
        return d

    @classmethod
    def from_dict(cls, data: dict) -> "AgentCapability":
        return cls(
            agent_name=data["agent_name"],
            agent_type=data["agent_type"],
            version=data["version"],
            supported_tasks=data["supported_tasks"],
            max_concurrency=data["max_concurrency"],
            description=data["description"],
            input_schema=data.get("input_schema"),
            output_schema=data.get("output_schema"),
        )


@dataclass
class TaskMessage:
    """任务消息 - 从调度器发往 Agent"""
    task_id: str
    agent_name: str
    task_type: str
    priority: int
    input_data: dict
    callback_url: str | None
    created_at: datetime
    timeout_seconds: int = 300
    conversation_id: str | None = None

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "agent_name": self.agent_name,
            "task_type": self.task_type,
            "priority": self.priority,
            "input_data": self.input_data,
            "callback_url": self.callback_url,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "timeout_seconds": self.timeout_seconds,
            "conversation_id": self.conversation_id,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "TaskMessage":
        created_at = data.get("created_at")
        if isinstance(created_at, str):
            created_at = datetime.fromisoformat(created_at)
        return cls(
            task_id=data["task_id"],
            agent_name=data["agent_name"],
            task_type=data["task_type"],
            priority=data.get("priority", 0),
            input_data=data.get("input_data", {}),
            callback_url=data.get("callback_url"),
            created_at=created_at or datetime.now(timezone.utc),
            timeout_seconds=data.get("timeout_seconds", 300),
            conversation_id=data.get("conversation_id"),
        )


@dataclass
class TaskResult:
    """任务结果 - 从 Agent 返回"""
    task_id: str
    agent_name: str
    status: str
    output_data: dict | None
    error_message: str | None
    started_at: datetime
    completed_at: datetime
    metrics: dict | None = None
    trace: Any | None = None

    def to_dict(self) -> dict:
        d = {
            "task_id": self.task_id,
            "agent_name": self.agent_name,
            "status": self.status,
            "output_data": self.output_data,
            "error_message": self.error_message,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "metrics": self.metrics,
        }
        if self.trace is not None:
            d["trace"] = self.trace.to_dict() if hasattr(self.trace, "to_dict") else self.trace
        return d

    @classmethod
    def from_dict(cls, data: dict) -> "TaskResult":
        started_at = data.get("started_at")
        if isinstance(started_at, str):
            started_at = datetime.fromisoformat(started_at)
        completed_at = data.get("completed_at")
        if isinstance(completed_at, str):
            completed_at = datetime.fromisoformat(completed_at)
        return cls(
            task_id=data["task_id"],
            agent_name=data["agent_name"],
            status=data["status"],
            output_data=data.get("output_data"),
            error_message=data.get("error_message"),
            started_at=started_at or datetime.now(timezone.utc),
            completed_at=completed_at or datetime.now(timezone.utc),
            metrics=data.get("metrics"),
            trace=data.get("trace"),
        )


@dataclass
class TaskProgress:
    """进度上报 - Agent 执行过程中上报"""
    task_id: str
    agent_name: str
    progress: float
    message: str
    updated_at: datetime

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "agent_name": self.agent_name,
            "progress": self.progress,
            "message": self.message,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
        }

    @classmethod
    def from_dict(cls, data: dict) -> "TaskProgress":
        updated_at = data.get("updated_at")
        if isinstance(updated_at, str):
            updated_at = datetime.fromisoformat(updated_at)
        return cls(
            task_id=data["task_id"],
            agent_name=data["agent_name"],
            progress=data.get("progress", 0.0),
            message=data.get("message", ""),
            updated_at=updated_at or datetime.now(timezone.utc),
        )


@dataclass
class HandoffMessage:
    """任务转交消息 - Agent 间 Handoff"""
    source_agent: str
    target_agent: str
    task_id: str
    task_type: str
    context: dict[str, Any]
    reason: str
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def to_dict(self) -> dict:
        return {
            "source_agent": self.source_agent,
            "target_agent": self.target_agent,
            "task_id": self.task_id,
            "task_type": self.task_type,
            "context": self.context,
            "reason": self.reason,
            "created_at": self.created_at.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: dict) -> "HandoffMessage":
        created_at = data.get("created_at")
        if isinstance(created_at, str):
            created_at = datetime.fromisoformat(created_at)
        return cls(
            source_agent=data["source_agent"],
            target_agent=data["target_agent"],
            task_id=data["task_id"],
            task_type=data["task_type"],
            context=data.get("context", {}),
            reason=data["reason"],
            created_at=created_at or datetime.now(timezone.utc),
        )


@dataclass
class EvolutionEvent:
    """进化事件 - 记录 Agent 的自我进化变更"""
    agent_name: str
    change_type: str  # prompt / strategy / pipeline
    before: dict[str, Any]
    after: dict[str, Any]
    metrics: dict[str, Any] | None = None
    event_id: str | None = None
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def to_dict(self) -> dict:
        return {
            "agent_name": self.agent_name,
            "change_type": self.change_type,
            "before": self.before,
            "after": self.after,
            "metrics": self.metrics,
            "event_id": self.event_id,
            "created_at": self.created_at.isoformat(),
        }


@dataclass
class CancellationToken:
    """协作式取消令牌，用于通知 ReAct 循环和 Agent 停止执行。

    由 BaseAgent 创建并存储在 _active_tokens 中，
    当外部调用 cancel_task() 时设置 cancelled 标志，
    ReAct 循环在每次迭代开始时检查该标志。
    """

    _cancelled: bool = field(default=False, repr=False)

    def cancel(self) -> None:
        """标记此令牌为已取消"""
        self._cancelled = True

    @property
    def is_cancelled(self) -> bool:
        """返回是否已取消"""
        return self._cancelled

    def check(self) -> None:
        """检查是否已取消，若已取消则抛出 TaskCancelledError"""
        if self._cancelled:
            raise TaskCancelledError(task_id="")
