from __future__ import annotations

import asyncio
import time
import uuid
from dataclasses import dataclass
from typing import Any, Awaitable, Callable

from flownet.core.request import FlownetPacket
from flownet.runtime.cluster_gossip_protocol import (
    MSG_DELTA_ACK,
    MSG_DELTA_PULL,
    MSG_SNAPSHOT,
    make_delta_pull,
    make_snapshot,
    parse_envelope,
    validate_envelope,
)


@dataclass(frozen=True)
class GossipServiceConfig:
    enabled: bool = True
    interval_s: float = 3.0
    snapshot_interval_s: float = 30.0
    strong_interval_s: float | None = None
    weak_interval_s: float | None = None
    strong_snapshot_interval_s: float | None = None
    weak_snapshot_interval_s: float | None = None
    strong_locality_tiers: tuple[str, ...] = ("zone", "region")
    aging_enabled: bool = False
    aging_interval_s: float = 5.0
    stale_after_s: float = 15.0
    remove_after_s: float = 60.0


class ClusterGossipService:
    """Periodic best-effort anti-entropy sender (pull + occasional snapshot)."""

    def __init__(
        self,
        *,
        cluster_view,
        local_address_fn: Callable[[], str],
        send_packet_fn: Callable[[str, FlownetPacket], Awaitable[Any]],
        config: GossipServiceConfig | None = None,
        monotonic_fn: Callable[[], float] | None = None,
        wallclock_fn: Callable[[], float] | None = None,
        log_fn: Callable[[str], None] | None = None,
    ) -> None:
        self._cluster_view = cluster_view
        self._local_address_fn = local_address_fn
        self._send_packet_fn = send_packet_fn
        self._config = config or GossipServiceConfig()
        self._monotonic_fn = monotonic_fn or time.monotonic
        self._wallclock_fn = wallclock_fn or time.time
        self._log_fn = log_fn or print

        self._task = None
        self._stopped = True
        self._peer_cursor = 0
        self._delta_seq = 0
        self._last_snapshot_at = float("-inf")
        self._last_aging_at = float("-inf")
        self._last_pull_at_by_peer: dict[str, float] = {}
        self._last_snapshot_at_by_peer: dict[str, float] = {}
        self._last_pull_reply_ack_seq_by_peer: dict[str, int] = {}
        self._last_pull_reply_snapshot_seq_by_peer: dict[str, int] = {}
        self._stats: dict[str, Any] = {
            "running": False,
            "sent_total": 0,
            MSG_DELTA_PULL: 0,
            MSG_SNAPSHOT: 0,
            "last_peer": None,
            "last_peer_tier": None,
            "last_error": None,
            "errors": 0,
            "last_sent_ms": None,
            "sent_by_tier": {"strong": 0, "weak": 0},
            "pull_replies": 0,
            "pull_reply_ack": 0,
            "pull_reply_ack_ignored": 0,
            "pull_reply_snapshot": 0,
            "pull_reply_snapshot_ignored": 0,
            "aging_degraded_total": 0,
            "aging_removed_total": 0,
            "last_aging_ms": None,
        }

    def start(self, submit_coro: Callable[[Awaitable[Any]], Any]) -> bool:
        if self._task is not None:
            return False
        if not self._config.enabled or self._tick_interval_s() <= 0:
            self._stats["running"] = False
            return False
        self._stopped = False
        self._stats["running"] = True
        self._task = submit_coro(self._run())
        return True

    def stop(self) -> bool:
        if self._task is None:
            return False
        self._stopped = True
        task = self._task
        self._task = None
        if hasattr(task, "cancel"):
            task.cancel()
        self._stats["running"] = False
        return True

    def stats_snapshot(self) -> dict[str, Any]:
        snapshot = dict(self._stats)
        sent_by_tier = snapshot.get("sent_by_tier")
        if isinstance(sent_by_tier, dict):
            snapshot["sent_by_tier"] = dict(sent_by_tier)
        return snapshot

    async def tick_once(self) -> None:
        await self._gossip_once()
        self._age_once_if_due()

    async def _run(self) -> None:
        try:
            while not self._stopped:
                await self._gossip_once()
                self._age_once_if_due()
                await asyncio.sleep(self._tick_interval_s())
        except asyncio.CancelledError:
            return
        except Exception as exc:  # pragma: no cover - defensive
            self._record_error(exc)
            self._log_fn(f"[ClusterGossip] service crashed: {exc}")

    async def _gossip_once(self) -> None:
        peers = self._peer_addresses()
        if not peers:
            return
        peer = self._next_peer(peers)
        if peer is None:
            return
        now_ms = int(self._wallclock_fn() * 1000.0)
        view_epoch = self._cluster_view.view_epoch()
        self._delta_seq += 1
        pull = make_delta_pull(
            from_node=self._cluster_view.local_node_id(),
            to_node=peer["node_id"],
            view_epoch=view_epoch,
            delta_seq=self._delta_seq,
            sent_at_ms=now_ms,
        )
        pull_reply = await self._send(
            peer["address"],
            pull.as_dict(),
            MSG_DELTA_PULL,
            peer_tier=peer["tier"],
            expect_reply=True,
        )
        self._handle_pull_reply(pull_reply)
        if self._pull_tiered_cadence_enabled():
            self._last_pull_at_by_peer[peer["node_id"]] = self._monotonic_fn()

        if self._should_send_snapshot(peer["node_id"], peer["tier"]):
            self._delta_seq += 1
            snapshot = make_snapshot(
                from_node=self._cluster_view.local_node_id(),
                to_node=peer["node_id"],
                view_epoch=view_epoch,
                delta_seq=self._delta_seq,
                sent_at_ms=now_ms,
                membership=self._membership_payload(),
            )
            await self._send(peer["address"], snapshot.as_dict(), MSG_SNAPSHOT, peer_tier=peer["tier"])
            now = self._monotonic_fn()
            if self._snapshot_tiered_cadence_enabled():
                self._last_snapshot_at_by_peer[peer["node_id"]] = now
            else:
                self._last_snapshot_at = now

    async def _send(
        self,
        address: str,
        payload: dict[str, Any],
        packet_type: str,
        *,
        peer_tier: str,
        expect_reply: bool = False,
    ) -> Any:
        packet_payload = payload
        if expect_reply:
            packet_payload = dict(payload)
            packet_payload.setdefault("reply_to", self._local_address_fn())
        packet = FlownetPacket(
            type=packet_type,
            uuid=uuid.uuid4().hex,
            reply=expect_reply,
            packet_content=packet_payload,
        )
        try:
            result = await self._send_packet_fn(address, packet)
        except Exception as exc:  # pragma: no cover - network errors are expected
            self._record_error(exc)
            self._log_fn(f"[ClusterGossip] send failed to {address}: {exc}")
            return None
        self._stats["sent_total"] = self._stats.get("sent_total", 0) + 1
        self._stats[packet_type] = self._stats.get(packet_type, 0) + 1
        self._stats["last_peer"] = address
        self._stats["last_peer_tier"] = peer_tier
        self._stats["last_sent_ms"] = int(self._wallclock_fn() * 1000.0)
        sent_by_tier = self._stats.get("sent_by_tier")
        if isinstance(sent_by_tier, dict):
            sent_by_tier[peer_tier] = int(sent_by_tier.get(peer_tier, 0)) + 1
        return result

    def _peer_addresses(self) -> list[dict[str, str]]:
        local_node = self._cluster_view.local_node_id()
        peers = []
        for node in self._cluster_view.list_nodes():
            if node.node_id == local_node:
                continue
            peers.append(
                {
                    "node_id": node.node_id,
                    "address": node.address,
                    "tier": self._peer_tier(node.node_id),
                }
            )
        peers.sort(key=lambda item: item["node_id"])
        return peers

    def _membership_payload(self) -> dict[str, dict[str, Any]]:
        payload: dict[str, dict[str, Any]] = {}
        for node in self._cluster_view.list_nodes():
            data = node.as_dict()
            data.pop("node_id", None)
            payload[node.node_id] = data
        return payload

    def _should_send_snapshot(self, peer_id: str, peer_tier: str) -> bool:
        now = self._monotonic_fn()
        if self._snapshot_tiered_cadence_enabled():
            interval = self._snapshot_interval_for_tier(peer_tier)
            if interval <= 0:
                return False
            last = self._last_snapshot_at_by_peer.get(peer_id, float("-inf"))
            return (now - last) >= interval
        if self._config.snapshot_interval_s <= 0:
            return False
        return (now - self._last_snapshot_at) >= self._config.snapshot_interval_s

    def _peer_tier(self, node_id: str) -> str:
        locality_tier_fn = getattr(self._cluster_view, "locality_tier", None)
        if callable(locality_tier_fn):
            try:
                tier = str(locality_tier_fn(node_id))
            except Exception:  # pragma: no cover - defensive against user hooks
                return "weak"
            if tier in self._config.strong_locality_tiers:
                return "strong"
        return "weak"

    def _tick_interval_s(self) -> float:
        if self._pull_tiered_cadence_enabled():
            candidates = [
                self._config.interval_s,
                self._config.strong_interval_s,
                self._config.weak_interval_s,
            ]
            positive = [float(value) for value in candidates if isinstance(value, (int, float)) and value > 0]
            if positive:
                return min(positive)
        return self._config.interval_s

    def _pull_tiered_cadence_enabled(self) -> bool:
        return self._config.strong_interval_s is not None or self._config.weak_interval_s is not None

    def _snapshot_tiered_cadence_enabled(self) -> bool:
        return (
            self._config.strong_snapshot_interval_s is not None
            or self._config.weak_snapshot_interval_s is not None
        )

    def _pull_interval_for_tier(self, peer_tier: str) -> float:
        if peer_tier == "strong" and self._config.strong_interval_s is not None:
            return self._config.strong_interval_s
        if peer_tier == "weak" and self._config.weak_interval_s is not None:
            return self._config.weak_interval_s
        return self._config.interval_s

    def _snapshot_interval_for_tier(self, peer_tier: str) -> float:
        if peer_tier == "strong" and self._config.strong_snapshot_interval_s is not None:
            return self._config.strong_snapshot_interval_s
        if peer_tier == "weak" and self._config.weak_snapshot_interval_s is not None:
            return self._config.weak_snapshot_interval_s
        return self._config.snapshot_interval_s

    def _next_peer(self, peers: list[dict[str, str]]) -> dict[str, str] | None:
        if not peers:
            return None
        if not self._pull_tiered_cadence_enabled():
            peer = peers[self._peer_cursor % len(peers)]
            self._peer_cursor += 1
            return peer
        now = self._monotonic_fn()
        start = self._peer_cursor % len(peers)
        for offset in range(len(peers)):
            idx = (start + offset) % len(peers)
            peer = peers[idx]
            interval = self._pull_interval_for_tier(peer["tier"])
            if interval <= 0:
                continue
            last = self._last_pull_at_by_peer.get(peer["node_id"], float("-inf"))
            if (now - last) >= interval:
                self._peer_cursor = idx + 1
                return peer
        return None

    def _record_error(self, exc: Exception) -> None:
        self._stats["errors"] = self._stats.get("errors", 0) + 1
        self._stats["last_error"] = str(exc)

    def _age_once_if_due(self) -> None:
        if not self._config.aging_enabled:
            return
        if self._config.aging_interval_s > 0:
            now_mono = self._monotonic_fn()
            if (now_mono - self._last_aging_at) < self._config.aging_interval_s:
                return
            self._last_aging_at = now_mono
        age_fn = getattr(self._cluster_view, "age_nodes", None)
        if not callable(age_fn):
            return
        try:
            result = age_fn(
                now=self._wallclock_fn(),
                stale_after_s=self._config.stale_after_s,
                remove_after_s=self._config.remove_after_s,
            )
        except Exception as exc:  # pragma: no cover - defensive
            self._record_error(exc)
            self._log_fn(f"[ClusterGossip] aging failed: {exc}")
            return
        if not isinstance(result, dict):
            return
        self._stats["aging_degraded_total"] = (
            self._stats.get("aging_degraded_total", 0) + int(result.get("degraded", 0))
        )
        self._stats["aging_removed_total"] = (
            self._stats.get("aging_removed_total", 0) + int(result.get("removed", 0))
        )
        self._stats["last_aging_ms"] = int(self._wallclock_fn() * 1000.0)

    def _handle_pull_reply(self, response: Any) -> None:
        if not isinstance(response, dict):
            return
        self._stats["pull_replies"] = self._stats.get("pull_replies", 0) + 1
        gossip_reply = response.get("gossip_reply")
        if not isinstance(gossip_reply, dict):
            return
        ack_payload = gossip_reply.get("ack")
        if isinstance(ack_payload, dict):
            try:
                ack_envelope = parse_envelope(ack_payload, msg_type=MSG_DELTA_ACK)
                validate_envelope(ack_envelope)
            except Exception as exc:  # pragma: no cover - malformed peer response
                self._record_error(exc)
                self._log_fn(f"[ClusterGossip] invalid pull ack payload: {exc}")
            else:
                peer = ack_envelope.from_node
                prev = int(self._last_pull_reply_ack_seq_by_peer.get(peer, -1))
                if ack_envelope.delta_seq <= prev:
                    self._stats["pull_reply_ack_ignored"] = (
                        self._stats.get("pull_reply_ack_ignored", 0) + 1
                    )
                else:
                    self._last_pull_reply_ack_seq_by_peer[peer] = int(ack_envelope.delta_seq)
                    self._stats["pull_reply_ack"] = self._stats.get("pull_reply_ack", 0) + 1
        snapshot_payload = gossip_reply.get("snapshot")
        if not isinstance(snapshot_payload, dict):
            return
        try:
            snapshot_envelope = parse_envelope(snapshot_payload, msg_type=MSG_SNAPSHOT)
            validate_envelope(snapshot_envelope)
        except Exception as exc:  # pragma: no cover - malformed peer response
            self._record_error(exc)
            self._log_fn(f"[ClusterGossip] invalid pull snapshot payload: {exc}")
            return
        peer = snapshot_envelope.from_node
        prev_snapshot = int(self._last_pull_reply_snapshot_seq_by_peer.get(peer, -1))
        if snapshot_envelope.delta_seq <= prev_snapshot:
            self._stats["pull_reply_snapshot_ignored"] = (
                self._stats.get("pull_reply_snapshot_ignored", 0) + 1
            )
            return
        self._last_pull_reply_snapshot_seq_by_peer[peer] = int(snapshot_envelope.delta_seq)
        try:
            delta = snapshot_envelope.delta
            if delta is not None:
                self._cluster_view.apply_gossip_delta(
                    membership=delta.membership,
                    tombstones=delta.tombstones,
                    delta_ts_ms=snapshot_envelope.sent_at_ms,
                    source_node=snapshot_envelope.from_node,
                    source_delta_seq=snapshot_envelope.delta_seq,
                )
        except Exception as exc:  # pragma: no cover - malformed peer response
            self._record_error(exc)
            self._log_fn(f"[ClusterGossip] invalid pull snapshot payload: {exc}")
            return
        self._stats["pull_reply_snapshot"] = self._stats.get("pull_reply_snapshot", 0) + 1


__all__ = [
    "GossipServiceConfig",
    "ClusterGossipService",
]
