import csv, hashlib, random, sys, pathlib
ROOT = pathlib.Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from wp34 import secp256k1 as k1, sha256 as S, ecdsa as D, schnorr as B
from wp34.ec import Curve


def test_secp256k1_parameters():
    f = k1.verify_parameters()
    for key in ["p_is_prime", "n_is_prime", "G_on_curve", "nG_is_infinity", "beta_cubed_is_1_mod_p",
                "lambda_cubed_is_1_mod_n", "phi_G_equals_lambda_G", "glv_basis_v1", "glv_basis_v2",
                "glv_basis_det_is_n", "disc_is_minus3_times_square", "hasse_exact"]:
        assert f[key], key
    assert f["j_invariant"] == 0 and f["p_mod_3"] == 1 and not f["anomalous"]


def test_sha256_matches_hashlib():
    rng = random.Random(7)
    for L in [0, 1, 55, 56, 63, 64, 65, 119, 120, 300]:
        m = bytes(rng.getrandbits(8) for _ in range(L))
        assert S.sha256(m) == hashlib.sha256(m).digest()


def test_sha256_step_inverse():
    rng = random.Random(3)
    for _ in range(50):
        st = tuple(rng.getrandbits(32) for _ in range(8))
        k, w = rng.getrandbits(32), rng.getrandbits(32)
        assert S.step_inverse(S.step(st, k, w), k, w) == st


def test_ecdsa_against_python_ecdsa_and_coincurve():
    import ecdsa as pyecdsa, coincurve
    rng = random.Random(11)
    for _ in range(5):
        d = rng.randrange(1, k1.n)
        msg = bytes(rng.getrandbits(8) for _ in range(40))
        h = hashlib.sha256(msg).digest()
        r, s = D.sign(d, h, low_s=False)
        sk = pyecdsa.SigningKey.from_secret_exponent(d, curve=pyecdsa.SECP256k1, hashfunc=hashlib.sha256)
        sig_ref = sk.sign_digest_deterministic(h, hashfunc=hashlib.sha256, sigencode=pyecdsa.util.sigencode_strings)
        assert (int.from_bytes(sig_ref[0], "big"), int.from_bytes(sig_ref[1], "big")) == (r, s)
        Q = D.pubkey(d)
        assert D.verify(Q, h, (r, s)) and D.verify_shamir(Q, h, (r, s))
        assert not D.verify(Q, h, (r, (s + 1) % k1.n))
        # coincurve (libsecp256k1) verifies our low-s signature in DER
        rl, sl = D.sign(d, h)
        der = coincurve.utils.der_encode if hasattr(coincurve.utils, "der_encode") else None
        pk = coincurve.PublicKey.from_secret(d.to_bytes(32, "big"))
        from coincurve.ecdsa import cdata_to_der, der_to_cdata  # noqa
        import coincurve._libsecp256k1 as lib  # noqa
        # use compact 64-byte form via coincurve's verify with hasher=None on the raw hash
        sig64 = rl.to_bytes(32, "big") + sl.to_bytes(32, "big")
        from coincurve.ecdsa import deserialize_compact, cdata_to_der
        assert pk.verify(cdata_to_der(deserialize_compact(sig64)), h, hasher=None)


def test_bip340_vectors():
    rows = list(csv.DictReader(open(ROOT / "data/vectors/bip340_test_vectors.csv")))
    assert len(rows) == 19
    for row in rows:
        pk = bytes.fromhex(row["public key"]); msg = bytes.fromhex(row["message"]); sig = bytes.fromhex(row["signature"])
        expected = row["verification result"] == "TRUE"
        if row["secret key"]:
            sk = bytes.fromhex(row["secret key"]); aux = bytes.fromhex(row["aux_rand"])
            assert B.pubkey_gen(sk) == pk
            assert B.sign(msg, sk, aux) == sig
        assert B.verify(msg, pk, sig) == expected, row["index"]


def test_bip340_batch():
    rng = random.Random(5)
    items = []
    for _ in range(4):
        sk = rng.randrange(1, k1.n).to_bytes(32, "big"); aux = bytes(rng.getrandbits(8) for _ in range(32))
        msg = bytes(rng.getrandbits(8) for _ in range(32))
        items.append((msg, B.pubkey_gen(sk), B.sign(msg, sk, aux)))
    assert B.batch_verify(items, random.Random(1))
    bad = list(items); m, pk, sig = bad[2]; bad[2] = (m, pk, sig[:63] + bytes([sig[63] ^ 1]))
    assert not B.batch_verify(bad, random.Random(1))


def test_glv_decomposition_bounds():
    rng = random.Random(2)
    for _ in range(200):
        k = rng.randrange(1, k1.n)
        k1_, k2_ = k1.glv_decompose(k)
        assert (k1_ + k2_ * k1.lam - k) % k1.n == 0
        assert abs(k1_) < 2**129 and abs(k2_) < 2**129
    E = k1.CURVE
    Q = E.mul(k, E.G)
    phi = lambda P: (k1.beta * P[0] % k1.p, P[1])
    assert E.add(E.mul(k1_, E.G), E.mul(k2_, phi(E.G))) == Q


def test_small_curve_group_law():
    E = Curve(97, 2, 3)
    N = E.count_points_naive()
    rng = random.Random(1)
    for _ in range(20):
        P, Q = E.random_point(rng), E.random_point(rng)
        assert E.add(P, Q) == E.add(Q, P)
        assert E.mul(N, P) is None
        a, b = rng.randrange(1, 300), rng.randrange(1, 300)
        assert E.mul(a + b, P) == E.add(E.mul(a, P), E.mul(b, P))
        assert E.mul_wnaf(a, P) == E.mul(a, P)