"""Tests for NetDB data store."""

import os
import time

import pytest


class TestNetDBEntry:
    def test_construct_router_info(self):
        from i2p_netdb.datastore import NetDBEntry, EntryType
        data = os.urandom(500)
        entry = NetDBEntry(key=os.urandom(32), entry_type=EntryType.ROUTER_INFO,
                          data=data, received_ms=1000)
        assert entry.entry_type == EntryType.ROUTER_INFO
        assert entry.data == data

    def test_construct_lease_set(self):
        from i2p_netdb.datastore import NetDBEntry, EntryType
        entry = NetDBEntry(key=os.urandom(32), entry_type=EntryType.LEASE_SET,
                          data=os.urandom(200), received_ms=1000)
        assert entry.entry_type == EntryType.LEASE_SET

    def test_is_expired(self):
        from i2p_netdb.datastore import NetDBEntry, EntryType
        entry = NetDBEntry(key=os.urandom(32), entry_type=EntryType.ROUTER_INFO,
                          data=b"", received_ms=1000, expiration_ms=5000)
        assert entry.is_expired(now_ms=6000)
        assert not entry.is_expired(now_ms=4000)


class TestDataStore:
    def test_put_and_get(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        key = os.urandom(32)
        entry = NetDBEntry(key=key, entry_type=EntryType.ROUTER_INFO,
                          data=b"router data", received_ms=1000)
        store.put(entry)
        result = store.get(key)
        assert result is not None
        assert result.data == b"router data"

    def test_get_missing(self):
        from i2p_netdb.datastore import DataStore
        store = DataStore()
        assert store.get(os.urandom(32)) is None

    def test_remove(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        key = os.urandom(32)
        store.put(NetDBEntry(key, EntryType.ROUTER_INFO, b"data", 1000))
        store.remove(key)
        assert store.get(key) is None

    def test_count(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        for i in range(5):
            store.put(NetDBEntry(os.urandom(32), EntryType.ROUTER_INFO, b"", 1000))
        assert store.count() == 5

    def test_count_by_type(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        for _ in range(3):
            store.put(NetDBEntry(os.urandom(32), EntryType.ROUTER_INFO, b"", 1000))
        for _ in range(2):
            store.put(NetDBEntry(os.urandom(32), EntryType.LEASE_SET, b"", 1000))
        assert store.count_by_type(EntryType.ROUTER_INFO) == 3
        assert store.count_by_type(EntryType.LEASE_SET) == 2

    def test_remove_expired(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        store.put(NetDBEntry(os.urandom(32), EntryType.ROUTER_INFO, b"", 1000, 500))
        store.put(NetDBEntry(os.urandom(32), EntryType.ROUTER_INFO, b"", 1000, 9999999999))
        store.remove_expired(now_ms=1000)
        assert store.count() == 1

    def test_get_all_keys(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        keys = [os.urandom(32) for _ in range(3)]
        for k in keys:
            store.put(NetDBEntry(k, EntryType.ROUTER_INFO, b"", 1000))
        all_keys = store.get_all_keys()
        assert len(all_keys) == 3
        for k in keys:
            assert k in all_keys

    def test_closest_keys(self):
        from i2p_netdb.datastore import DataStore, NetDBEntry, EntryType
        store = DataStore()
        for _ in range(10):
            store.put(NetDBEntry(os.urandom(32), EntryType.ROUTER_INFO, b"", 1000))
        target = os.urandom(32)
        closest = store.closest_keys(target, n=3)
        assert len(closest) == 3