"""A2A Protocol Types - Agent2Agent v1.0 data structures."""

from __future__ import annotations
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, List, Optional, Union

class TaskState(str, Enum):
    SUBMITTED = "TASK_STATE_SUBMITTED"
    WORKING = "TASK_STATE_WORKING"
    INPUT_REQUIRED = "TASK_STATE_INPUT_REQUIRED"
    COMPLETED = "TASK_STATE_COMPLETED"
    FAILED = "TASK_STATE_FAILED"
    CANCELED = "TASK_STATE_CANCELED"
    REJECTED = "TASK_STATE_REJECTED"
    @property
    def terminal(self) -> bool:
        return self in {TaskState.COMPLETED, TaskState.FAILED, TaskState.CANCELED, TaskState.REJECTED}

@dataclass
class TextPart:
    text: str
    metadata: Optional[Dict[str, Any]] = None
    def to_dict(self) -> dict:
        d = {"text": self.text}
        if self.metadata: d["metadata"] = self.metadata
        return d
    @classmethod
    def from_dict(cls, d): return cls(text=d["text"], metadata=d.get("metadata"))

@dataclass
class FilePart:
    media_type: str = "application/octet-stream"
    raw: Optional[str] = None
    url: Optional[str] = None
    filename: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    def to_dict(self) -> dict:
        d = {"mediaType": self.media_type}
        if self.raw is not None: d["raw"] = self.raw
        if self.url is not None: d["url"] = self.url
        if self.filename: d["filename"] = self.filename
        if self.metadata: d["metadata"] = self.metadata
        return d
    @classmethod
    def from_dict(cls, d):
        return cls(media_type=d.get("mediaType","application/octet-stream"), raw=d.get("raw"), url=d.get("url"), filename=d.get("filename"), metadata=d.get("metadata"))

@dataclass
class DataPart:
    data: Dict[str, Any]
    media_type: str = "application/json"
    metadata: Optional[Dict[str, Any]] = None
    def to_dict(self) -> dict:
        d = {"data": self.data, "mediaType": self.media_type}
        if self.metadata: d["metadata"] = self.metadata
        return d
    @classmethod
    def from_dict(cls, d): return cls(data=d["data"], media_type=d.get("mediaType","application/json"), metadata=d.get("metadata"))

Part = Union[TextPart, FilePart, DataPart]

def part_from_dict(d):
    if "text" in d: return TextPart.from_dict(d)
    if "raw" in d or "url" in d: return FilePart.from_dict(d)
    if "data" in d: return DataPart.from_dict(d)
    raise ValueError(f"Cannot discriminate Part type from keys: {list(d.keys())}")

@dataclass
class Message:
    role: str
    parts: List[Part]
    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    context_id: Optional[str] = None
    task_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    def to_dict(self):
        d = {"role": self.role, "messageId": self.message_id, "parts": [p.to_dict() for p in self.parts]}
        if self.context_id: d["contextId"] = self.context_id
        if self.task_id: d["taskId"] = self.task_id
        if self.metadata: d["metadata"] = self.metadata
        return d
    @classmethod
    def from_dict(cls, d):
        return cls(role=d["role"], parts=[part_from_dict(p) for p in d["parts"]], message_id=d.get("messageId",str(uuid.uuid4())), context_id=d.get("contextId"), task_id=d.get("taskId"), metadata=d.get("metadata"))

@dataclass
class Artifact:
    artifact_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    parts: List[Part] = field(default_factory=list)
    name: Optional[str] = None
    description: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    def to_dict(self):
        d = {"artifactId": self.artifact_id, "parts": [p.to_dict() for p in self.parts]}
        if self.name: d["name"] = self.name
        if self.description: d["description"] = self.description
        if self.metadata: d["metadata"] = self.metadata
        return d
    @classmethod
    def from_dict(cls, d):
        return cls(artifact_id=d.get("artifactId",str(uuid.uuid4())), parts=[part_from_dict(p) for p in d.get("parts",[])], name=d.get("name"), description=d.get("description"), metadata=d.get("metadata"))

@dataclass
class TaskStatus:
    state: TaskState
    message: Optional[Message] = None
    timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
    def to_dict(self):
        d = {"state": self.state.value, "timestamp": self.timestamp}
        if self.message: d["message"] = self.message.to_dict()
        return d
    @classmethod
    def from_dict(cls, d):
        msg = d.get("message")
        return cls(state=TaskState(d["state"]), message=Message.from_dict(msg) if msg else None, timestamp=d.get("timestamp",datetime.now(timezone.utc).isoformat()))

@dataclass
class Task:
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    context_id: Optional[str] = None
    status: TaskStatus = field(default_factory=lambda: TaskStatus(state=TaskState.SUBMITTED))
    artifacts: List[Artifact] = field(default_factory=list)
    history: List[Message] = field(default_factory=list)
    metadata: Optional[Dict[str, Any]] = None
    def to_dict(self):
        d = {"id": self.id, "status": self.status.to_dict()}
        if self.context_id: d["contextId"] = self.context_id
        if self.artifacts: d["artifacts"] = [a.to_dict() for a in self.artifacts]
        if self.history: d["history"] = [m.to_dict() for m in self.history]
        if self.metadata: d["metadata"] = self.metadata
        return d
    @classmethod
    def from_dict(cls, d):
        return cls(id=d.get("id",str(uuid.uuid4())), context_id=d.get("contextId"), status=TaskStatus.from_dict(d["status"]) if "status" in d else TaskStatus(TaskState.SUBMITTED), artifacts=[Artifact.from_dict(a) for a in d.get("artifacts",[])], history=[Message.from_dict(m) for m in d.get("history",[])], metadata=d.get("metadata"))

@dataclass
class AgentSkill:
    id: str
    name: str
    description: str = ""
    tags: List[str] = field(default_factory=list)
    examples: List[str] = field(default_factory=list)
    input_modes: List[str] = field(default_factory=lambda: ["text"])
    output_modes: List[str] = field(default_factory=lambda: ["text"])
    def to_dict(self):
        d = {"id": self.id, "name": self.name, "description": self.description}
        if self.tags: d["tags"] = self.tags
        if self.examples: d["examples"] = self.examples
        if self.input_modes != ["text"]: d["inputModes"] = self.input_modes
        if self.output_modes != ["text"]: d["outputModes"] = self.output_modes
        return d
    @classmethod
    def from_dict(cls, d):
        return cls(id=d["id"], name=d.get("name",d["id"]), description=d.get("description",""), tags=d.get("tags",[]), examples=d.get("examples",[]), input_modes=d.get("inputModes",["text"]), output_modes=d.get("outputModes",["text"]))

@dataclass
class AgentCard:
    name: str
    description: str = ""
    version: str = "1.0.0"
    url: str = ""
    skills: List[AgentSkill] = field(default_factory=list)
    capabilities: Dict[str, bool] = field(default_factory=dict)
    provider: Optional[Dict[str, str]] = None
    def to_dict(self):
        d = {"name": self.name, "description": self.description, "version": self.version, "url": self.url, "skills": [s.to_dict() for s in self.skills]}
        if self.capabilities: d["capabilities"] = self.capabilities
        if self.provider: d["provider"] = self.provider
        return d
    @classmethod
    def from_dict(cls, d):
        return cls(name=d["name"], description=d.get("description",""), version=d.get("version","1.0.0"), url=d.get("url",""), skills=[AgentSkill.from_dict(s) for s in d.get("skills",[])], capabilities=d.get("capabilities",{}), provider=d.get("provider"))

@dataclass
class JSONRPCError:
    code: int
    message: str
    data: Optional[Any] = None
    def to_dict(self):
        d = {"code": self.code, "message": self.message}
        if self.data is not None: d["data"] = self.data
        return d

@dataclass
class JSONRPCRequest:
    method: str
    params: Optional[Dict[str, Any]] = None
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    jsonrpc: str = "2.0"
    def to_dict(self):
        d = {"jsonrpc": self.jsonrpc, "method": self.method, "id": self.id}
        if self.params is not None: d["params"] = self.params
        return d

@dataclass
class JSONRPCResponse:
    id: str
    result: Optional[Any] = None
    error: Optional[JSONRPCError] = None
    jsonrpc: str = "2.0"
    def to_dict(self):
        d = {"jsonrpc": self.jsonrpc, "id": self.id}
        if self.error: d["error"] = self.error.to_dict()
        else: d["result"] = self.result
        return d
    @classmethod
    def from_dict(cls, d):
        err = d.get("error")
        return cls(id=d["id"], result=d.get("result"), error=JSONRPCError(err["code"],err["message"],err.get("data")) if err else None)

class A2AError:
    @staticmethod
    def parse_error(data=None): return JSONRPCError(-32700, "Parse error", data)
    @staticmethod
    def invalid_request(data=None): return JSONRPCError(-32600, "Invalid Request", data)
    @staticmethod
    def method_not_found(data=None): return JSONRPCError(-32601, "Method not found", data)
    @staticmethod
    def invalid_params(data=None): return JSONRPCError(-32602, "Invalid params", data)
    @staticmethod
    def internal_error(data=None): return JSONRPCError(-32603, "Internal error", data)
    @staticmethod
    def task_not_found(task_id): return JSONRPCError(-32001, f"Task not found: {task_id}")
    @staticmethod
    def task_not_cancelable(task_id): return JSONRPCError(-32002, f"Task not cancelable: {task_id}")
    @staticmethod
    def agent_not_found(name): return JSONRPCError(-32009, f"Agent not found: {name}")