"""NetDB lookup service — query peers for RouterInfo/LeaseSet."""

import enum


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


class LookupState(enum.Enum):
    NEW = 0
    QUERYING = 1
    FOUND = 2
    NOT_FOUND = 3
    FAILED = 4


class LookupRequest:
    """A single lookup request for a key in the NetDB."""

    def __init__(self, key: bytes, timeout_ms: int = 10000):
        self.key = key
        self.timeout_ms = timeout_ms
        self.state = LookupState.NEW
        self.result: bytes | None = None
        self.error: str | None = None
        self.queried_peers: set[bytes] = set()

    def start(self):
        self.state = LookupState.QUERYING

    def set_found(self, data: bytes):
        self.result = data
        self.state = LookupState.FOUND

    def set_not_found(self):
        self.state = LookupState.NOT_FOUND

    def set_failed(self, error: str):
        self.error = error
        self.state = LookupState.FAILED

    def add_queried_peer(self, peer_hash: bytes):
        self.queried_peers.add(peer_hash)

    def already_queried(self, peer_hash: bytes) -> bool:
        return peer_hash in self.queried_peers


class FloodfillSet:
    """Set of known floodfill routers."""

    def __init__(self):
        self._floodfills: dict[bytes, bytes | None] = {}

    def add(self, router_hash: bytes, router_info: bytes | None = None):
        self._floodfills[router_hash] = router_info

    def remove(self, router_hash: bytes):
        self._floodfills.pop(router_hash, None)

    def is_floodfill(self, router_hash: bytes) -> bool:
        return router_hash in self._floodfills

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

    def closest(self, target: bytes, n: int) -> list[bytes]:
        hashes = list(self._floodfills.keys())
        hashes.sort(key=lambda h: _xor_distance(h, target))
        return hashes[:n]