"""Tests for electrum.btcq_keyderiv — BTCQ HD key derivation."""
import importlib.util
import os
import sys
import hmac
import hashlib

_ELECTRUM_DIR = os.path.join(os.path.dirname(__file__), '..', 'electrum')

def _load(name, filename):
    spec = importlib.util.spec_from_file_location(name, os.path.join(_ELECTRUM_DIR, filename))
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod

# Load in dependency order
_load('electrum.mldsa_keygen', 'mldsa_keygen.py')
_load('electrum.mldsa', 'mldsa.py')
btcq_crypto = _load('electrum.btcq_crypto', 'btcq_crypto.py')
keyderiv_mod = _load('electrum.btcq_keyderiv', 'btcq_keyderiv.py')
BtcqKeyDerivation = keyderiv_mod.BtcqKeyDerivation


def test_from_seed_deterministic():
    seed = bytes(range(64))
    k1 = BtcqKeyDerivation.from_seed(seed)
    k2 = BtcqKeyDerivation.from_seed(seed)
    assert k1.public_key == k2.public_key
    assert k1.secret_key == k2.secret_key
    assert k1.chaincode == k2.chaincode
    assert k1.depth == 0


def test_from_seed_sizes():
    seed = bytes(range(64))
    k = BtcqKeyDerivation.from_seed(seed)
    assert len(k.public_key) == 1952
    assert len(k.secret_key) == 4032
    assert len(k.chaincode) == 32


def test_child_derivation_different_keys():
    seed = bytes(range(64))
    master = BtcqKeyDerivation.from_seed(seed)
    child = master.derive_child(0)
    assert child.public_key != master.public_key
    assert child.secret_key != master.secret_key
    assert child.depth == 1


def test_different_indices_different_children():
    seed = bytes(range(64))
    master = BtcqKeyDerivation.from_seed(seed)
    c0 = master.derive_child(0)
    c1 = master.derive_child(1)
    assert c0.public_key != c1.public_key
    assert c0.secret_key != c1.secret_key
    assert c0.chaincode != c1.chaincode


def test_derivation_uses_rho():
    """Verify derivation uses rho (sk[0:32]) by manually computing HMAC."""
    from electrum.mldsa import mldsa_keygen_from_seed
    seed = bytes(range(64))
    master = BtcqKeyDerivation.from_seed(seed)
    rho = master.secret_key[:32]
    index = 42
    data = b'\x00' + rho + index.to_bytes(4, 'big')
    hmac_out = hmac.new(master.chaincode, data, hashlib.sha512).digest()
    expected_xi = hmac_out[:32]
    expected_cc = hmac_out[32:]
    expected_pk, expected_sk = mldsa_keygen_from_seed(expected_xi)
    child = master.derive_child(index)
    assert child.public_key == expected_pk
    assert child.secret_key == expected_sk
    assert child.chaincode == expected_cc
    assert child.seed_xi == expected_xi


def test_five_level_derivation():
    seed = bytes(range(64))
    key = BtcqKeyDerivation.from_seed(seed)
    for depth in range(5):
        key = key.derive_child(depth)
        assert key.depth == depth + 1
        assert len(key.public_key) == 1952
        assert len(key.secret_key) == 4032
    assert key.depth == 5


def test_child_fingerprint():
    seed = bytes(range(64))
    master = BtcqKeyDerivation.from_seed(seed)
    child = master.derive_child(0)
    expected_fp = btcq_crypto.hash_160_btcq(master.public_key)[:4]
    assert child.parent_fingerprint == expected_fp


if __name__ == "__main__":
    tests = [f for f in dir() if f.startswith('test_')]
    for t in sorted(tests):
        print(f"  {t} ... ", end="", flush=True)
        globals()[t]()
        print("OK")
    print(f"\nAll {len(tests)} btcq_keyderiv tests passed.")
