#!/usr/bin/env python3
"""orcrdt -- dependency-free state-based CRDT kit (OR-Set, PN-Counter, LWW-Register).

Design rules
------------
* State-based (CvRDT): every replica exposes ``merge(other, ...) -> new`` implementing
  a join over a semilattice. ``merge`` is PURE -- it never mutates ``self`` or its
  arguments, so a replica can re-deliver a snapshot any number of times.
* A merge that is not commutative / associative / idempotent makes replicas diverge.
  That is not a matter of opinion, so ``test_crdt.py`` checks it two ways: random
  gossip simulation over many seeds, and mutation -- deliberately broken merges must
  be *detected* by the same check.
* Elements of ``ORSet`` must be JSON scalars (str/int/float/bool/None). This is a
  deliberate constraint: it makes ``to_state()`` lossless instead of repr-based.
* No third-party imports. Python 3.9+.

Public surface
--------------
ORSet         add-wins observed-remove set
PNCounter     increment/decrement counter (per-replica vectors, pointwise max)
LWWRegister   last-writer-wins register, stamp = (timestamp, replica)
CRDTError     raised on malformed state or a broken local operation
"""

from __future__ import annotations

import json

__all__ = ["ORSet", "PNCounter", "LWWRegister", "CRDTError"]

_SCALARS = (str, int, float, bool, type(None))


class CRDTError(ValueError):
    """Malformed state, non-scalar element, or a non-monotone local operation."""


def _check_scalar(elem, where):
    if isinstance(elem, _SCALARS):
        return elem
    raise CRDTError(
        "%s: element must be a JSON scalar (str/int/float/bool/None), got %r"
        % (where, type(elem).__name__)
    )


def _canon(elem):
    """Canonical, order-stable, comparable form of a scalar element."""
    return json.dumps(elem, sort_keys=True)


def _sorted_tags(tags):
    return sorted(tags, key=lambda t: (t[0], t[1]))


def _pmax(a, b):
    """Pointwise max over the union of two non-negative-int vectors."""
    out = dict(a)
    for k, v in b.items():
        if v > out.get(k, 0):
            out[k] = v
    return out


# ---------------------------------------------------------------------------
# OR-Set
# ---------------------------------------------------------------------------
class ORSet:
    """Observed-Remove Set with add-wins semantics.

    State = (``_adds``: elem -> set of unique tags, ``_removed``: set of tags).
    A tag is ``(replica, counter)`` and is unique per local ``add``.

    ``contains(e)`` is true iff ``e`` has at least one tag that no ``remove`` has
    tombstoned. ``remove(e)`` tombstones only the tags the caller had *observed*,
    so a concurrent ``add`` survives the merge -- that is the add-wins rule.
    """

    __slots__ = ("replica", "_adds", "_removed", "_clock")

    def __init__(self, replica="r0"):
        if not isinstance(replica, str) or not replica:
            raise CRDTError("replica id must be a non-empty string")
        self.replica = replica
        self._adds = {}        # elem -> set(tag)
        self._removed = set()  # set(tag)
        self._clock = 0

    # -- local mutations ----------------------------------------------------
    def add(self, elem):
        _check_scalar(elem, "ORSet.add")
        self._clock += 1
        tag = (self.replica, self._clock)
        self._adds.setdefault(elem, set()).add(tag)
        return tag

    def remove(self, elem):
        """Tombstone every tag of ``elem`` observed locally. Returns tags removed."""
        _check_scalar(elem, "ORSet.remove")
        tags = self._adds.get(elem)
        if not tags:
            return 0
        fresh = [t for t in tags if t not in self._removed]
        self._removed.update(fresh)
        return len(fresh)

    # -- queries ------------------------------------------------------------
    def contains(self, elem):
        return any(t not in self._removed for t in self._adds.get(elem, ()))

    def elements(self):
        return frozenset(e for e in self._adds if self.contains(e))

    def value_key(self):
        return tuple(sorted(_canon(e) for e in self.elements()))

    def __contains__(self, elem):
        return self.contains(elem)

    def __len__(self):
        return len(self.elements())

    def __iter__(self):
        return iter(sorted(self.elements(), key=_canon))

    # -- merge / state ------------------------------------------------------
    def merge(self, *others):
        out = ORSet(self.replica)
        states = (self,) + others
        for st in states:
            if not isinstance(st, ORSet):
                raise CRDTError("ORSet.merge: expected ORSet, got %r" % type(st).__name__)
            for e, tags in st._adds.items():
                out._adds.setdefault(e, set()).update(tags)
            out._removed.update(st._removed)
        out._clock = max(st._clock for st in states)
        return out

    def state_key(self):
        adds = tuple(
            sorted(
                (_canon(e), tuple(_sorted_tags(t))) for e, t in self._adds.items()
            )
        )
        return (adds, tuple(_sorted_tags(self._removed)))

    def __eq__(self, other):
        return isinstance(other, ORSet) and self.state_key() == other.state_key()

    def __hash__(self):
        return hash(self.state_key())

    def to_state(self):
        return {
            "type": "ORSet",
            "replica": self.replica,
            "clock": self._clock,
            "adds": [
                [_canon(e), [[t[0], t[1]] for t in _sorted_tags(tags)]]
                for e, tags in sorted(self._adds.items(), key=lambda kv: _canon(kv[0]))
            ],
            "removed": [[t[0], t[1]] for t in _sorted_tags(self._removed)],
        }

    @classmethod
    def from_state(cls, d):
        try:
            if d["type"] != "ORSet":
                raise CRDTError("ORSet.from_state: wrong type tag %r" % (d["type"],))
            out = cls(d["replica"])
            out._clock = int(d["clock"])
            for raw_e, raw_tags in d["adds"]:
                elem = json.loads(raw_e)
                _check_scalar(elem, "ORSet.from_state")
                out._adds.setdefault(elem, set()).update(
                    (t[0], int(t[1])) for t in raw_tags
                )
            out._removed = {(t[0], int(t[1])) for t in d["removed"]}
            if out._clock < 0 or any(c <= 0 for _, c in out._removed):
                raise CRDTError("ORSet.from_state: non-positive clock/counter")
            return out
        except (KeyError, TypeError, ValueError) as exc:
            if isinstance(exc, CRDTError):
                raise
            raise CRDTError("ORSet.from_state: malformed state: %s" % exc) from exc

    def dumps(self):
        return json.dumps(self.to_state(), sort_keys=True)

    @classmethod
    def loads(cls, s):
        try:
            d = json.loads(s)
        except ValueError as exc:
            raise CRDTError("ORSet.loads: not JSON: %s" % exc) from exc
        if not isinstance(d, dict):
            raise CRDTError("ORSet.loads: state must be a JSON object")
        return cls.from_state(d)


# ---------------------------------------------------------------------------
# PN-Counter
# ---------------------------------------------------------------------------
class PNCounter:
    """Increment/decrement counter.

    State = two per-replica vectors (``inc``, ``dec``); merge = pointwise max.
    Each replica writes only its own slot, which is what makes pointwise max a
    correct join (a replica's own slot only ever grows).
    """

    __slots__ = ("replica", "inc", "dec")

    def __init__(self, replica="r0"):
        if not isinstance(replica, str) or not replica:
            raise CRDTError("replica id must be a non-empty string")
        self.replica = replica
        self.inc = {}
        self.dec = {}

    def _bump(self, vec, n):
        if not isinstance(n, int) or isinstance(n, bool) or n < 0:
            raise CRDTError("PNCounter: delta must be a non-negative int, got %r" % (n,))
        vec[self.replica] = vec.get(self.replica, 0) + n

    def increment(self, n=1):
        self._bump(self.inc, n)

    def decrement(self, n=1):
        self._bump(self.dec, n)

    def value(self):
        return sum(self.inc.values()) - sum(self.dec.values())

    def value_key(self):
        return self.value()

    def merge(self, *others):
        out = PNCounter(self.replica)
        for st in (self,) + others:
            if not isinstance(st, PNCounter):
                raise CRDTError("PNCounter.merge: expected PNCounter, got %r" % type(st).__name__)
            out.inc = _pmax(out.inc, st.inc)
            out.dec = _pmax(out.dec, st.dec)
        return out

    def state_key(self):
        return (tuple(sorted(self.inc.items())), tuple(sorted(self.dec.items())))

    def __eq__(self, other):
        return isinstance(other, PNCounter) and self.state_key() == other.state_key()

    def __hash__(self):
        return hash(self.state_key())

    def to_state(self):
        return {
            "type": "PNCounter",
            "replica": self.replica,
            "inc": sorted(self.inc.items()),
            "dec": sorted(self.dec.items()),
        }

    @classmethod
    def from_state(cls, d):
        try:
            if d["type"] != "PNCounter":
                raise CRDTError("PNCounter.from_state: wrong type tag %r" % (d["type"],))
            out = cls(d["replica"])
            for key in ("inc", "dec"):
                vec = {}
                for k, v in d[key]:
                    if not isinstance(k, str) or not isinstance(v, int) or v < 0:
                        raise CRDTError("PNCounter.from_state: bad entry in %s" % key)
                    vec[k] = v
                setattr(out, key, vec)
            return out
        except (KeyError, TypeError, ValueError) as exc:
            if isinstance(exc, CRDTError):
                raise
            raise CRDTError("PNCounter.from_state: malformed state: %s" % exc) from exc

    def dumps(self):
        return json.dumps(self.to_state(), sort_keys=True)

    @classmethod
    def loads(cls, s):
        try:
            d = json.loads(s)
        except ValueError as exc:
            raise CRDTError("PNCounter.loads: not JSON: %s" % exc) from exc
        if not isinstance(d, dict):
            raise CRDTError("PNCounter.loads: state must be a JSON object")
        return cls.from_state(d)


# ---------------------------------------------------------------------------
# LWW-Register
# ---------------------------------------------------------------------------
class LWWRegister:
    """Last-writer-wins register.

    Stamp = ``(timestamp: int, replica: str)``. Merge takes the maximum stamp, so
    the winner is total-order-determined and independent of argument order. Local
    writes must advance the stamp strictly (``set`` raises otherwise), and two
    different values at an *identical* stamp are refused rather than silently
    broken -- that only happens if two nodes reuse one replica id.
    """

    __slots__ = ("replica", "_value", "_stamp")

    def __init__(self, replica="r0"):
        if not isinstance(replica, str) or not replica:
            raise CRDTError("replica id must be a non-empty string")
        self.replica = replica
        self._value = None
        self._stamp = (-1, "")

    def set(self, value, timestamp):
        if not isinstance(timestamp, int) or isinstance(timestamp, bool):
            raise CRDTError("LWWRegister.set: timestamp must be an int")
        stamp = (timestamp, self.replica)
        if stamp <= self._stamp:
            raise CRDTError(
                "LWWRegister.set: non-monotone stamp %r <= %r" % (stamp, self._stamp)
            )
        self._value = value
        self._stamp = stamp
        return stamp

    def get(self):
        return self._value

    @property
    def stamp(self):
        return self._stamp

    def value_key(self):
        return self._value

    def merge(self, *others):
        out = LWWRegister(self.replica)
        out._value, out._stamp = self._value, self._stamp
        for st in (self,) + others:
            if not isinstance(st, LWWRegister):
                raise CRDTError(
                    "LWWRegister.merge: expected LWWRegister, got %r" % type(st).__name__
                )
            if st._stamp > out._stamp:
                out._value, out._stamp = st._value, st._stamp
            elif st._stamp == out._stamp and st._value != out._value:
                raise CRDTError(
                    "LWWRegister.merge: conflicting values at identical stamp %r"
                    % (st._stamp,)
                )
        return out

    def state_key(self):
        return (self._stamp, _canon(self._value) if isinstance(self._value, _SCALARS) else repr(self._value))

    def __eq__(self, other):
        return isinstance(other, LWWRegister) and self.state_key() == other.state_key()

    def __hash__(self):
        return hash(self.state_key())

    def to_state(self):
        return {
            "type": "LWWRegister",
            "replica": self.replica,
            "stamp": [self._stamp[0], self._stamp[1]],
            "value": self._value,
        }

    @classmethod
    def from_state(cls, d):
        try:
            if d["type"] != "LWWRegister":
                raise CRDTError("LWWRegister.from_state: wrong type tag %r" % (d["type"],))
            out = cls(d["replica"])
            out._stamp = (int(d["stamp"][0]), d["stamp"][1])
            out._value = d["value"]
            return out
        except (KeyError, TypeError, ValueError, IndexError) as exc:
            if isinstance(exc, CRDTError):
                raise
            raise CRDTError("LWWRegister.from_state: malformed state: %s" % exc) from exc

    def dumps(self):
        return json.dumps(self.to_state(), sort_keys=True)

    @classmethod
    def loads(cls, s):
        try:
            d = json.loads(s)
        except ValueError as exc:
            raise CRDTError("LWWRegister.loads: not JSON: %s" % exc) from exc
        if not isinstance(d, dict):
            raise CRDTError("LWWRegister.loads: state must be a JSON object")
        return cls.from_state(d)


if __name__ == "__main__":
    # 12-line smoke illustration: concurrent remove loses to concurrent add.
    a, b = ORSet("a"), ORSet("b")
    a.add("x")
    b = ORSet.loads(a.dumps())
    a.remove("x")
    b.add("x")
    m = a.merge(b)
    print("orcrdt smoke: a.elements=%r b.elements=%r merged=%r"
          % (sorted(a), sorted(b), sorted(m)))