"""Tests for key derivation."""

from x402_stacks.stacks_tx.keys import (
    private_key_to_public_key,
    hash160,
    public_key_to_address,
    private_key_to_address,
    get_address_version,
)

# Known test key
TEST_PRIVATE_KEY = "edf9aee84d9b7abc145504dde6726c64f369d37ee34ced614f76c2dcb7004497"


class TestKeyDerivation:
    def test_public_key_is_33_bytes(self):
        pub = private_key_to_public_key(TEST_PRIVATE_KEY)
        assert len(pub) == 33
        # Compressed key starts with 02 or 03
        assert pub[0] in (0x02, 0x03)

    def test_hash160_is_20_bytes(self):
        pub = private_key_to_public_key(TEST_PRIVATE_KEY)
        h = hash160(pub)
        assert len(h) == 20

    def test_testnet_address_prefix(self):
        addr = private_key_to_address(TEST_PRIVATE_KEY, "testnet")
        assert addr.startswith("ST")

    def test_mainnet_address_prefix(self):
        addr = private_key_to_address(TEST_PRIVATE_KEY, "mainnet")
        assert addr.startswith("SP")

    def test_deterministic(self):
        addr1 = private_key_to_address(TEST_PRIVATE_KEY, "testnet")
        addr2 = private_key_to_address(TEST_PRIVATE_KEY, "testnet")
        assert addr1 == addr2

    def test_hex_prefix_handling(self):
        addr1 = private_key_to_address(TEST_PRIVATE_KEY, "testnet")
        addr2 = private_key_to_address("0x" + TEST_PRIVATE_KEY, "testnet")
        assert addr1 == addr2

    def test_compression_flag_handling(self):
        addr1 = private_key_to_address(TEST_PRIVATE_KEY, "testnet")
        addr2 = private_key_to_address(TEST_PRIVATE_KEY + "01", "testnet")
        assert addr1 == addr2

    def test_address_version(self):
        assert get_address_version("mainnet") == 22
        assert get_address_version("testnet") == 26

    def test_public_key_to_address_roundtrip(self):
        pub = private_key_to_public_key(TEST_PRIVATE_KEY)
        addr = public_key_to_address(pub, "testnet")
        assert addr == private_key_to_address(TEST_PRIVATE_KEY, "testnet")
