"""BTCQ extended key derivation — ML-DSA-65 with hardened-only BIP32 tree.

Implements the BTCQ HD wallet derivation specified in
``doc/btcq-derivation.md`` in the btcq daemon repository.

Differences from upstream BIP32 (which Electrum's :mod:`electrum.bip32`
implements for Bitcoin):

- Leaf primitive is ML-DSA-65 (FIPS 204), not secp256k1. The "private
  key" stored in an extended-key node is a 32-byte ``xi`` seed (RFC 9881);
  the "public key" is the 1952-byte ML-DSA-65 ``pk``.
- All derivation is hardened — there is no public-only child derivation.
  ML-DSA-65 has no group homomorphism enabling that. ``subkey_at_public_derivation``
  raises :class:`NotImplementedError`.
- Extended public keys carry the full 1952-byte ML-DSA pk, so xpub
  serialization is **1997 bytes** (vs BIP32's 78). xprv keeps BIP32's
  78-byte size since the 32-byte xi fits in the original secret-key slot.

The tree structure (HMAC-SHA512 → IL/IR split, chain code, depth, child
fingerprint) follows BIP32 unchanged. The non-key utilities
(``convert_bip32_*``, :class:`KeyOriginInfo`, :data:`BIP32_PRIME`) live in
:mod:`electrum.bip32` and are imported here.

Replaces the previous ``BIP32Node`` shim that wrapped ``electrum_ecc``'s
secp256k1 ``ECPubkey`` / ``ECPrivkey`` classes around ML-DSA byte buffers.
That shim silently produced wrong bytes when ``get_public_key_bytes()``
was called on a node holding a 32-byte ``xi`` seed (which secp256k1 then
treated as a valid private scalar). This module removes the secp256k1
abstraction entirely from BTCQ's HD derivation path.
"""
import hashlib
import hmac
import struct
from typing import NamedTuple, Optional, Iterable, Sequence, Tuple, Union, List

from . import constants
from .bip32 import (BIP32_PRIME, UINT32_MAX,
                    convert_bip32_strpath_to_intpath,
                    convert_bip32_intpath_to_strpath,
                    KeyOriginInfo)
from .bitcoin import EncodeBase58Check, DecodeBase58Check
from .crypto import hash_160
from .logging import get_logger
from .util import bfh, BitcoinException


_logger = get_logger(__name__)


# ML-DSA-65 / BIP32 byte sizes
MLDSA65_PK_SIZE = 1952
MLDSA65_XI_SIZE = 32
CHAINCODE_SIZE = 32

# Serialized payload sizes (without the trailing 4-byte base58check checksum):
#   common header: version(4) + depth(1) + fingerprint(4) + child_number(4) + chaincode(32) = 45
#   xprv suffix:   marker(1) + xi(32)                                                       = 33  → 78 total
#   xpub suffix:   pk(1952)                                                                 = 1952 → 1997 total
XPRV_PAYLOAD_SIZE = 78
XPUB_PAYLOAD_SIZE = 4 + 1 + 4 + 4 + 32 + MLDSA65_PK_SIZE  # 1997


class InvalidMasterKeyVersionBytes(BitcoinException):
    """Raised when an extended key's version bytes don't match any known network."""


def xprv_header(xtype: str, *, net=None) -> bytes:
    """Look up the network's xprv version-bytes for ``xtype``.

    BTCQ has exactly one signature scheme (ML-DSA-65) and one address
    format (bech32), so unlike upstream Electrum we don't keep separate
    xprv-header byte sequences for p2wpkh / p2pkh / p2sh / etc. Every
    xtype falls back to ``'standard'``.
    """
    if net is None:
        net = constants.net
    header = net.XPRV_HEADERS.get(xtype) or net.XPRV_HEADERS['standard']
    return header.to_bytes(length=4, byteorder="big")


def xpub_header(xtype: str, *, net=None) -> bytes:
    """Look up the network's xpub version-bytes for ``xtype``. See :func:`xprv_header`."""
    if net is None:
        net = constants.net
    header = net.XPUB_HEADERS.get(xtype) or net.XPUB_HEADERS['standard']
    return header.to_bytes(length=4, byteorder="big")


class BtcqExtKey(NamedTuple):
    """A BTCQ extended key node, public or private.

    Replaces upstream Electrum's ``BIP32Node``. Holds:

    - ``pk``: 1952-byte ML-DSA-65 public key. **Always populated.**
      For nodes that originate from a private path, ``pk`` is computed
      via ``mldsa_keygen_from_seed(xi)`` at construction time.
    - ``chaincode``: 32-byte HMAC-SHA512 chain code (BIP32-shaped).
    - ``xtype, depth, fingerprint, child_number``: BIP32 metadata.
    - ``xi``: 32-byte ML-DSA-65 seed (FIPS 204 KeyGen_internal input).
      ``None`` for public-only nodes; populated for private nodes.

    A node ``is_private()`` iff ``xi is not None``. Public-only nodes
    still carry ``pk`` and can compute fingerprints, sign-verify, etc.
    — they just cannot derive children.

    The expanded 4032-byte ML-DSA secret key is **not** stored — it is
    recomputed via ``mldsa_keygen_from_seed(xi)`` during each derivation
    step. This costs one keygen per step (~1-2s on Zen 4) but keeps node
    objects compact (1952+32+misc bytes vs 1952+32+4032).
    """
    xtype: str
    pk: bytes
    chaincode: bytes
    depth: int = 0
    fingerprint: bytes = b'\x00' * 4   # serialized layout: this is the *parent's* fingerprint
    child_number: bytes = b'\x00' * 4
    xi: Optional[bytes] = None

    # ---- construction --------------------------------------------------

    @classmethod
    def from_rootseed(cls, seed: bytes, *, xtype: str) -> 'BtcqExtKey':
        """Create a depth-0 master node from a BIP32 seed.

        Mirrors the daemon's ``CExtKey::SetSeed``:
        ``HMAC-SHA512("Bitcoin seed", seed) → IL || IR``, then expand
        ``IL`` to the master ML-DSA keypair.
        """
        from .mldsa import mldsa_keygen_from_seed
        I = hmac.new(b"Bitcoin seed", seed, hashlib.sha512).digest()
        master_xi, master_chaincode = I[:32], I[32:]
        master_pk, _ = mldsa_keygen_from_seed(master_xi)
        return cls(xtype=xtype, pk=master_pk, chaincode=master_chaincode,
                   xi=master_xi)

    @classmethod
    def from_xkey(cls, xkey: str, *, net=None,
                  allow_custom_headers: bool = True) -> 'BtcqExtKey':
        """Parse a base58check-encoded xprv (78 B) or xpub (1997 B) string."""
        if net is None:
            net = constants.net
        raw = DecodeBase58Check(xkey)
        if len(raw) not in (XPRV_PAYLOAD_SIZE, XPUB_PAYLOAD_SIZE):
            raise BitcoinException(
                f'Invalid length for extended key: {len(raw)} '
                f'(expected {XPRV_PAYLOAD_SIZE} for xprv or {XPUB_PAYLOAD_SIZE} for xpub)'
            )
        depth = raw[4]
        fingerprint = raw[5:9]
        child_number = raw[9:13]
        chaincode = raw[13:13 + CHAINCODE_SIZE]
        header = int.from_bytes(raw[0:4], byteorder='big')
        if header in net.XPRV_HEADERS_INV:
            headers_inv = net.XPRV_HEADERS_INV
            is_private = True
        elif header in net.XPUB_HEADERS_INV:
            headers_inv = net.XPUB_HEADERS_INV
            is_private = False
        else:
            raise InvalidMasterKeyVersionBytes(f'Invalid extended key format: {hex(header)}')
        xtype = headers_inv[header]
        if not allow_custom_headers and xtype != "standard":
            raise ValueError(f"only standard xpub/xprv allowed. found custom xtype={xtype}")
        if is_private:
            # 0x00 marker byte at offset 45, xi starts at offset 46.
            if len(raw) != XPRV_PAYLOAD_SIZE:
                raise BitcoinException(
                    f'xprv header but length {len(raw)} != {XPRV_PAYLOAD_SIZE}'
                )
            if raw[45] != 0x00:
                raise BitcoinException(
                    f'xprv has non-zero marker byte: {raw[45]:#x} (expected 0x00)'
                )
            xi = raw[46:46 + MLDSA65_XI_SIZE]
            from .mldsa import mldsa_keygen_from_seed
            pk, _ = mldsa_keygen_from_seed(xi)
            return cls(xtype=xtype, pk=pk, chaincode=chaincode,
                       depth=depth, fingerprint=fingerprint,
                       child_number=child_number, xi=xi)
        else:
            # Public-only: pk starts immediately after chaincode at offset 45.
            if len(raw) != XPUB_PAYLOAD_SIZE:
                raise BitcoinException(
                    f'xpub header but length {len(raw)} != {XPUB_PAYLOAD_SIZE}'
                )
            pk = raw[45:45 + MLDSA65_PK_SIZE]
            return cls(xtype=xtype, pk=pk, chaincode=chaincode,
                       depth=depth, fingerprint=fingerprint,
                       child_number=child_number, xi=None)

    @classmethod
    def from_bytes(cls, b: bytes) -> 'BtcqExtKey':
        if len(b) not in (XPRV_PAYLOAD_SIZE, XPUB_PAYLOAD_SIZE):
            raise BitcoinException(
                f"unexpected xkey raw bytes len {len(b)}; "
                f"expected {XPRV_PAYLOAD_SIZE} or {XPUB_PAYLOAD_SIZE}"
            )
        return cls.from_xkey(EncodeBase58Check(b))

    # ---- serialization -------------------------------------------------

    def to_xprv(self, *, net=None) -> str:
        return EncodeBase58Check(self.to_xprv_bytes(net=net))

    def to_xprv_bytes(self, *, net=None) -> bytes:
        if not self.is_private():
            raise Exception("cannot serialize as xprv; xi missing")
        payload = (xprv_header(self.xtype, net=net) +
                   bytes([self.depth]) +
                   self.fingerprint +
                   self.child_number +
                   self.chaincode +
                   bytes([0]) +
                   self.xi)
        assert len(payload) == XPRV_PAYLOAD_SIZE, \
            f"unexpected xprv payload len {len(payload)}"
        return payload

    def to_xpub(self, *, net=None) -> str:
        return EncodeBase58Check(self.to_xpub_bytes(net=net))

    def to_xpub_bytes(self, *, net=None) -> bytes:
        payload = (xpub_header(self.xtype, net=net) +
                   bytes([self.depth]) +
                   self.fingerprint +
                   self.child_number +
                   self.chaincode +
                   self.pk)
        assert len(payload) == XPUB_PAYLOAD_SIZE, \
            f"unexpected xpub payload len {len(payload)}"
        return payload

    def to_xkey(self, *, net=None) -> str:
        return self.to_xprv(net=net) if self.is_private() else self.to_xpub(net=net)

    def to_bytes(self, *, net=None) -> bytes:
        return self.to_xprv_bytes(net=net) if self.is_private() else self.to_xpub_bytes(net=net)

    # ---- inspection ----------------------------------------------------

    def is_private(self) -> bool:
        return self.xi is not None

    def convert_to_public(self) -> 'BtcqExtKey':
        """Return a public-only copy, dropping the xi seed."""
        if not self.is_private():
            return self
        return self._replace(xi=None)

    def calc_fingerprint_of_this_node(self) -> bytes:
        """Returns the 4-byte fingerprint of this node (hash of its own pk).

        Note that ``self.fingerprint`` is the *parent's* fingerprint
        (per BIP32 layout). Use this method to compute the child-side
        fingerprint to embed in derived nodes.
        """
        return hash_160(self.pk)[:4]

    # ---- derivation ----------------------------------------------------

    def subkey_at_private_derivation(
        self, path: Union[str, Iterable[int]]
    ) -> 'BtcqExtKey':
        """Walk down the BIP32 tree privately. All steps must be hardened in BTCQ."""
        if path is None:
            raise Exception("derivation path must not be None")
        if isinstance(path, str):
            path = convert_bip32_strpath_to_intpath(path)
        path = list(path)
        if not self.is_private():
            raise Exception("cannot do private derivation; xi missing")
        if not path:
            return self

        from .mldsa import mldsa_keygen_from_seed

        xi = self.xi
        pk = self.pk
        chaincode = self.chaincode
        depth = self.depth
        parent_pk_for_fp = pk

        for child_index in path:
            if child_index < 0 or child_index > UINT32_MAX:
                raise ValueError(f"bip32 child index out of range: {child_index}")
            # Expand parent xi to the full ML-DSA sk to extract rho (sk[:32]),
            # per the BTCQ derivation spec's Representation Invariant.
            _, parent_sk = mldsa_keygen_from_seed(xi)
            rho = parent_sk[:32]

            data = bytes([0]) + rho + child_index.to_bytes(4, 'big', signed=False)
            I = hmac.new(chaincode, data, hashlib.sha512).digest()
            child_xi = I[:32]
            child_chaincode = I[32:]

            parent_pk_for_fp = pk
            child_pk, _ = mldsa_keygen_from_seed(child_xi)

            xi = child_xi
            pk = child_pk
            chaincode = child_chaincode
            depth += 1

        fingerprint = hash_160(parent_pk_for_fp)[:4]
        final_index = path[-1]
        child_number = final_index.to_bytes(4, 'big', signed=False)
        return BtcqExtKey(
            xtype=self.xtype,
            pk=pk,
            chaincode=chaincode,
            depth=depth,
            fingerprint=fingerprint,
            child_number=child_number,
            xi=xi,
        )

    def subkey_at_public_derivation(
        self, path: Union[str, Iterable[int]]
    ) -> 'BtcqExtKey':
        """Always raises — ML-DSA-65 has no public-only child derivation.

        Callers that previously used this for unhardened-path optimization
        (e.g. ``Xpub._derive_pubkey``'s base implementation) must walk
        privately from the ``xprv`` instead. ``BIP32_KeyStore._derive_pubkey``
        is the canonical override.
        """
        if isinstance(path, str):
            path = convert_bip32_strpath_to_intpath(path)
        path = list(path) if path is not None else []
        if not path:
            # No-op path: just return a public-only copy of self.
            return self.convert_to_public()
        raise NotImplementedError(
            "BTCQ has no public-only child derivation — ML-DSA-65 lacks "
            "the group homomorphism that BIP32's CKD_pub relies on. "
            "Walk privately from the xprv instead."
        )


# ---- helpers reused across the codebase --------------------------------

def xpub_type(x: str) -> str:
    assert x is not None
    return BtcqExtKey.from_xkey(x).xtype


def is_xpub(text: str) -> bool:
    try:
        node = BtcqExtKey.from_xkey(text)
        return not node.is_private()
    except Exception:
        return False


def is_xprv(text: str) -> bool:
    try:
        node = BtcqExtKey.from_xkey(text)
        return node.is_private()
    except Exception:
        return False


def xpub_from_xprv(xprv: str) -> str:
    return BtcqExtKey.from_xkey(xprv).to_xpub()


def root_fp_and_der_prefix_from_xkey(xkey: str) -> Tuple[Optional[str], Optional[str]]:
    """Returns the root bip32 fingerprint and the derivation path from the
    root to the given xkey, if they can be determined. Otherwise (None, None).

    Mirrors the upstream-Electrum function of the same name in :mod:`electrum.bip32`,
    ported to ``BtcqExtKey``.
    """
    node = BtcqExtKey.from_xkey(xkey)
    derivation_prefix = None
    root_fingerprint = None
    assert node.depth >= 0, node.depth
    if node.depth == 0:
        derivation_prefix = 'm'
        root_fingerprint = node.calc_fingerprint_of_this_node().hex().lower()
    elif node.depth == 1:
        child_number_int = int.from_bytes(node.child_number, 'big')
        derivation_prefix = convert_bip32_intpath_to_strpath([child_number_int])
        root_fingerprint = node.fingerprint.hex()
    return root_fingerprint, derivation_prefix


def is_xkey_consistent_with_key_origin_info(
    xkey: str, *,
    derivation_prefix: Optional[str] = None,
    root_fingerprint: Optional[str] = None,
) -> bool:
    """Cross-check ``xkey`` against optional ``derivation_prefix`` / ``root_fingerprint``."""
    node = BtcqExtKey.from_xkey(xkey)
    int_path = None
    if derivation_prefix is not None:
        int_path = convert_bip32_strpath_to_intpath(derivation_prefix)
    if int_path is not None and len(int_path) != node.depth:
        return False
    if node.depth == 0:
        if bfh(root_fingerprint) != node.calc_fingerprint_of_this_node():
            return False
        if node.child_number != bytes(4):
            return False
    if int_path is not None and node.depth > 0:
        if int.from_bytes(node.child_number, 'big') != int_path[-1]:
            return False
    if node.depth == 1:
        if bfh(root_fingerprint) != node.fingerprint:
            return False
    return True
