#!/usr/bin/env python3
# Copyright (c) 2026 The MoonBite developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test the SOS timelock claw-back vault RPCs (createvault/vaultstatus/sos/unvault).

Vault witnessScript (P2WSH):
    OP_IF   <recoveryPubKey> OP_CHECKSIG
    OP_ELSE <delayBlocks> OP_CHECKSEQUENCEVERIFY OP_DROP <hotPubKey> OP_CHECKSIG
    OP_ENDIF

Recovery key can spend immediately; hot key is CSV-frozen for <delayBlocks>.
"""

from test_framework.test_framework import BitcoinTestFramework
from test_framework.util import (
    assert_equal,
    assert_greater_than,
    assert_raises_rpc_error,
    softfork_active,
)

NOT_FINAL_ERROR = "non-BIP68-final"
DELAY = 10  # small CSV delay so the test can mine past it quickly


class VaultTest(BitcoinTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        self.setup_clean_chain = True
        # -acceptnonstdtxn: the vault witnessScript is a non-standard IF/ELSE/CSV
        # template, so we must accept non-standard txs into the mempool.
        self.extra_args = [[
            "-acceptnonstdtxn=1",
            "-vbparams=mweb:-2:0",
        ]]

    def skip_test_if_missing_module(self):
        self.skip_if_no_wallet()

    def new_pubkey(self):
        addr = self.node.getnewaddress()
        return self.node.getaddressinfo(addr)["pubkey"]

    def fund_vault(self, vault_addr, amount):
        txid = self.node.sendtoaddress(vault_addr, amount)
        self.node.generate(1)
        # locate the vout paying the vault
        raw = self.node.getrawtransaction(txid, True)
        vout = None
        for o in raw["vout"]:
            spk = o["scriptPubKey"]
            if vault_addr in spk.get("addresses", []) or spk.get("address") == vault_addr:
                vout = o["n"]
                break
        assert vout is not None, "vault vout not found"
        return txid, vout

    def run_test(self):
        self.node = self.nodes[0]

        self.log.info("Activate CSV (regtest activates at height 432)")
        self.node.generate(432)
        assert softfork_active(self.node, 'csv')
        self.node.generate(20)  # extra spendable coinbase maturity headroom

        recovery_pub = self.new_pubkey()
        hot_pub = self.new_pubkey()

        self.log.info("createvault builds a P2WSH vault and registers it watch-only")
        v = self.node.createvault(recovery_pub, hot_pub, DELAY, "sos-vault")
        vault_addr = v["address"]
        wscript = v["witness_script"]
        assert_equal(v["delay_blocks"], DELAY)
        assert vault_addr.startswith("rmoon"), "expected regtest bech32 rmoon HRP, got %s" % vault_addr

        # bad witnessScript rejected
        assert_raises_rpc_error(-8, "Not a valid vault witnessScript",
                                self.node.vaultstatus, "00")

        self.log.info("Case 1: recovery key claws back immediately")
        txid, vout = self.fund_vault(vault_addr, 5)
        st = self.node.vaultstatus(wscript)
        assert_equal(st["address"], vault_addr)
        assert_equal(st["delay_blocks"], DELAY)
        assert_greater_than(len(st["utxos"]), 0)
        u = st["utxos"][0]
        assert_equal(u["recovery_spendable"], True)
        assert_greater_than(u["hot_spendable_after"], 0)  # hot still frozen

        dest = self.node.getnewaddress()
        signed = self.node.unvault(wscript, txid, vout, dest, "recovery")
        rtxid = self.node.sendrawtransaction(signed["hex"])
        assert_equal(rtxid, signed["txid"])
        self.node.generate(1)
        # UTXO now spent -> gone from vaultstatus
        st2 = self.node.vaultstatus(wscript)
        assert all(x["txid"] != txid or x["vout"] != vout for x in st2["utxos"])
        self.log.info("  recovery claw-back confirmed")

        self.log.info("Case 2: hot key is rejected before the CSV delay matures")
        txid2, vout2 = self.fund_vault(vault_addr, 5)
        early = self.node.unvault(wscript, txid2, vout2, dest, "hot")
        assert_raises_rpc_error(-26, NOT_FINAL_ERROR,
                                self.node.sendrawtransaction, early["hex"])
        self.log.info("  hot spend correctly rejected as non-BIP68-final")

        self.log.info("Case 3: hot key spends after the CSV delay matures")
        # UTXO already has 1 confirmation from fund_vault; mine DELAY-1 more.
        self.node.generate(DELAY - 1)
        st3 = self.node.vaultstatus(wscript)
        u2 = next(x for x in st3["utxos"] if x["txid"] == txid2 and x["vout"] == vout2)
        assert_equal(u2["hot_spendable_after"], 0)
        mature = self.node.unvault(wscript, txid2, vout2, dest, "hot")
        htxid = self.node.sendrawtransaction(mature["hex"])
        assert_equal(htxid, mature["txid"])
        self.node.generate(1)
        self.log.info("  hot spend confirmed after delay")

        self.log.info("Case 4: sos sweeps the spendable balance into a vault")
        bal_before = self.node.getbalance()
        assert_greater_than(bal_before, 0)
        sweep_txid = self.node.sos(vault_addr)
        self.node.generate(1)
        # the swept coins now sit in the vault as watch-only, not spendable balance
        st4 = self.node.vaultstatus(wscript)
        assert_greater_than(st4["total"], 0)
        assert any(x["recovery_spendable"] for x in st4["utxos"])
        self.log.info("  sos sweep landed in vault (txid=%s)" % sweep_txid)

        self.log.info("All vault scenarios passed.")


if __name__ == '__main__':
    VaultTest().main()