"""
aaiclick.snowflake.snowflake_id - Snowflake ID generation.

Globally unique, time-ordered 64-bit identifiers in the Twitter Snowflake
layout. IDs are pre-fetched in batches for efficiency, served one at a time
from an in-memory buffer until empty, then refilled.

In chdb mode, IDs are minted in pure Python from a process-wide
(timestamp, sequence) pair. This deliberately avoids touching the shared
chdb Session: callers running under ``orch_context(with_ch=False)`` must be
able to mint IDs without acquiring the chdb file lock.

In distributed mode, IDs are fetched from ClickHouse via clickhouse-connect
(``generateSnowflakeID()``).

Snowflake ID format (64 bits):
- Bit 63: Sign bit (always 0 for positive integers)
- Bits 62-22: Timestamp in milliseconds (41 bits)
- Bits 21-12: Machine/worker ID (10 bits)
- Bits 11-0: Sequence number (12 bits)
"""

import os
import threading
import time
from collections import deque

from ..backend import is_chdb, parse_ch_url

# Bit allocation (Wikipedia Snowflake ID standard)
MACHINE_ID_BITS = 10  # Bits 21-12: supports 1024 machines
SEQUENCE_BITS = 12  # Bits 11-0: supports 4096 IDs per millisecond

# Maximum values
MAX_SEQUENCE = (1 << SEQUENCE_BITS) - 1  # 4095

# Bit shifts for decoding the 64-bit ID
TIMESTAMP_SHIFT = MACHINE_ID_BITS + SEQUENCE_BITS  # 22
MACHINE_ID_SHIFT = SEQUENCE_BITS  # 12

# Default batch size for pre-fetching IDs from ClickHouse
_BUFFER_SIZE = 100

# Process-wide state for in-memory generation (chdb mode). Shared across
# every SnowflakeGenerator instance and every thread so the (timestamp,
# sequence) pair stays monotonic for the lifetime of the process.
_IN_MEMORY_MACHINE_ID = os.getpid() & ((1 << MACHINE_ID_BITS) - 1)
_in_memory_lock = threading.Lock()
_in_memory_last_ms = 0
_in_memory_sequence = -1


def _generate_in_memory(count: int) -> list[int]:
    """Mint ``count`` Snowflake IDs in pure Python.

    Used in chdb mode to avoid the shared chdb Session — mint ID flows
    must work even when ``orch_context(with_ch=False)`` has suppressed
    the ClickHouse client.
    """
    global _in_memory_last_ms, _in_memory_sequence
    out: list[int] = []
    with _in_memory_lock:
        for _ in range(count):
            now_ms = int(time.time() * 1000)
            if now_ms > _in_memory_last_ms:
                _in_memory_last_ms = now_ms
                _in_memory_sequence = 0
            else:
                _in_memory_sequence += 1
                if _in_memory_sequence > MAX_SEQUENCE:
                    # Sequence overflow inside one ms — borrow from next ms.
                    _in_memory_last_ms += 1
                    _in_memory_sequence = 0
            out.append(
                (_in_memory_last_ms << TIMESTAMP_SHIFT)
                | (_IN_MEMORY_MACHINE_ID << MACHINE_ID_SHIFT)
                | _in_memory_sequence
            )
    return out


class SnowflakeGenerator:
    """Snowflake ID generator.

    In chdb mode, IDs come from a process-wide in-memory counter.
    In distributed mode, IDs are pre-fetched from ClickHouse's
    ``generateSnowflakeID()`` in batches and served from an in-memory
    buffer; the buffer refills automatically when exhausted.
    """

    def __init__(self, buffer_size: int = _BUFFER_SIZE):
        self._buffer_size = buffer_size
        self._buffer: deque[int] = deque()

    def _fetch_ids(self, count: int) -> list[int]:
        """Fetch a batch of Snowflake IDs."""
        if is_chdb():
            return _generate_in_memory(count)
        return self._fetch_ids_remote(count)

    @staticmethod
    def _fetch_ids_remote(count: int) -> list[int]:
        from clickhouse_connect import get_client

        client = get_client(**parse_ch_url())
        try:
            result = client.query(f"SELECT generateSnowflakeID() FROM numbers({count})")
            return [row[0] for row in result.result_rows]
        finally:
            client.close()

    def generate(self) -> int:
        """Generate a single Snowflake ID."""
        if not self._buffer:
            self._buffer.extend(self._fetch_ids(self._buffer_size))
        return self._buffer.popleft()

    def generate_bulk(self, count: int) -> list[int]:
        """Generate multiple Snowflake IDs.

        Args:
            count: Number of IDs to generate (must be >= 1)

        Returns:
            list[int]: List of unique Snowflake IDs in ascending order
        """
        if count < 1:
            raise ValueError(f"Count must be at least 1, got {count}")
        return self.get(count)

    def get(self, size: int) -> list[int]:
        """Get multiple Snowflake IDs.

        Args:
            size: Number of IDs to generate (must be >= 0)

        Returns:
            list[int]: List of unique Snowflake IDs in ascending order
        """
        if size < 0:
            raise ValueError(f"Size must be >= 0, got {size}")
        if size == 0:
            return []

        if size <= len(self._buffer):
            return [self._buffer.popleft() for _ in range(size)]

        # Drain buffer, fetch remaining + refill from CH
        result = list(self._buffer)
        self._buffer.clear()
        remaining = size - len(result)
        fetched = self._fetch_ids(remaining + self._buffer_size)
        result.extend(fetched[:remaining])
        self._buffer.extend(fetched[remaining:])
        return result


# Global generator instance (lazy CH connection on first use)
_generator = SnowflakeGenerator()


def get_snowflake_id() -> int:
    """Get a single Snowflake ID.

    Returns:
        int: Unique 64-bit Snowflake ID
    """
    return _generator.generate()


def get_snowflake_ids(size: int) -> list[int]:
    """Get multiple Snowflake IDs.

    Args:
        size: Number of IDs to generate (must be >= 0)

    Returns:
        list[int]: List of unique 64-bit Snowflake IDs in ascending order
    """
    return _generator.get(size)


def decode_snowflake_id(id_val: int) -> tuple[int, int, int]:
    """Decode a Snowflake ID into its component parts.

    Args:
        id_val: 64-bit Snowflake ID to decode

    Returns:
        tuple: (timestamp, machine_id, sequence)
            - timestamp: Milliseconds since epoch (bits 62-22)
            - machine_id: Machine/worker ID (bits 21-12)
            - sequence: Sequence number (bits 11-0)

    Example:
        >>> id_val = get_snowflake_id()
        >>> timestamp, machine_id, sequence = decode_snowflake_id(id_val)
    """
    timestamp = id_val >> TIMESTAMP_SHIFT
    machine_id = (id_val >> MACHINE_ID_SHIFT) & ((1 << MACHINE_ID_BITS) - 1)
    sequence = id_val & MAX_SEQUENCE
    return timestamp, machine_id, sequence
