from __future__ import annotations

from dataclasses import dataclass
from typing import Any


MSG_DELTA_PUSH = "cv.delta.push"
MSG_DELTA_PULL = "cv.delta.pull"
MSG_DELTA_ACK = "cv.delta.ack"
MSG_SNAPSHOT = "cv.snapshot"

SUPPORTED_SCHEMA = "v1"
SUPPORTED_MSG_TYPES = frozenset(
    {
        MSG_DELTA_PUSH,
        MSG_DELTA_PULL,
        MSG_DELTA_ACK,
        MSG_SNAPSHOT,
    }
)


@dataclass(frozen=True)
class ClusterViewDelta:
    membership: dict[str, dict[str, Any]]
    tombstones: dict[str, int] | None = None
    leases_summary: dict[str, Any] | None = None

    def as_dict(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "membership": dict(self.membership),
        }
        if self.tombstones is not None:
            payload["tombstones"] = dict(self.tombstones)
        if self.leases_summary is not None:
            payload["leases_summary"] = dict(self.leases_summary)
        return payload


@dataclass(frozen=True)
class ClusterViewGossipEnvelope:
    msg_type: str
    schema_version: str
    from_node: str
    to_node: str
    view_epoch: int
    delta_seq: int
    sent_at_ms: int
    delta: ClusterViewDelta | None = None

    def as_dict(self) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "msg_type": self.msg_type,
            "schema_version": self.schema_version,
            "from_node": self.from_node,
            "to_node": self.to_node,
            "view_epoch": self.view_epoch,
            "delta_seq": self.delta_seq,
            "sent_at_ms": self.sent_at_ms,
        }
        if self.delta is not None:
            payload["delta"] = self.delta.as_dict()
        return payload


def parse_envelope(
    payload: ClusterViewGossipEnvelope | dict[str, Any],
    *,
    msg_type: str | None = None,
) -> ClusterViewGossipEnvelope:
    if isinstance(payload, ClusterViewGossipEnvelope):
        return payload
    if not isinstance(payload, dict):
        raise TypeError("Cluster gossip payload must be dict or ClusterViewGossipEnvelope.")

    envelope_type = payload.get("msg_type", msg_type)
    if not isinstance(envelope_type, str):
        raise ValueError("cluster gossip payload missing msg_type.")

    raw_delta = payload.get("delta")
    delta: ClusterViewDelta | None = None
    if raw_delta is not None:
        if not isinstance(raw_delta, dict):
            raise TypeError("cluster gossip delta must be a dict.")
        membership = raw_delta.get("membership")
        if not isinstance(membership, dict):
            raise ValueError("cluster gossip delta missing membership map.")
        tombstones = raw_delta.get("tombstones")
        leases_summary = raw_delta.get("leases_summary")
        delta = ClusterViewDelta(
            membership=membership,
            tombstones=tombstones if isinstance(tombstones, dict) else None,
            leases_summary=leases_summary if isinstance(leases_summary, dict) else None,
        )

    return ClusterViewGossipEnvelope(
        msg_type=envelope_type,
        schema_version=str(payload.get("schema_version", "")),
        from_node=str(payload.get("from_node", "")),
        to_node=str(payload.get("to_node", "")),
        view_epoch=int(payload.get("view_epoch", -1)),
        delta_seq=int(payload.get("delta_seq", -1)),
        sent_at_ms=int(payload.get("sent_at_ms", 0)),
        delta=delta,
    )


def make_delta_push(
    *,
    from_node: str,
    to_node: str,
    view_epoch: int,
    delta_seq: int,
    sent_at_ms: int,
    membership: dict[str, dict[str, Any]],
    tombstones: dict[str, int] | None = None,
    leases_summary: dict[str, Any] | None = None,
    schema_version: str = SUPPORTED_SCHEMA,
) -> ClusterViewGossipEnvelope:
    return ClusterViewGossipEnvelope(
        msg_type=MSG_DELTA_PUSH,
        schema_version=schema_version,
        from_node=from_node,
        to_node=to_node,
        view_epoch=view_epoch,
        delta_seq=delta_seq,
        sent_at_ms=sent_at_ms,
        delta=ClusterViewDelta(
            membership=membership,
            tombstones=tombstones,
            leases_summary=leases_summary,
        ),
    )


def make_delta_pull(
    *,
    from_node: str,
    to_node: str,
    view_epoch: int,
    delta_seq: int,
    sent_at_ms: int,
    schema_version: str = SUPPORTED_SCHEMA,
) -> ClusterViewGossipEnvelope:
    return ClusterViewGossipEnvelope(
        msg_type=MSG_DELTA_PULL,
        schema_version=schema_version,
        from_node=from_node,
        to_node=to_node,
        view_epoch=view_epoch,
        delta_seq=delta_seq,
        sent_at_ms=sent_at_ms,
        delta=None,
    )


def make_delta_ack(
    *,
    from_node: str,
    to_node: str,
    view_epoch: int,
    delta_seq: int,
    sent_at_ms: int,
    schema_version: str = SUPPORTED_SCHEMA,
) -> ClusterViewGossipEnvelope:
    return ClusterViewGossipEnvelope(
        msg_type=MSG_DELTA_ACK,
        schema_version=schema_version,
        from_node=from_node,
        to_node=to_node,
        view_epoch=view_epoch,
        delta_seq=delta_seq,
        sent_at_ms=sent_at_ms,
        delta=None,
    )


def make_snapshot(
    *,
    from_node: str,
    to_node: str,
    view_epoch: int,
    delta_seq: int,
    sent_at_ms: int,
    membership: dict[str, dict[str, Any]],
    tombstones: dict[str, int] | None = None,
    leases_summary: dict[str, Any] | None = None,
    schema_version: str = SUPPORTED_SCHEMA,
) -> ClusterViewGossipEnvelope:
    return ClusterViewGossipEnvelope(
        msg_type=MSG_SNAPSHOT,
        schema_version=schema_version,
        from_node=from_node,
        to_node=to_node,
        view_epoch=view_epoch,
        delta_seq=delta_seq,
        sent_at_ms=sent_at_ms,
        delta=ClusterViewDelta(
            membership=membership,
            tombstones=tombstones,
            leases_summary=leases_summary,
        ),
    )


def validate_envelope(envelope: ClusterViewGossipEnvelope) -> None:
    if envelope.msg_type not in SUPPORTED_MSG_TYPES:
        raise ValueError(f"Unsupported gossip msg_type: {envelope.msg_type}")
    if envelope.schema_version != SUPPORTED_SCHEMA:
        raise ValueError(
            f"Unsupported gossip schema_version: {envelope.schema_version}"
        )
    if envelope.view_epoch < 0:
        raise ValueError("view_epoch must be >= 0")
    if envelope.delta_seq < 0:
        raise ValueError("delta_seq must be >= 0")
    if envelope.sent_at_ms <= 0:
        raise ValueError("sent_at_ms must be > 0")
    if envelope.msg_type in {MSG_DELTA_PUSH, MSG_SNAPSHOT} and envelope.delta is None:
        raise ValueError(f"{envelope.msg_type} requires delta payload")


__all__ = [
    "MSG_DELTA_PUSH",
    "MSG_DELTA_PULL",
    "MSG_DELTA_ACK",
    "MSG_SNAPSHOT",
    "SUPPORTED_SCHEMA",
    "SUPPORTED_MSG_TYPES",
    "ClusterViewDelta",
    "ClusterViewGossipEnvelope",
    "parse_envelope",
    "make_delta_push",
    "make_delta_pull",
    "make_delta_ack",
    "make_snapshot",
    "validate_envelope",
]
