"""Small A2A-compatible message models used by the MQTT transport."""

from __future__ import annotations

from dataclasses import dataclass
import json
from typing import Any, ClassVar

from a2a_agent.invite import validate_pairing_code


JsonDict = dict[str, Any]
VALID_TASK_STATUSES = {"pending", "accepted", "working", "completed", "failed"}


@dataclass(frozen=True)
class AgentCard:
    name: str

    def to_dict(self) -> JsonDict:
        return {"name": self.name}

    @classmethod
    def from_dict(cls, data: JsonDict) -> "AgentCard":
        return cls(name=str(data["name"]))

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, sort_keys=True)

    @classmethod
    def from_json(cls, payload: str | bytes) -> "AgentCard":
        return cls.from_dict(json.loads(payload))


@dataclass(frozen=True)
class PairingRequest:
    code: str

    def to_dict(self) -> JsonDict:
        return {"code": validate_pairing_code(self.code)}

    @classmethod
    def from_dict(cls, data: JsonDict) -> "PairingRequest":
        return cls(code=validate_pairing_code(str(data["code"])))

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, sort_keys=True)

    @classmethod
    def from_json(cls, payload: str | bytes) -> "PairingRequest":
        return cls.from_dict(json.loads(payload))


@dataclass(frozen=True)
class PairingReply:
    ok: bool
    room: str | None = None
    device_name: str | None = None
    token: str | None = None
    broker: JsonDict | None = None
    share: JsonDict | None = None
    error: str | None = None

    def to_dict(self) -> JsonDict:
        data: JsonDict = {"ok": self.ok}
        if self.room is not None:
            data["room"] = self.room
        if self.device_name is not None:
            data["device_name"] = self.device_name
        if self.token is not None:
            data["token"] = self.token
        if self.broker is not None:
            data["broker"] = dict(self.broker)
        if self.share is not None:
            data["share"] = dict(self.share)
        if self.error is not None:
            data["error"] = self.error
        return data

    @classmethod
    def from_dict(cls, data: JsonDict) -> "PairingReply":
        broker = data.get("broker")
        if broker is not None and not isinstance(broker, dict):
            raise ValueError("pairing reply broker must be a mapping")
        share = data.get("share")
        if share is not None and not isinstance(share, dict):
            raise ValueError("pairing reply share must be a mapping")
        return cls(
            ok=bool(data["ok"]),
            room=str(data["room"]) if data.get("room") is not None else None,
            device_name=(
                str(data["device_name"])
                if data.get("device_name") is not None
                else None
            ),
            token=str(data["token"]) if data.get("token") is not None else None,
            broker=dict(broker) if broker is not None else None,
            share=dict(share) if share is not None else None,
            error=str(data["error"]) if data.get("error") is not None else None,
        )

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, sort_keys=True)

    @classmethod
    def from_json(cls, payload: str | bytes) -> "PairingReply":
        return cls.from_dict(json.loads(payload))


@dataclass(frozen=True)
class Message:
    role: str
    parts: list[JsonDict]

    def to_dict(self) -> JsonDict:
        return {
            "role": self.role,
            "parts": [dict(part) for part in self.parts],
        }

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

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, sort_keys=True)

    @classmethod
    def from_json(cls, payload: str | bytes) -> "Message":
        return cls.from_dict(json.loads(payload))


@dataclass(frozen=True)
class FilePart:
    name: str
    path: str
    size: int

    def __post_init__(self) -> None:
        if not self.name:
            raise ValueError("file part name must not be empty")
        if not self.path or self.path.startswith("/") or ".." in self.path.split("/"):
            raise ValueError("file part path must be a share-relative path")
        if self.size < 0:
            raise ValueError("file part size must be non-negative")

    def to_dict(self) -> JsonDict:
        return {
            "type": "file",
            "name": self.name,
            "path": self.path,
            "size": self.size,
        }

    @classmethod
    def from_dict(cls, data: JsonDict) -> "FilePart":
        if data.get("type") != "file":
            raise ValueError("file part type must be 'file'")
        name = str(data.get("name", ""))
        path = str(data.get("path", ""))
        size = int(data.get("size", -1))
        return cls(name=name, path=path, size=size)


def build_file_part(name: str, path: str, size: int) -> JsonDict:
    return FilePart(name=name, path=path, size=size).to_dict()


def parse_file_part(part: JsonDict) -> FilePart:
    return FilePart.from_dict(part)


@dataclass(frozen=True)
class Task:
    id: str
    from_: str
    to: str
    status: str
    messages: list[Message]
    share: JsonDict | None = None
    control: str | None = None

    valid_statuses: ClassVar[set[str]] = VALID_TASK_STATUSES

    def __post_init__(self) -> None:
        if self.status not in self.valid_statuses:
            allowed = ", ".join(sorted(self.valid_statuses))
            raise ValueError(f"status must be one of: {allowed}")

    def to_dict(self) -> JsonDict:
        data: JsonDict = {
            "id": self.id,
            "from": self.from_,
            "to": self.to,
            "status": self.status,
            "messages": [message.to_dict() for message in self.messages],
        }
        if self.share is not None:
            data["share"] = dict(self.share)
        if self.control is not None:
            data["control"] = self.control
        return data

    @classmethod
    def from_dict(cls, data: JsonDict) -> "Task":
        share = data.get("share")
        if share is not None and not isinstance(share, dict):
            raise ValueError("task share must be a mapping")
        return cls(
            id=str(data["id"]),
            from_=str(data["from"]),
            to=str(data["to"]),
            status=str(data["status"]),
            messages=[
                Message.from_dict(message)
                for message in data.get("messages", [])
            ],
            share=dict(share) if share is not None else None,
            control=(
                str(data["control"])
                if data.get("control") is not None
                else None
            ),
        )

    def to_json(self) -> str:
        return json.dumps(self.to_dict(), ensure_ascii=False, sort_keys=True)

    @classmethod
    def from_json(cls, payload: str | bytes) -> "Task":
        return cls.from_dict(json.loads(payload))
