import random, sys, pathlib
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1] / "src"))
from wp34.ec import Curve, random_prime_order_curve
from wp34 import dlp


def _instance(bits, seed, j0=False):
    rng = random.Random(seed)
    E = random_prime_order_curve(bits, rng, j0=j0)
    d = rng.randrange(1, E.n)
    return E, d, E.mul(d, E.G)


def test_bsgs_rho_kangaroo_agree():
    for seed in range(3):
        E, d, Q = _instance(22, seed)
        assert dlp.bsgs(E, E.G, Q)[0] == d
        assert dlp.pollard_rho(E, E.G, Q, seed=seed)[0] == d
        assert dlp.pollard_rho_dp(E, E.G, Q, dp_bits=3, walks=4, seed=seed)[0] == d
        assert dlp.pollard_rho_dp(E, E.G, Q, dp_bits=3, walks=4, seed=seed, neg_map=True)[0] == d
        assert dlp.kangaroo(E, E.G, Q, 0, E.n - 1, seed=seed)[0] == d
        lo = max(0, d - 2**12); hi = lo + 2**13
        assert dlp.kangaroo(E, E.G, Q, lo, hi, seed=seed)[0] == d


def test_bsgs_memory_tradeoff():
    E, d, Q = _instance(20, 9)
    for m in [2**6, 2**8, 2**10, 2**12]:
        dd, st = dlp.bsgs(E, E.G, Q, m=m)
        assert dd == d and st["memory_entries"] == m


def test_pohlig_hellman_small():
    E = Curve(10007, 1, 6)
    N = E.count_points_naive()
    rng = random.Random(4)
    P = E.random_point(rng)
    order = E.order_of_point(P, bound=N + 1)
    d = rng.randrange(1, order)
    Q = E.mul(d, P)
    assert dlp.pohlig_hellman(E, P, Q, order)[0] == d