"""PlanDelta CRDT primitives: OR-Set and LWW registers for deterministic merging.

This file implements a minimal set of CRDT-like types used to make PlanDelta
merges deterministic on reconnection. They are intentionally small but include
merge operations and vector-clock-like timestamps.
"""
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
import time


@dataclass
class ORSet:
    # store JSON-keyed entries so unhashable objects (dicts) can be used
    add_set: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    rem_set: Dict[str, Dict[str, Any]] = field(default_factory=dict)

    def _key(self, element: Any) -> str:
        import json as _json

        return _json.dumps(element, sort_keys=True, separators=(',', ':'))

    def add(self, element: Any, ts: Optional[float] = None):
        k = self._key(element)
        self.add_set[k] = {"val": element, "ts": time.time() if ts is None else ts}

    def remove(self, element: Any, ts: Optional[float] = None):
        k = self._key(element)
        self.rem_set[k] = {"val": element, "ts": time.time() if ts is None else ts}

    def value(self) -> List[Any]:
        res: List[Any] = []
        for k, info in self.add_set.items():
            add_ts = info["ts"]
            rem_ts = self.rem_set.get(k, {}).get("ts", -1)
            if add_ts > rem_ts:
                # original object
                res.append(self.add_set[k]["val"])
        return res

    def merge(self, other: "ORSet") -> "ORSet":
        res = ORSet()
        # merge add_set
        for k, info in {**self.add_set, **other.add_set}.items():
            # pick newest
            left = self.add_set.get(k)
            right = other.add_set.get(k)
            if left and right:
                res.add_set[k] = left if left["ts"] >= right["ts"] else right
            else:
                res.add_set[k] = info
        # merge rem_set similarly
        for k, info in {**self.rem_set, **other.rem_set}.items():
            left = self.rem_set.get(k)
            right = other.rem_set.get(k)
            if left and right:
                res.rem_set[k] = left if left["ts"] >= right["ts"] else right
            else:
                res.rem_set[k] = info
        return res


@dataclass
class LWWRegister:
    value: Any = None
    ts: float = 0.0

    def write(self, value: Any, ts: Optional[float] = None):
        self.value = value
        self.ts = time.time() if ts is None else ts

    def read(self):
        return self.value

    def merge(self, other: "LWWRegister") -> "LWWRegister":
        if other.ts >= self.ts:
            return LWWRegister(value=other.value, ts=other.ts)
        return LWWRegister(value=self.value, ts=self.ts)


@dataclass
class PlanDelta:
    # a minimal plan delta carrying a set of commitments and a priority register
    commitments: ORSet = field(default_factory=ORSet)
    priority: LWWRegister = field(default_factory=LWWRegister)
    origin: str = ""

    def merge(self, other: "PlanDelta") -> "PlanDelta":
        p = PlanDelta()
        p.commitments = self.commitments.merge(other.commitments)
        p.priority = self.priority.merge(other.priority)
        # origin is not merged; keep latest priority's origin heuristically
        p.origin = other.origin if other.priority.ts >= self.priority.ts else self.origin
        return p

    def snapshot(self) -> Dict[str, Any]:
        return {
            "commitments": self.commitments.value(),
            "priority": self.priority.read(),
            "priority_ts": self.priority.ts,
            "origin": self.origin,
        }

    @classmethod
    def from_snapshot(cls, snapshot: Dict[str, Any]) -> "PlanDelta":
        delta = cls(origin=snapshot.get("origin", ""))
        for commitment in snapshot.get("commitments", []):
            delta.commitments.add(commitment, ts=0.0)
        delta.priority = LWWRegister(value=snapshot.get("priority"), ts=snapshot.get("priority_ts", 0.0))
        return delta