"""Snowflake IDs (架构文档 §6.1 原则 2).

Auto-increment keys invite enumeration attacks and collide when the database is
later sharded; UUID primary keys have poor index locality. A 64-bit snowflake
avoids both.

Layout: 41 bits ms since epoch | 10 bits node id | 12 bits sequence.
"""

from __future__ import annotations

import threading
import time
from collections.abc import Callable

#: 2025-01-01T00:00:00Z in milliseconds. Do not change once IDs are issued.
DEFAULT_EPOCH_MS = 1_735_689_600_000

_NODE_BITS = 10
_SEQUENCE_BITS = 12
_MAX_NODE_ID = (1 << _NODE_BITS) - 1
_SEQUENCE_MASK = (1 << _SEQUENCE_BITS) - 1


def _system_clock_ms() -> int:
    return int(time.time() * 1000)


class SnowflakeGenerator:
    """Thread-safe monotonic 64-bit id generator."""

    __slots__ = ("_clock", "_epoch_ms", "_last_ms", "_lock", "_node_id", "_sequence")

    def __init__(
        self,
        node_id: int = 1,
        epoch_ms: int = DEFAULT_EPOCH_MS,
        clock: Callable[[], int] | None = None,
    ) -> None:
        if not 0 <= node_id <= _MAX_NODE_ID:
            msg = f"node_id must be within 0..{_MAX_NODE_ID}"
            raise ValueError(msg)
        self._node_id = node_id
        self._epoch_ms = epoch_ms
        self._sequence = 0
        self._last_ms = -1
        self._lock = threading.Lock()
        # 时钟可注入：序号回绕与时钟回拨这两条分支只能靠假时钟覆盖，
        # 真等一毫秒的测试既慢又不确定。
        self._clock = clock or _system_clock_ms

    def _now_ms(self) -> int:
        return self._clock()

    def next_id(self) -> int:
        with self._lock:
            now = self._now_ms()
            if now < self._last_ms:
                # Clock moved backwards: wait it out rather than issue duplicates.
                now = self._last_ms
            if now == self._last_ms:
                self._sequence = (self._sequence + 1) & _SEQUENCE_MASK
                if self._sequence == 0:
                    while now <= self._last_ms:
                        now = self._now_ms()
            else:
                self._sequence = 0
            self._last_ms = now
            elapsed = now - self._epoch_ms
            return (
                (elapsed << (_NODE_BITS + _SEQUENCE_BITS))
                | (self._node_id << _SEQUENCE_BITS)
                | self._sequence
            )


_default_generator = SnowflakeGenerator()


def next_id() -> int:
    """Process-wide default generator."""
    return _default_generator.next_id()


def timestamp_ms_of(snowflake_id: int, epoch_ms: int = DEFAULT_EPOCH_MS) -> int:
    """Extract the creation timestamp from an id (useful for partition pruning)."""
    return (snowflake_id >> (_NODE_BITS + _SEQUENCE_BITS)) + epoch_ms