"""Hybrid Logical Clock (HLC) for causal ordering across nodes.

An HLC timestamp combines wall-clock time with a logical counter to
provide a monotonically increasing, causally consistent timestamp that
tolerates clock skew between nodes.

Properties:
  - Always >= physical wall time (bounded drift)
  - Monotonically increasing within a node
  - Causally consistent: if event A happened-before event B, hlc(A) < hlc(B)
  - Comparable across nodes (unlike pure Lamport clocks)
  - Serializable as a 64-bit integer or ISO-like string

Format:
  Physical time (ms since epoch) in upper 44 bits + logical counter in lower 20 bits.
  This gives ~557 years of range and 1M events per millisecond.

  String format: "{wall_ms}.{counter:06d}.{node_id}"
  Example: "1741459200000.000001.node-1"

Reference: Kulkarni et al., "Logical Physical Clocks and Consistent
Snapshots in Globally Distributed Databases" (2014).
"""
from __future__ import annotations

import time
from dataclasses import dataclass


# Maximum allowed clock skew before we refuse to merge (5 minutes in ms).
# If a remote timestamp is more than this far in the future, something
# is very wrong with NTP and we should flag it rather than silently accept.
MAX_CLOCK_SKEW_MS = 5 * 60 * 1000


@dataclass(frozen=True, order=True)
class HLCTimestamp:
    """A Hybrid Logical Clock timestamp.

    Ordered by (wall_ms, counter, node_id) — wall time dominates,
    counter breaks ties within the same millisecond, node_id is
    the final tiebreaker for deterministic ordering.
    """
    wall_ms: int
    counter: int
    node_id: str = ""

    def __str__(self) -> str:
        return f"{self.wall_ms}.{self.counter:06d}.{self.node_id}"

    @classmethod
    def from_str(cls, s: str) -> HLCTimestamp:
        """Parse a string-encoded HLC timestamp."""
        parts = s.split(".", 2)
        if len(parts) < 2:
            raise ValueError(f"Invalid HLC timestamp: {s}")
        wall_ms = int(parts[0])
        counter = int(parts[1])
        node_id = parts[2] if len(parts) > 2 else ""
        return cls(wall_ms=wall_ms, counter=counter, node_id=node_id)

    def to_int(self) -> int:
        """Encode as a 64-bit integer (wall_ms << 20 | counter).

        Useful for SQLite INTEGER columns. Drops node_id.
        """
        return (self.wall_ms << 20) | (self.counter & 0xFFFFF)

    @classmethod
    def from_int(cls, val: int, node_id: str = "") -> HLCTimestamp:
        """Decode from a 64-bit integer."""
        wall_ms = val >> 20
        counter = val & 0xFFFFF
        return cls(wall_ms=wall_ms, counter=counter, node_id=node_id)

    @classmethod
    def zero(cls) -> HLCTimestamp:
        """Return the zero/origin timestamp."""
        return cls(wall_ms=0, counter=0, node_id="")


class HLC:
    """A Hybrid Logical Clock instance for a single node.

    Usage:
        clock = HLC(node_id="node-1")
        ts = clock.now()           # Local event
        ts = clock.receive(remote) # Received a message with remote timestamp
    """

    def __init__(self, node_id: str, wall_clock: callable | None = None):
        """Initialize the HLC.

        Args:
            node_id: Unique identifier for this node.
            wall_clock: Optional function returning current time in ms.
                        Defaults to time.time() * 1000.
        """
        self.node_id = node_id
        self._wall_clock = wall_clock or (lambda: int(time.time() * 1000))
        self._last = HLCTimestamp.zero()

    def now(self) -> HLCTimestamp:
        """Generate a timestamp for a local event.

        Returns a timestamp that is:
          - Greater than the last timestamp issued
          - Greater than or equal to the current wall clock
        """
        physical = self._wall_clock()

        if physical > self._last.wall_ms:
            # Wall clock advanced — reset counter
            ts = HLCTimestamp(wall_ms=physical, counter=0, node_id=self.node_id)
        else:
            # Wall clock hasn't advanced (same ms or clock went backward)
            # Increment counter on the last known timestamp
            ts = HLCTimestamp(
                wall_ms=self._last.wall_ms,
                counter=self._last.counter + 1,
                node_id=self.node_id,
            )

        self._last = ts
        return ts

    def receive(self, remote: HLCTimestamp) -> HLCTimestamp:
        """Update the clock upon receiving a remote timestamp.

        Merges the remote timestamp with the local clock to maintain
        the causal ordering property: the returned timestamp is
        greater than both the local last timestamp and the remote one.

        Args:
            remote: Timestamp from a remote node.

        Returns:
            A new timestamp reflecting the merged state.

        Raises:
            ClockSkewError: If the remote timestamp is too far in
                            the future (exceeds MAX_CLOCK_SKEW_MS).
        """
        physical = self._wall_clock()

        # Check for excessive clock skew
        if remote.wall_ms > physical + MAX_CLOCK_SKEW_MS:
            raise ClockSkewError(
                f"Remote timestamp {remote.wall_ms}ms is "
                f"{remote.wall_ms - physical}ms ahead of local clock "
                f"(max allowed: {MAX_CLOCK_SKEW_MS}ms). "
                f"Remote node: {remote.node_id}"
            )

        if physical > self._last.wall_ms and physical > remote.wall_ms:
            # Local wall clock is ahead of everything — reset counter
            ts = HLCTimestamp(wall_ms=physical, counter=0, node_id=self.node_id)
        elif self._last.wall_ms == remote.wall_ms:
            # Same wall time — advance counter past both
            ts = HLCTimestamp(
                wall_ms=self._last.wall_ms,
                counter=max(self._last.counter, remote.counter) + 1,
                node_id=self.node_id,
            )
        elif self._last.wall_ms > remote.wall_ms:
            # Local is ahead — just increment local counter
            ts = HLCTimestamp(
                wall_ms=self._last.wall_ms,
                counter=self._last.counter + 1,
                node_id=self.node_id,
            )
        else:
            # Remote is ahead — adopt remote wall time, increment counter
            ts = HLCTimestamp(
                wall_ms=remote.wall_ms,
                counter=remote.counter + 1,
                node_id=self.node_id,
            )

        self._last = ts
        return ts

    @property
    def last(self) -> HLCTimestamp:
        """The most recent timestamp issued by this clock."""
        return self._last


class ClockSkewError(Exception):
    """Raised when a remote timestamp exceeds the maximum allowed clock skew."""
    pass


class HLCWriter:
    """Manages HLC timestamp generation for database writes.

    A singleton-per-node that wraps an HLC clock and provides methods
    to generate timestamps for INSERT/UPDATE operations on synced tables.

    Usage:
        writer = HLCWriter.get_instance("node-1")
        hlc_str = writer.generate()  # For local writes
        hlc_str = writer.receive_and_generate(remote_hlc_str)  # During sync

    The generated string is stored in the _hlc_timestamp column.
    """

    _instances: dict[str, "HLCWriter"] = {}

    def __init__(self, node_id: str, wall_clock: callable | None = None):
        self._clock = HLC(node_id, wall_clock=wall_clock)
        self._node_id = node_id

    @classmethod
    def get_instance(cls, node_id: str) -> "HLCWriter":
        """Get or create a singleton HLCWriter for a given node_id."""
        if node_id not in cls._instances:
            cls._instances[node_id] = cls(node_id)
        return cls._instances[node_id]

    @classmethod
    def reset_instances(cls) -> None:
        """Clear all cached instances (for testing)."""
        cls._instances.clear()

    @property
    def node_id(self) -> str:
        return self._node_id

    @property
    def clock(self) -> HLC:
        return self._clock

    def generate(self) -> str:
        """Generate an HLC timestamp string for a local write event.

        Returns:
            String in format "wall_ms.counter.node_id"
        """
        ts = self._clock.now()
        return str(ts)

    def receive_and_generate(self, remote_hlc: str) -> str:
        """Receive a remote HLC timestamp and generate a new one that
        causally follows it.

        Used when applying a synced row from a remote node — the local
        timestamp must be > the remote one for causal ordering.

        Args:
            remote_hlc: String-encoded HLC timestamp from remote node.

        Returns:
            New HLC timestamp string that is causally after remote_hlc.

        Raises:
            ClockSkewError: If remote timestamp is too far in the future.
        """
        remote_ts = HLCTimestamp.from_str(remote_hlc)
        ts = self._clock.receive(remote_ts)
        return str(ts)

    @property
    def last(self) -> HLCTimestamp:
        """The most recent timestamp issued."""
        return self._clock.last

