# -*- coding: utf-8 -*-
"""Regression tests for BTCQ extended-key derivation (electrum.btcq_extkey).

Replaces the old tests/test_btcq_bip32.py BIP32Node-shaped tests after
the rewrite that moved derivation logic out of the secp256k1-typed
BIP32Node shim into the BTCQ-native :class:`BtcqExtKey`.

The wallet wizard's create-new-wallet flow calls:
  keystore.from_seed(seed, ...)
  -> keystore.add_xprv_from_seed(bip32_seed, xtype, derivation)
  -> BtcqExtKey.from_rootseed(bip32_seed, xtype)
  -> rootnode.subkey_at_private_derivation(derivation)
  -> per-step _CKD_priv (rho extraction, HMAC-SHA512, mldsa_keygen)
  -> node.to_xprv() and node.to_xpub() for storage and master-pubkey display

ML-DSA-65 keygen is ~1-2s per call; one derivation step does multiple
keygens (parent rho + child pk + xpub generation if requested), so these
tests run on the order of seconds, not milliseconds.
"""
import unittest

from electrum.btcq_extkey import (
    BtcqExtKey,
    is_xpub,
    is_xprv,
    xpub_from_xprv,
    xpub_type,
    root_fp_and_der_prefix_from_xkey,
    xprv_header,
    xpub_header,
    XPRV_PAYLOAD_SIZE,
    XPUB_PAYLOAD_SIZE,
    MLDSA65_PK_SIZE,
    MLDSA65_XI_SIZE,
)
from electrum.keystore import BIP32_KeyStore, CannotDerivePubkey


# 64-byte BIP32 seed (deterministic test vector — first 64 bytes of "BTCQ" repeated).
TEST_BIP32_SEED = (b"BTCQ" * 16)[:64]


class BtcqExtKeyDerivation(unittest.TestCase):
    """End-to-end derivation matching what the wizard does on
    create-new-wallet. Exercises rho-extraction, HMAC-SHA512, and
    ML-DSA-65 keygen at each step."""

    def test_from_rootseed_populates_xi_and_pk(self):
        root = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        self.assertEqual(root.depth, 0)
        self.assertIsNotNone(root.xi)
        self.assertEqual(len(root.xi), MLDSA65_XI_SIZE)
        self.assertEqual(len(root.pk), MLDSA65_PK_SIZE)
        self.assertEqual(len(root.chaincode), 32)
        self.assertEqual(root.fingerprint, b"\x00" * 4)
        self.assertTrue(root.is_private())

    def test_one_hardened_step_updates_fingerprint(self):
        root = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        node = root.subkey_at_private_derivation([0x80000000])
        self.assertEqual(node.depth, 1)
        self.assertEqual(len(node.fingerprint), 4)
        self.assertNotEqual(node.fingerprint, b"\x00" * 4,
                            "fingerprint should differ from default after a derivation step")
        self.assertTrue(node.is_private())

    def test_derivation_is_deterministic(self):
        a = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard").subkey_at_private_derivation([0x80000000])
        b = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard").subkey_at_private_derivation([0x80000000])
        self.assertEqual(a.xi, b.xi)
        self.assertEqual(a.pk, b.pk)
        self.assertEqual(a.chaincode, b.chaincode)


class BtcqExtKeySerialization(unittest.TestCase):
    """xprv (78 B) and xpub (1997 B) round-trips."""

    def setUp(self):
        self.root = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        self.node = self.root.subkey_at_private_derivation([0x80000000])

    def test_xprv_size_matches_BIP32_shape(self):
        xprv_bytes = self.node.to_xprv_bytes()
        self.assertEqual(len(xprv_bytes), XPRV_PAYLOAD_SIZE)
        self.assertEqual(len(xprv_bytes), 78)

    def test_xpub_size_carries_full_mldsa_pk(self):
        xpub_bytes = self.node.to_xpub_bytes()
        self.assertEqual(len(xpub_bytes), XPUB_PAYLOAD_SIZE)
        self.assertEqual(len(xpub_bytes), 1997,
                         "BTCQ xpub must hold the full 1952-byte ML-DSA pk")

    def test_xprv_round_trip_preserves_xi_pk_chaincode(self):
        # The latent crash this test prevents: BIP32Node.from_xkey on a
        # 78-byte xprv would call ECPrivkey(32 bytes), which works; but
        # any post-parse path that reaches into eckey for pk would
        # silently produce 33 bytes of secp256k1 garbage. With BtcqExtKey,
        # parsing eagerly recomputes pk via ML-DSA keygen from xi.
        xprv = self.node.to_xprv()
        parsed = BtcqExtKey.from_xkey(xprv)
        self.assertEqual(parsed.xi, self.node.xi)
        self.assertEqual(parsed.pk, self.node.pk)
        self.assertEqual(parsed.chaincode, self.node.chaincode)
        self.assertTrue(parsed.is_private())

    def test_xpub_round_trip_preserves_pk_chaincode(self):
        # The latent crash this test prevents: BIP32Node.from_xkey on a
        # 1997-byte xpub would call ECPubkey(1952 bytes) → InvalidECPointException.
        # With BtcqExtKey, the 1952-byte pk loads cleanly into the pk slot.
        xpub = self.node.to_xpub()
        parsed = BtcqExtKey.from_xkey(xpub)
        self.assertIsNone(parsed.xi, "xpub-parsed node should have xi=None")
        self.assertEqual(parsed.pk, self.node.pk)
        self.assertEqual(parsed.chaincode, self.node.chaincode)
        self.assertFalse(parsed.is_private())

    def test_to_xkey_branches_correctly(self):
        # Private node serializes as xprv; public-only node as xpub.
        self.assertEqual(self.node.to_xkey(), self.node.to_xprv())
        pub = self.node.convert_to_public()
        self.assertEqual(pub.to_xkey(), pub.to_xpub())

    def test_convert_to_public_strips_xi(self):
        pub = self.node.convert_to_public()
        self.assertIsNone(pub.xi)
        self.assertEqual(pub.pk, self.node.pk, "pk must survive conversion")
        self.assertFalse(pub.is_private())

    def test_fingerprint_stable_across_xpub_round_trip(self):
        xpub = self.node.to_xpub()
        parsed = BtcqExtKey.from_xkey(xpub)
        self.assertEqual(parsed.calc_fingerprint_of_this_node(),
                         self.node.calc_fingerprint_of_this_node())


class BtcqExtKeyPublicDerivationRaises(unittest.TestCase):
    """ML-DSA-65 has no group homomorphism — public-key derivation
    must raise loudly, not silently produce wrong bytes (which is what
    the secp256k1-shim code path did historically)."""

    def test_subkey_at_public_derivation_raises_on_hardened_path(self):
        node = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        with self.assertRaises(NotImplementedError) as ctx:
            node.subkey_at_public_derivation([0x80000000])
        self.assertIn("ML-DSA-65", str(ctx.exception))

    def test_subkey_at_public_derivation_raises_on_unhardened_path(self):
        node = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        with self.assertRaises(NotImplementedError):
            node.subkey_at_public_derivation([0])

    def test_empty_path_returns_public_copy(self):
        # Edge case: empty path should still be a no-op (matches upstream
        # semantics of "convert_to_public on empty path").
        node = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        result = node.subkey_at_public_derivation([])
        self.assertIsNone(result.xi)
        self.assertEqual(result.pk, node.pk)


class XprvXpubHeaderResolution(unittest.TestCase):
    """The wizard defaults new wallets to xtype='p2wpkh' (modern segwit
    seed type). Upstream Electrum has separate header bytes per xtype
    (xprv/yprv/zprv/...). BTCQ has one signature scheme (ML-DSA-65) and
    one address format (bech32), so all xtypes must collapse to 'standard'.
    Without this fallback xprv_header crashes with `KeyError: 'p2wpkh'`
    immediately after the user types their seed in the wizard."""

    def test_all_common_xtypes_resolve_to_standard_header(self):
        from electrum import constants
        net = constants.BtcqTestnet
        standard = net.XPRV_HEADERS['standard'].to_bytes(4, 'big')
        for xtype in ('standard', 'p2pkh', 'p2sh', 'p2wpkh', 'p2wpkh-p2sh', 'p2wsh', 'p2wsh-p2sh'):
            self.assertEqual(
                xprv_header(xtype, net=net),
                standard,
                f"xprv_header({xtype!r}) should fall back to 'standard' on BTCQ",
            )
            self.assertEqual(
                xpub_header(xtype, net=net),
                net.XPUB_HEADERS['standard'].to_bytes(4, 'big'),
                f"xpub_header({xtype!r}) should fall back to 'standard' on BTCQ",
            )


class XkeyHelperFunctions(unittest.TestCase):
    """Smoke tests for the helper functions (is_xpub, is_xprv, etc.)
    that are imported across the codebase from btcq_extkey."""

    def setUp(self):
        root = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        self.node = root.subkey_at_private_derivation([0x80000000])
        self.xprv = self.node.to_xprv()
        self.xpub = self.node.to_xpub()

    def test_is_xpub_classifies_correctly(self):
        self.assertTrue(is_xpub(self.xpub))
        self.assertFalse(is_xpub(self.xprv))
        self.assertFalse(is_xpub("not-an-xkey"))

    def test_is_xprv_classifies_correctly(self):
        self.assertTrue(is_xprv(self.xprv))
        self.assertFalse(is_xprv(self.xpub))
        self.assertFalse(is_xprv("not-an-xkey"))

    def test_xpub_from_xprv_matches_direct_serialize(self):
        self.assertEqual(xpub_from_xprv(self.xprv), self.xpub)

    def test_xpub_type_returns_standard(self):
        self.assertEqual(xpub_type(self.xpub), "standard")
        self.assertEqual(xpub_type(self.xprv), "standard")


class Bip32KeystoreDerivePubkey(unittest.TestCase):
    """Regression: wallet creation succeeded but address generation
    failed because Xpub._derive_pubkey calls subkey_at_public_derivation,
    which raises NotImplementedError on BTCQ (ML-DSA-65 has no
    public-key derivation). BIP32_KeyStore overrides _derive_pubkey to
    walk privately and return the leaf node's pk directly."""

    def _build_keystore(self):
        rootnode = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        accountnode = rootnode.subkey_at_private_derivation([0x80000000])
        xprv = accountnode.to_xprv()
        return BIP32_KeyStore({
            'xpub': accountnode.to_xpub(),
            'xprv': xprv,
            'derivation': "m/0'",
            'root_fingerprint': accountnode.fingerprint.hex(),
        })

    def test_derive_pubkey_returns_mldsa_pk(self):
        ks = self._build_keystore()
        pk = ks._derive_pubkey(for_change=0, n=0)
        self.assertEqual(len(pk), MLDSA65_PK_SIZE,
                         f"expected 1952-byte ML-DSA pk, got {len(pk)}")

    def test_derive_pubkey_change_and_index_distinct(self):
        ks = self._build_keystore()
        receive_0 = ks._derive_pubkey(for_change=0, n=0)
        receive_1 = ks._derive_pubkey(for_change=0, n=1)
        change_0 = ks._derive_pubkey(for_change=1, n=0)
        self.assertNotEqual(receive_0, receive_1)
        self.assertNotEqual(receive_0, change_0)
        self.assertNotEqual(receive_1, change_0)

    def test_derive_pubkey_deterministic(self):
        ks_a = self._build_keystore()
        ks_b = self._build_keystore()
        self.assertEqual(ks_a._derive_pubkey(0, 0), ks_b._derive_pubkey(0, 0))

    def test_derive_pubkey_rejects_forbidden_branch(self):
        ks = self._build_keystore()
        with self.assertRaises(CannotDerivePubkey):
            ks._derive_pubkey(for_change=2, n=0)


class DescriptorPubkeyExpansion(unittest.TestCase):
    """Regression: clicking Details on an address (which goes through
    wallet.get_redeem_script → descriptor.expand → PubkeyProvider.get_pubkey_bytes)
    crashed because the descriptor used subkey_at_public_derivation —
    same kind of bug as keystore._derive_pubkey, fixed the same way."""

    def _build_pubkey_provider(self, deriv_path):
        from electrum.descriptor import PubkeyProvider
        rootnode = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        return PubkeyProvider(origin=None, pubkey=rootnode.to_xprv(),
                              deriv_path=deriv_path)

    def test_get_pubkey_bytes_with_derivation_returns_mldsa_pk(self):
        pp = self._build_pubkey_provider("/0/0")
        pk = pp.get_pubkey_bytes()
        self.assertEqual(len(pk), MLDSA65_PK_SIZE,
                         f"expected 1952-byte ML-DSA pk, got {len(pk)}")

    def test_get_pubkey_bytes_no_derivation_private_node(self):
        # No derivation path on a private-bearing node — must return the
        # node's pk directly, NOT call get_public_key_bytes() (which
        # would silently produce a 33-byte secp256k1 pubkey from xi in
        # the pre-rewrite code).
        pp = self._build_pubkey_provider(None)
        pk = pp.get_pubkey_bytes()
        self.assertEqual(len(pk), MLDSA65_PK_SIZE)


class RootFpAndDerPrefix(unittest.TestCase):
    """root_fp_and_der_prefix_from_xkey is called by keystore.add_xpub
    on every wallet load. Cover the depth-0 and depth-1 cases."""

    def test_depth_0_returns_self_fingerprint_and_m(self):
        root = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        xpub = root.to_xpub()
        fp, prefix = root_fp_and_der_prefix_from_xkey(xpub)
        self.assertEqual(prefix, "m")
        self.assertEqual(fp, root.calc_fingerprint_of_this_node().hex().lower())

    def test_depth_1_returns_parent_fingerprint_and_one_step_path(self):
        root = BtcqExtKey.from_rootseed(TEST_BIP32_SEED, xtype="standard")
        node = root.subkey_at_private_derivation([0x80000000])
        xpub = node.to_xpub()
        fp, prefix = root_fp_and_der_prefix_from_xkey(xpub)
        self.assertEqual(fp, node.fingerprint.hex())
        # 0x80000000 is hardened index 0; intpath_to_strpath emits "0h".
        self.assertEqual(prefix, "m/0h")


class WalletGetRedeemScript(unittest.IsolatedAsyncioTestCase):
    """Regression: GUI's right-click → Details on an address triggers
    ``wallet.get_redeem_script(addr)`` → ``descriptor.expand()`` →
    ``PubkeyProvider.get_pubkey_bytes(pos=...)``. Pre-fix, this raised
    ``NotImplementedError("BTCQ cannot derive pubkeys from xpub alone")``
    because ``Xpub.get_pubkey_provider`` handed the xpub-string to
    PubkeyProvider, whose internal extkey ended up public-only.

    The fix: ``BIP32_KeyStore.get_pubkey_provider`` overrides the base
    Xpub method to ALSO pass an xprv-loaded BtcqExtKey via the new
    ``extkey`` kwarg of PubkeyProvider. The xpub-as-string remains in
    PubkeyProvider.pubkey for safe ``to_string()`` serialization.
    """

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        from electrum import constants
        cls._prev_net = constants.net
        constants.BtcqTestnet.set_as_network()

    @classmethod
    def tearDownClass(cls):
        from electrum import constants
        cls._prev_net.set_as_network()
        super().tearDownClass()

    def setUp(self):
        import tempfile
        from electrum import util
        super().setUp()
        self.tmpdir = tempfile.mkdtemp(prefix="btcq-extkey-test-")
        assert util._asyncio_event_loop is None, "global event loop already set"

    async def asyncSetUp(self):
        from electrum import util
        await super().asyncSetUp()
        util._asyncio_event_loop = util.get_asyncio_loop()

    async def asyncTearDown(self):
        await super().asyncTearDown()

    def tearDown(self):
        import shutil
        from electrum import util
        shutil.rmtree(self.tmpdir, ignore_errors=True)
        util._asyncio_event_loop = None
        super().tearDown()

    def _build_wallet(self):
        import os
        from electrum.simple_config import SimpleConfig
        from electrum.wallet import Standard_Wallet
        from electrum.wallet_db import WalletDB
        from electrum.storage import WalletStorage
        from electrum.keystore import from_seed

        config = SimpleConfig({})
        seed = ('square blur patch tail unveil road bridge digital require '
                'pumpkin genuine subway')
        ks = from_seed(seed, passphrase='', for_multisig=False)

        storage = WalletStorage(os.path.join(self.tmpdir, 'wallet'))
        db = WalletDB(s='', storage=storage, upgrade=True)
        db.put('keystore', ks.dump())
        db.put('wallet_type', 'standard')
        db.put('seed_type', 'segwit')
        return Standard_Wallet(db, config=config)

    async def test_get_redeem_script_does_not_raise_NotImplementedError(self):
        # THE PRECISE crash path from the GUI Details click traceback.
        # Just calling without raising NotImplementedError is the success
        # criterion. Pre-fix this raised inside PubkeyProvider.get_pubkey_bytes.
        wallet = self._build_wallet()
        addr = wallet.get_receiving_addresses()[0]
        redeem_script = wallet.get_redeem_script(addr)
        # P2WPKH addresses have no redeem_script — the descriptor returns
        # None. Non-None means the wallet inferred the wrong type.
        self.assertIsNone(redeem_script,
                          "P2WPKH addresses have no redeem_script; "
                          "non-None means the wallet inferred the wrong type")

    async def test_keystore_uses_bip44_derivation_for_btcq(self):
        # BTCQ aligns with the daemon's m/44'/<coin>'/0' convention from
        # doc/btcq-derivation.md. Upstream Electrum used m/0' for native
        # segwit; that produced different addresses than btcqd from the
        # same seed and broke cross-wallet seed portability.
        wallet = self._build_wallet()
        # Testnet BIP44 coin_type is 1 (universal SLIP-0044 testnet).
        self.assertEqual(wallet.keystore.get_derivation_prefix(),
                         "m/44h/1h/0h",
                         "BTCQ keystore must use BIP44 path matching the daemon")

    async def test_address_path_str_shows_full_bip44_path(self):
        # The Address Details dialog shows wallet.get_address_path_str(addr).
        # Pre-fix this returned just the keystore-relative suffix "m/0/0".
        # Post-fix it returns the full BIP44 path.
        wallet = self._build_wallet()
        addr = wallet.get_receiving_addresses()[0]
        path = wallet.get_address_path_str(addr)
        self.assertEqual(path, "m/44h/1h/0h/0/0",
                         "Details dialog must show the full derivation path")


if __name__ == "__main__":
    unittest.main()
