"""
SSP (Self-Sovereign Protocol) 2-of-2 multisig wallet integration.

Acts as the wallet side (key #1) in the SSP protocol. Coordinates with
SSP Key (mobile app, key #2) via the SSP Relay server for transaction
co-signing.

Usage — new wallet:
    ssp = await SSPWallet.create_paired("my_ssp", mnemonic, "relay.sspwallet.io")
    # ... scan QR codes in SSP Key ...
    wallet = await ssp.complete_pairing()
    txid = await ssp.send([(addr, amount)])

Usage — re-pair existing:
    ssp = await SSPWallet.repair(wallet, "relay.sspwallet.io")
    # ... scan QR codes in SSP Key ...
    await ssp.complete_repair()

Usage — send from existing:
    ssp = await SSPWallet.from_existing(wallet, "relay.sspwallet.io")
    txid = await ssp.send([(addr, amount)])
"""

from __future__ import annotations

import hashlib
import logging

from fluxwallet.db import Db
from fluxwallet.encoding import change_base, pubkeyhash_to_addr_bech32
from fluxwallet.keys import HDKey
from fluxwallet.mnemonic import Mnemonic
from fluxwallet.networks import Network
from fluxwallet.services.relay import SSPRelayClient, RelayError
from fluxwallet.wallet import Wallet, WalletError

_logger = logging.getLogger(__name__)

# BIP48 key path for SSP multisig (Flux uses P2SH = scriptType 0)
SSP_KEY_PATH = [
    "m", "purpose'", "coin_type'", "account'", "script_type'", "change", "address_index"
]
SSP_PURPOSE = 48
SSP_SCRIPT_TYPE = 0  # P2SH


class SSPWallet:
    """SSP 2-of-2 multisig wallet (initiator / key #1 side).

    Wraps a standard Wallet with SSP relay integration for co-signing
    with the SSP Key mobile app.
    """

    def __init__(
        self,
        wallet: Wallet | None,
        relay: SSPRelayClient,
        master_key: HDKey,
    ):
        self.wallet = wallet
        self.relay = relay
        self._master_key = master_key
        self._wk_identity: str | None = None
        self._peer_xpub: HDKey | None = None
        self._peer_btc_xpub: HDKey | None = None
        self._peer_btc_identity_pubkey: str | None = None

    # --- Construction ---

    @classmethod
    async def create_paired(
        cls,
        name: str,
        keys: str | HDKey,
        relay_url: str,
        network: str = "flux",
    ) -> SSPWallet:
        """Start creating a new SSP multisig wallet.

        Generates the master key from a mnemonic (or accepts an HDKey).
        The multisig wallet is created by complete_pairing() after SSP Key
        responds with its xpub.

        :param name: Name for the multisig wallet to be created
        :param keys: Mnemonic seed phrase or HDKey
        :param relay_url: SSP Relay server URL
        :param network: Network name (default: 'flux')
        :returns: SSPWallet instance ready for pairing
        """
        await Db.start()

        if isinstance(keys, str) and len(keys.split()) > 1:
            seed = Mnemonic().to_seed(keys)
            master_key = HDKey.from_seed(seed, network=network)
        elif isinstance(keys, HDKey):
            master_key = keys
        else:
            master_key = HDKey(keys, network=network)

        identity_key = _derive_identity_key(master_key)
        relay = SSPRelayClient(relay_url, identity_key)

        ssp = cls(None, relay, master_key)
        ssp._wallet_name = name
        ssp._network = network
        return ssp

    @classmethod
    async def repair(
        cls,
        wallet: Wallet,
        relay_url: str,
    ) -> SSPWallet:
        """Re-pair an existing multisig wallet with SSP Key.

        Extracts our master key from the wallet's private cosigner.
        After scanning xpubs in SSP Key, call complete_repair() to
        update the stored ssp_peer_pubkey.

        :param wallet: Existing multisig Wallet
        :param relay_url: SSP Relay server URL
        :returns: SSPWallet instance ready for re-pairing
        """
        if not wallet.multisig:
            raise WalletError("repair() requires a multisig wallet")

        private_cosigner = None
        for cs in wallet.cosigner:
            if cs.main_key and cs.main_key.is_private:
                private_cosigner = cs
                break
        if not private_cosigner:
            raise WalletError("No private key found in multisig cosigners")

        master_key = private_cosigner.main_key.key()
        identity_key = _derive_identity_key(master_key)
        relay = SSPRelayClient(relay_url, identity_key)

        return cls(wallet, relay, master_key)

    @classmethod
    async def from_existing(
        cls,
        wallet: Wallet,
        relay_url: str,
    ) -> SSPWallet:
        """Open an existing paired multisig wallet for sending.

        Extracts our master key, derives relay auth from stored
        ssp_peer_pubkey. Ready to send transactions immediately.

        :param wallet: Existing multisig Wallet (must have ssp_peer_pubkey)
        :param relay_url: SSP Relay server URL
        :returns: SSPWallet instance ready to send
        """
        if not wallet.multisig:
            raise WalletError("from_existing() requires a multisig wallet")

        private_cosigner = None
        for cs in wallet.cosigner:
            if cs.main_key and cs.main_key.is_private:
                private_cosigner = cs
                break
        if not private_cosigner:
            raise WalletError("No private key found in multisig cosigners")
        master_key = private_cosigner.main_key.key()

        if not wallet.ssp_peer_pubkey:
            raise WalletError("Wallet missing ssp_peer_pubkey. Re-pair to fix.")

        wk_identity, witness_script, our_pubkey = _derive_relay_auth(
            master_key, wallet.ssp_peer_pubkey,
        )

        identity_key = _derive_identity_key(master_key)
        relay = SSPRelayClient(
            relay_url, identity_key,
            witness_script=witness_script,
            identity_pubkey=our_pubkey,
        )

        ssp = cls(wallet, relay, master_key)
        ssp._wk_identity = wk_identity
        return ssp

    # --- Key Derivation ---

    def identity_key(self) -> HDKey:
        """Internal identity key for relay authentication.

        Derivation: m/48'/0'/0'/0'/10/0 (BTC identity chain)
        """
        return _derive_identity_key(self._master_key)

    def external_identity(self) -> HDKey:
        """External identity key (public FluxID).

        Derivation: m/48'/0'/0'/0'/11/0 (Bitcoin address)
        """
        return _derive_identity_key(self._master_key, type_index=11)

    def xpub(self, network: str | None = None) -> str:
        """BIP48 public master key (xpub) for sharing with SSP Key.

        Returns the xpub at ``m/48'/<coin>'/0'/0'``.

        :param network: Network to derive for (default: wallet's network)
        :returns: xpub WIF string
        """
        if network is None:
            network = self.wallet.network.name if self.wallet else self._network
        coin_type = Network(network).bip44_cointype
        account_key = self._master_key.subkey_for_path(
            f"m/48'/{coin_type}'/0'/0'"
        )
        return account_key.wif_public()

    def xpub_identity(self) -> str:
        """BTC identity xpub for SSP Key initial pairing.

        :returns: BTC xpub WIF string
        """
        return self.xpub(network="bitcoin")

    # --- Pairing ---

    async def pair_with_key(self) -> str:
        """Initiate pairing with SSP Key.

        Returns our BIP48 xpub for SSP Key to scan/input.

        :returns: Our BIP48 xpub WIF string
        """
        _logger.info("SSP pairing initiated, share this xpub with SSP Key")
        return self.xpub()

    async def wait_for_btc_sync(self, timeout: float = 300) -> str:
        """Wait for SSP Key to respond to BTC identity sync.

        Captures the peer's BTC xpub for relay identity derivation.

        :param timeout: Maximum seconds to wait
        :returns: The peer's BTC xpub WIF string
        :raises RelayError: On timeout
        """
        identity = self.identity_key()
        wallet_identity = identity.address()

        _logger.info("Polling relay for BTC sync (identity: %s)", wallet_identity)
        sync_data = await self.relay.poll_sync(wallet_identity, timeout=timeout, chain="btc")

        if not sync_data:
            raise RelayError("Timeout waiting for SSP Key BTC sync")

        peer_xpub_wif = sync_data.get("keyXpub", "")
        if not peer_xpub_wif:
            raise RelayError("No xpub received from SSP Key")

        # SSP Key sends BTC xpub with custom version bytes (Zpub).
        # Convert to standard xpub format for parsing.
        peer_xpub_wif = _normalize_xpub(peer_xpub_wif)
        self._peer_btc_xpub = HDKey(peer_xpub_wif, network="bitcoin")

        # Store the BTC identity pubkey
        peer_id = self._peer_btc_xpub.subkey_for_path("10/0")
        self._peer_btc_identity_pubkey = peer_id.public_hex

        _logger.info("Received SSP Key BTC xpub")
        return peer_xpub_wif

    async def wait_for_chain_sync(self, timeout: float = 300) -> str:
        """Wait for SSP Key to respond to chain (Flux) sync.

        Captures the peer's chain xpub for multisig wallet creation.

        :param timeout: Maximum seconds to wait
        :returns: The peer's chain xpub WIF string
        :raises RelayError: On timeout
        """
        identity = self.identity_key()
        wallet_identity = identity.address()

        network = self.wallet.network.name if self.wallet else self._network
        _logger.info("Polling relay for chain sync (identity: %s)", wallet_identity)
        sync_data = await self.relay.poll_sync(wallet_identity, timeout=timeout, chain=network)

        if not sync_data:
            raise RelayError("Timeout waiting for SSP Key chain sync")

        peer_xpub_wif = sync_data.get("keyXpub", "")
        if not peer_xpub_wif:
            raise RelayError("No xpub received from SSP Key")

        self._peer_xpub = HDKey(peer_xpub_wif, network=network)

        _logger.info("Received SSP Key chain xpub")
        return peer_xpub_wif

    async def complete_pairing(self) -> Wallet:
        """Create the 2-of-2 multisig wallet from both xpubs.

        Must be called after wait_for_pairing() succeeds.
        Only for new wallets created via create_paired().

        :returns: The new multisig Wallet
        """
        if not self._peer_xpub:
            raise WalletError("Not paired yet — call wait_for_pairing() first")
        if self.wallet and self.wallet.multisig:
            raise WalletError("Wallet already exists. Use complete_repair() instead.")

        name = self._wallet_name
        network = self._network

        from fluxwallet.wallets import wallet_exists
        if await wallet_exists(name):
            raise WalletError(f"Wallet '{name}' already exists")

        ms_wallet = await Wallet.create(
            name,
            keys=[self._master_key, self._peer_xpub.public()],
            sigs_required=2,
            network=network,
            cosigner_id=0,
            key_path=SSP_KEY_PATH,
            purpose=SSP_PURPOSE,
        )

        if not self._peer_btc_identity_pubkey:
            raise WalletError("BTC identity not synced. Call wait_for_btc_sync() first.")
        await ms_wallet.set_ssp_peer_pubkey(self._peer_btc_identity_pubkey)

        self.wallet = ms_wallet
        _logger.info("Created SSP multisig wallet: %s", name)
        return ms_wallet

    async def complete_repair(self) -> None:
        """Update the existing wallet's ssp_peer_pubkey after re-pairing.

        Must be called after wait_for_btc_sync() succeeds.
        Only for existing wallets via repair().
        """
        if not self._peer_btc_identity_pubkey:
            raise WalletError("BTC identity not synced. Call wait_for_btc_sync() first.")
        if not self.wallet or not self.wallet.multisig:
            raise WalletError("No existing multisig wallet. Use complete_pairing() instead.")

        await self.wallet.set_ssp_peer_pubkey(self._peer_btc_identity_pubkey)

        _logger.info("Updated ssp_peer_pubkey for wallet: %s", self.wallet.name)

    # --- Transactions ---

    async def send(
        self,
        outputs: list[tuple[str, int]],
        fee: int | None = None,
        timeout: float = 300,
    ) -> str:
        """Create, partially sign, and co-sign a transaction via SSP Key.

        :param outputs: List of (address, amount_satoshis) tuples
        :param fee: Fee in satoshis (default: auto-calculate)
        :param timeout: Max seconds to wait for SSP Key to co-sign
        :returns: Transaction ID after broadcast
        :raises RelayError: On timeout or rejection
        """
        if not self.wallet or not self.wallet.multisig:
            raise WalletError("Cannot send from non-multisig wallet.")

        if not self._wk_identity:
            raise WalletError("wkIdentity not available. Wallet missing ssp_peer_pubkey.")

        # Ensure relay is connected
        await self.relay.connect(self._wk_identity)

        # Build and partially sign
        wt = await self.wallet.transaction_create(outputs, fee=fee)
        wt.sign(partial=True)

        raw_hex = wt.tx.raw_hex()
        _logger.info("Partially signed tx (%d bytes), sending to SSP Key", len(raw_hex) // 2)

        # Build UTXO metadata for SSP Key
        utxos = []
        for inp in wt.tx.inputs:
            utxos.append({
                "txid": inp.prev_txid.hex() if isinstance(inp.prev_txid, bytes) else inp.prev_txid,
                "vout": inp.output_n_int,
                "satoshis": str(inp.value),
                "scriptPubKey": "",
                "confirmations": 0,
                "coinbase": False,
            })

        # Determine derivation path from the first input's key path.
        # Key path is like "m/48'/19167'/0'/0'/0/1" — we need the last
        # two components: typeIndex (0=receive, 1=change) and addressIndex.
        inp_key_path = wt.tx.inputs[0].key_path
        if inp_key_path and "/" in inp_key_path:
            parts = inp_key_path.split("/")
            path = f"{parts[-2]}-{parts[-1]}"
        else:
            path = "0-0"

        # Post to relay
        await self.relay.post_action(
            action="tx",
            payload=raw_hex,
            chain="flux",
            path=path,
            wk_identity=self._wk_identity,
            utxos=utxos,
        )

        # Wait for SSP Key to co-sign, finalize, and broadcast
        txid = await self.relay.wait_for_txid(timeout=timeout)
        _logger.info("Transaction broadcast by SSP Key: %s", txid)

        return txid

    # --- Cleanup ---

    async def close(self) -> None:
        """Close relay connections."""
        await self.relay.close()


XPUB_VERSION = bytes.fromhex("0488b21e")


def _normalize_xpub(xpub_wif: str) -> str:
    """Convert any extended public key (Zpub, ypub, etc.) to standard xpub format.

    SSP Key uses custom BIP48 version bytes for BTC (Zpub prefix, 0x02aa7ed3).
    This converts to standard xpub (0x0488b21e) so our HDKey can parse it.
    The underlying key data (chain code + public key) is unchanged.

    :param xpub_wif: Extended public key in any format
    :returns: Standard xpub WIF string
    """
    if xpub_wif.startswith("xpub"):
        return xpub_wif
    raw = change_base(xpub_wif, 58, 256)
    if raw is None or len(raw) < 78:
        raise WalletError(f"Invalid extended key: {xpub_wif[:20]}...")
    standard = XPUB_VERSION + raw[4:]
    return change_base(standard, 256, 58)


def _derive_relay_auth(
    our_master_key: HDKey,
    peer_identity_pubkey_hex: str,
) -> tuple[str, str, str]:
    """Derive wkIdentity, witnessScript, and our identity pubkey for relay auth.

    :param our_master_key: Our private master key (depth 0)
    :param peer_identity_pubkey_hex: Peer's BTC identity compressed pubkey (hex)
    :returns: (wk_identity, witness_script_hex, our_identity_pubkey_hex)
    """
    # Our BTC identity pubkey at m/48'/0'/0'/0'/10/0
    our_btc_account = our_master_key.subkey_for_path("48'/0'/0'/0'")
    our_id = our_btc_account.subkey_for_path("10/0")
    our_id_pubkey = our_id.public_hex

    # Build 2-of-2 witness script with sorted pubkeys
    sorted_pubkeys = sorted([our_id_pubkey, peer_identity_pubkey_hex])
    pubkey_bufs = [bytes.fromhex(pk) for pk in sorted_pubkeys]
    ws = bytes([0x52])  # OP_2
    for pk in pubkey_bufs:
        ws += bytes([len(pk)]) + pk
    ws += bytes([0x52, 0xae])  # OP_2 OP_CHECKMULTISIG

    # wkIdentity = P2WSH bech32 address (Bitcoin network)
    ws_hash = hashlib.sha256(ws).digest()
    wk_identity = pubkeyhash_to_addr_bech32(ws_hash, prefix="bc", witver=0)

    return wk_identity, ws.hex(), our_id_pubkey


def _derive_identity_key(
    master_key: HDKey,
    type_index: int = 10,
) -> HDKey:
    """Derive the SSP identity key.

    SSP always uses the **BTC identity chain** (``identityChain = 'btc'``,
    coin_type=0) for identity, regardless of which blockchain is being
    used. The identity is derived from the BTC xpub at
    ``m/48'/0'/0'/0'`` → ``child(typeIndex)/child(0)``, producing a
    Bitcoin P2PKH address.

    :param master_key: HD master key (depth 0, with private key)
    :param type_index: 10 for internal identity, 11 for external (public FluxID)
    :returns: Identity HDKey with Bitcoin-network address
    """
    # Always BTC coin_type=0 for identity
    account_key = master_key.subkey_for_path("m/48'/0'/0'/0'")
    identity_key = account_key.subkey_for_path(f"{type_index}/0")

    # Re-wrap on Bitcoin network for P2PKH address format
    return HDKey(
        key=identity_key.private_byte,
        chain=identity_key.chain,
        network="bitcoin",
    )
