#!/usr/bin/env python3
"""Load a founder web-wallet key into a temporary node wallet, for the ADR-015 sweep.

The web wallet can't pay the vault (it only builds ordinary P2WPKH outputs),
so the sweep runs from the node instead: this script puts the founder key into
a node wallet named "founder-sweep", then sweep_to_vault.py moves the coins.

    python tools/vault/load_founder_key.py --cli "moonbite-cli -datadir=... -conf=..."

You type the recovery phrase yourself, hidden. It is never printed, logged,
written to disk by this script, or put on a command line (the key goes to the
node over stdin). It derives the key exactly like the web wallet
(wallet.derive_from_seed_phrase), and stops unless the address it opens is
one of the founder addresses in sweep_to_vault.py.

When the sweep is done:  moonbite-cli ... unloadwallet founder-sweep
and delete the founder-sweep folder in the node's wallets directory: it holds
the key unencrypted. The web wallet keeps working as before.
"""

from __future__ import annotations

import argparse
import getpass
import hashlib
import json
import os
import shlex
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", ".."))
sys.path.insert(0, HERE)

from sweep_to_vault import FOUNDER_ADDRESSES  # noqa: E402

WALLET = "founder-sweep"
_B58 = "123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz"


def b58encode(b: bytes) -> str:
    n = int.from_bytes(b, "big")
    out = ""
    while n:
        n, r = divmod(n, 58)
        out = _B58[r] + out
    return "1" * (len(b) - len(b.lstrip(b"\0"))) + out


def b58decode(s: str) -> bytes:
    n = 0
    for ch in s:
        n = n * 58 + _B58.index(ch)
    raw = n.to_bytes((n.bit_length() + 7) // 8, "big")
    return b"\0" * (len(s) - len(s.lstrip("1"))) + raw


def _check(b: bytes) -> bytes:
    return hashlib.sha256(hashlib.sha256(b).digest()).digest()[:4]


def wif_version(wif: str) -> int:
    raw = b58decode(wif)
    if _check(raw[:-4]) != raw[-4:]:
        raise ValueError("bad checksum in the node's sample key")
    return raw[0]


def to_wif(priv_hex: str, version: int) -> str:
    body = bytes([version]) + bytes.fromhex(priv_hex) + b"\x01"  # compressed key
    return b58encode(body + _check(body))


class Cli:
    def __init__(self, cmd: str, wallet: str | None = None):
        # Windows paths: forward slashes work for the node and survive shlex.
        self.argv = shlex.split(cmd.replace(chr(92), "/")) + ([f"-rpcwallet={wallet}"] if wallet else [])

    def __call__(self, method: str, *params, stdin: str | None = None, ok_fail: bool = False):
        args = [p if isinstance(p, str) else json.dumps(p) for p in params]
        argv = self.argv + (["-stdin"] if stdin is not None else []) + [method, *args]
        out = subprocess.run(argv, input=stdin, capture_output=True, text=True)
        if out.returncode != 0:
            if ok_fail:
                return None
            raise SystemExit(f"{method} failed: {out.stderr.strip() or out.stdout.strip()}")
        text = out.stdout.strip()
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return text


def main(argv=None, *, ask=getpass.getpass, make_cli=Cli) -> int:
    ap = argparse.ArgumentParser(description="Load a founder key into a temporary node wallet (ADR-015).")
    ap.add_argument("--cli", required=True, help='moonbite-cli command, e.g. "moonbite-cli -datadir=..."')
    args = ap.parse_args(argv)

    import wallet  # the repo's reference derivation, identical to the web wallet

    phrase = ask("Recovery phrase (hidden): ")
    extra = ask("Extra password, if you set one (Enter for none): ")
    try:
        info = wallet.derive_from_seed_phrase(phrase, extra)
    except ValueError as e:
        raise SystemExit(f"That phrase didn't work: {e}")
    finally:
        phrase = extra = None
    address = info["address"]
    if address not in FOUNDER_ADDRESSES:
        raise SystemExit(f"This phrase opens {address}, which is not a founder address. Nothing loaded.")
    print(f"Phrase opens founder address {address}.")

    node = make_cli(args.cli)
    if WALLET not in (node("listwallets") or []):
        if node("loadwallet", WALLET, ok_fail=True) is None:
            # Legacy wallet (descriptors=false) so importprivkey works.
            node("createwallet", WALLET, "false", "false", "", "false", "false")
    w = make_cli(args.cli, WALLET)

    # Ask the node for its own key format instead of assuming one.
    version = wif_version(w("dumpprivkey", w("getnewaddress", "probe", "bech32")))
    wif = to_wif(info["private_key"], version)
    info = None
    print("Importing the key and rescanning the chain (a minute or two)...")
    w("importprivkey", stdin=f"{wif}\nfounder\ntrue\n")
    wif = None

    mine = w("getaddressinfo", address)
    if not mine.get("ismine"):
        raise SystemExit("The node doesn't recognise the address as its own. Nothing to sweep; stop here.")
    bal = w("getbalances").get("mine", {})
    print(f"Loaded. Spendable now: {bal.get('trusted')} MBITE, still maturing: {bal.get('immature')} MBITE.")
    print("\nNext, a dry run of the sweep (sends nothing):")
    print(f'  python tools/vault/sweep_to_vault.py --cli "{args.cli} -rpcwallet={WALLET}"')
    return 0


if __name__ == "__main__":
    sys.exit(main())