from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

MergeFn = Callable[[Any, Any], Any]


@dataclass
class SharedValue:
    value: Any
    version: int = 0
    merge: MergeFn | None = None


@dataclass
class SharedState:
    """Replicated key/value map with merge rules.

    Great for blockchain-like replicated state where every peer keeps and
    periodically exchanges its copy.
    """

    _values: dict[str, SharedValue] = field(default_factory=dict)

    def set(self, key: str, value: Any, merge: MergeFn | None = None) -> None:
        current = self._values.get(key)
        if current is None:
            self._values[key] = SharedValue(value=value, version=1, merge=merge)
            return

        current.version += 1
        current.value = value
        if merge is not None:
            current.merge = merge

    def get(self, key: str, default: Any = None) -> Any:
        node = self._values.get(key)
        return node.value if node is not None else default

    def snapshot(self) -> dict[str, dict[str, Any]]:
        return {
            k: {"value": v.value, "version": v.version}
            for k, v in self._values.items()
        }

    def apply_remote_snapshot(self, remote: dict[str, dict[str, Any]]) -> None:
        for key, entry in remote.items():
            remote_version = int(entry.get("version", 0))
            remote_value = entry.get("value")
            local = self._values.get(key)

            if local is None:
                self._values[key] = SharedValue(value=remote_value, version=remote_version)
                continue

            if remote_version > local.version:
                local.version = remote_version
                local.value = remote_value
                continue

            if remote_version == local.version and local.merge is not None:
                local.value = local.merge(local.value, remote_value)
