"""Task, Message, Artifact, and Part data models — A2A compatible."""

from __future__ import annotations

import asyncio
import enum
from collections.abc import Callable, Coroutine
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any


class TaskStatus(enum.Enum):
    """A2A task lifecycle states."""

    SUBMITTED = "submitted"
    WORKING = "working"
    INPUT_REQUIRED = "input_required"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELED = "canceled"
    REJECTED = "rejected"

    @property
    def is_terminal(self) -> bool:
        return self in (
            TaskStatus.COMPLETED,
            TaskStatus.FAILED,
            TaskStatus.CANCELED,
            TaskStatus.REJECTED,
        )


@dataclass
class Part:
    """The smallest content unit in a message or artifact."""

    text: str | None = None
    data: dict[str, Any] | None = None
    url: str | None = None
    raw: bytes | None = None
    media_type: str | None = None
    metadata: dict[str, str] | None = None

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {}
        if self.text is not None:
            d["text"] = self.text
        if self.data is not None:
            d["data"] = self.data
        if self.url is not None:
            d["url"] = self.url
        if self.raw is not None:
            d["raw"] = self.raw.hex()
        if self.media_type:
            d["media_type"] = self.media_type
        if self.metadata:
            d["metadata"] = self.metadata
        return d

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Part:
        raw = data.get("raw")
        return cls(
            text=data.get("text"),
            data=data.get("data"),
            url=data.get("url"),
            raw=bytes.fromhex(raw) if isinstance(raw, str) else raw,
            media_type=data.get("media_type"),
            metadata=data.get("metadata"),
        )


@dataclass
class Message:
    """A communication turn between agents."""

    role: str  # "user" or "agent"
    parts: list[Part] = field(default_factory=list)
    message_id: str = ""
    created_at: datetime | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "role": self.role,
            "parts": [p.to_dict() for p in self.parts],
            "message_id": self.message_id,
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Message:
        parts = [Part.from_dict(p) for p in data.get("parts", [])]
        return cls(
            role=data["role"],
            parts=parts,
            message_id=data.get("message_id", ""),
        )


@dataclass
class Artifact:
    """A concrete output generated by an agent."""

    artifact_id: str = ""
    name: str = ""
    parts: list[Part] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "artifact_id": self.artifact_id,
            "name": self.name,
            "parts": [p.to_dict() for p in self.parts],
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> Artifact:
        parts = [Part.from_dict(p) for p in data.get("parts", [])]
        return cls(
            artifact_id=data.get("artifact_id", ""),
            name=data.get("name", ""),
            parts=parts,
        )


@dataclass
class Task:
    """Represents a complete A2A task with all state."""

    task_id: str
    context_id: str = ""
    status: TaskStatus = TaskStatus.SUBMITTED
    messages: list[Message] = field(default_factory=list)
    artifacts: list[Artifact] = field(default_factory=list)
    target_skill_id: str = ""
    originator_peer_id: str = ""
    created_at: datetime | None = None
    updated_at: datetime | None = None


class TaskHandle:
    """Handle for tracking a remote task's progress.

    Returned by Node.send_task(). Allows waiting for completion,
    checking status, and canceling.
    """

    def __init__(self, task: Task, cancel_fn: Callable[[], Coroutine[Any, Any, None]]) -> None:
        self._task = task
        self._cancel_fn = cancel_fn
        self._done_event = asyncio.Event()
        if task.status.is_terminal:
            self._done_event.set()

    @property
    def task_id(self) -> str:
        return self._task.task_id

    @property
    def status(self) -> TaskStatus:
        return self._task.status

    @property
    def artifacts(self) -> list[Artifact]:
        return self._task.artifacts

    def _update(self, status: TaskStatus, artifacts: list[Artifact] | None = None) -> None:
        """Called internally when a status update is received from the daemon."""
        self._task.status = status
        if artifacts:
            self._task.artifacts = artifacts
        if status.is_terminal:
            self._done_event.set()

    async def wait(self, timeout: float | None = None) -> Task:
        """Wait for the task to reach a terminal state.

        Args:
            timeout: Maximum seconds to wait. None means wait forever.

        Returns:
            The completed Task.

        Raises:
            TaskTimeoutError: If timeout is exceeded.
            TaskFailedError: If the task failed.
            TaskCanceledError: If the task was canceled.
            TaskRejectedError: If the task was rejected.
        """
        from agentanycast.exceptions import (
            TaskCanceledError,
            TaskFailedError,
            TaskRejectedError,
            TaskTimeoutError,
        )

        try:
            await asyncio.wait_for(self._done_event.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            raise TaskTimeoutError(f"Task {self.task_id} timed out after {timeout}s")

        if self._task.status == TaskStatus.FAILED:
            raise TaskFailedError(f"Task {self.task_id} failed")
        if self._task.status == TaskStatus.CANCELED:
            raise TaskCanceledError(f"Task {self.task_id} was canceled")
        if self._task.status == TaskStatus.REJECTED:
            raise TaskRejectedError(f"Task {self.task_id} was rejected")

        return self._task

    async def cancel(self) -> None:
        """Request cancellation of this task."""
        await self._cancel_fn()


class IncomingTask:
    """A task received from a remote agent for local processing.

    Provides methods to update status, complete, or fail the task.
    """

    def __init__(
        self,
        task: Task,
        sender_card: Any | None,
        update_fn: Callable[
            [str, TaskStatus, list[Artifact] | None, str | None],
            Coroutine[Any, Any, None],
        ],
    ) -> None:
        self._task = task
        self.sender_card = sender_card
        self._update_fn = update_fn

    @property
    def task_id(self) -> str:
        return self._task.task_id

    @property
    def peer_id(self) -> str:
        return self._task.originator_peer_id

    @property
    def messages(self) -> list[Message]:
        return self._task.messages

    @property
    def target_skill_id(self) -> str:
        return self._task.target_skill_id

    async def update_status(self, status: str) -> None:
        """Update the task status (e.g., 'working')."""
        ts = TaskStatus(status)
        await self._update_fn(self.task_id, ts, None, None)

    async def complete(
        self,
        artifacts: list[dict[str, Any]] | list[Artifact] | None = None,
    ) -> None:
        """Mark the task as completed with optional artifacts."""
        parsed: list[Artifact] = []
        if artifacts:
            for a in artifacts:
                if isinstance(a, Artifact):
                    parsed.append(a)
                else:
                    parsed.append(Artifact.from_dict(a))
        await self._update_fn(self.task_id, TaskStatus.COMPLETED, parsed, None)

    async def fail(self, error: str) -> None:
        """Mark the task as failed."""
        await self._update_fn(self.task_id, TaskStatus.FAILED, None, error)

    async def request_input(self, message: dict[str, Any] | Message | None = None) -> None:
        """Request additional input from the client.

        Args:
            message: Optional message describing what input is needed.
        """
        msg_text: str | None = None
        if message is not None:
            if isinstance(message, Message):
                texts = [p.text for p in message.parts if p.text]
                msg_text = " ".join(texts) if texts else None
            elif isinstance(message, dict):
                parts = message.get("parts", [])
                texts = [p.get("text", "") for p in parts if p.get("text")]
                msg_text = " ".join(texts) if texts else None
        await self._update_fn(self.task_id, TaskStatus.INPUT_REQUIRED, None, msg_text)
