"""NetDB data store — in-memory storage for RouterInfo and LeaseSet entries."""

import enum
import time


class EntryType(enum.IntEnum):
    ROUTER_INFO = 0
    LEASE_SET = 1


class NetDBEntry:
    """A single entry in the network database."""

    def __init__(self, key: bytes, entry_type: EntryType, data: bytes,
                 received_ms: int, expiration_ms: int | None = None):
        self.key = key
        self.entry_type = entry_type
        self.data = data
        self.received_ms = received_ms
        self.expiration_ms = expiration_ms

    def is_expired(self, now_ms: int | None = None) -> bool:
        if self.expiration_ms is None:
            return False
        if now_ms is None:
            now_ms = int(time.time() * 1000)
        return now_ms >= self.expiration_ms


def _xor_distance(a: bytes, b: bytes) -> bytes:
    return bytes(x ^ y for x, y in zip(a, b))


class DataStore:
    """In-memory key-value store for NetDB entries."""

    def __init__(self):
        self._entries: dict[bytes, NetDBEntry] = {}

    def put(self, entry: NetDBEntry):
        self._entries[entry.key] = entry

    def get(self, key: bytes) -> NetDBEntry | None:
        return self._entries.get(key)

    def remove(self, key: bytes):
        self._entries.pop(key, None)

    def count(self) -> int:
        return len(self._entries)

    def count_by_type(self, entry_type: EntryType) -> int:
        return sum(1 for e in self._entries.values() if e.entry_type == entry_type)

    def remove_expired(self, now_ms: int | None = None):
        expired = [k for k, e in self._entries.items() if e.is_expired(now_ms)]
        for k in expired:
            del self._entries[k]

    def get_all(self) -> list[NetDBEntry]:
        """Return all entries."""
        return list(self._entries.values())

    def get_all_keys(self) -> list[bytes]:
        return list(self._entries.keys())

    def closest_keys(self, target: bytes, n: int) -> list[bytes]:
        keys = list(self._entries.keys())
        keys.sort(key=lambda k: _xor_distance(k, target))
        return keys[:n]