#!/usr/bin/env python3
# Copyright (c) 2016-2024 The Hush developers
# Copyright (c) 2018 The Zcash developers
# Distributed under the GPLv3 software license, see the accompanying
# file COPYING or https://www.gnu.org/licenses/gpl-3.0.en.html

from test_framework.test_framework import BitcoinTestFramework
from test_framework.authproxy import JSONRPCException
from test_framework.util import (
    assert_equal,
    initialize_chain_clean,
    p2p_port,
    set_node_times,
    start_node,
    sync_blocks,
    wait_and_assert_operationid_status,
)

import os
import stat

from decimal import Decimal

# Test wallet behaviour with Sapling addresses
class WalletSaplingTest(BitcoinTestFramework):

    def setup_chain(self):
        # The shared initialize_chain() cache builder in test_framework/util.py has not been
        # repaired for DragonX (it spawns the cache nodes by bare name off PATH, drives the
        # CLI without -regtest so it dials the assetchain RPC port, and deletes debug.log
        # from the non-net-specific datadir). Build the same starting state here instead --
        # see _generate_starting_chain() -- so this test does not depend on it.
        print("Initializing test directory " + self.options.tmpdir)
        initialize_chain_clean(self.options.tmpdir, 4)

    # !!! test_framework/util.py:start_node() hardcodes "-connect=0" into every regtest
    # node's argv. Modern Bitcoin Core special-cases that value to mean "make no automatic
    # connections", but THIS codebase does not (net.cpp ThreadOpenConnections just iterates
    # mapMultiArgs["-connect"]), so "0" is dialled as a hostname: it resolves to 0.0.0.0,
    # which on Linux connects to localhost on Params().GetDefaultPort() -- 21768, the live
    # DRAGONX p2p port. Observed directly: a regtest node started by the unmodified
    # framework peered with the production dragonxd on this host and with seven public
    # mainnet nodes (heights ~3.25M) and began ingesting mainnet headers. "-connect" also
    # soft-sets "-listen=0", so the framework's own connect_nodes_bi() can never establish
    # the local links a multi-node test needs.
    #
    # Both problems are in the shared framework, which this port is not allowed to touch, so
    # they are worked around per-node here: the daemon is launched through a tiny wrapper
    # that strips the "-connect=0" argument, and the real topology/listening flags are passed
    # as extra_args (which start_node appends after its own).
    def _daemon_wrapper(self):
        srcdir = os.path.join(
            os.path.dirname(os.path.abspath(__file__)), "..", "..", "src", "dragonxd")
        path = os.path.join(self.options.tmpdir, "dragonxd-no-connect0")
        with open(path, "w") as f:
            f.write("#!/usr/bin/env bash\n")
            f.write("args=()\n")
            f.write('for a in "$@"; do\n')
            f.write('  if [ "$a" = "-connect=0" ]; then continue; fi\n')
            f.write('  args+=("$a")\n')
            f.write("done\n")
            f.write('exec %s "${args[@]}"\n' % os.path.realpath(srcdir))
        os.chmod(path, os.stat(path).st_mode | stat.S_IXUSR)
        return path

    def setup_nodes(self):
        binary = self._daemon_wrapper()
        nodes = []
        for i in range(4):
            extra_args = [
                #'-nuparams=5ba81b19:201', # Overwinter
                #'-nuparams=76b809bb:203', # Sapling
                #'-experimentalfeatures', '-zmergetoaddress',
                # Listen on this test's PID-keyed port so the nodes can actually peer with
                # each other, and only ever dial each other -- never the public network.
                '-listen=1',
                '-bind=127.0.0.1',
                '-port=%d' % p2p_port(i),
            ] + ['-connect=127.0.0.1:%d' % p2p_port(j) for j in range(4) if j != i]
            nodes.append(start_node(i, self.options.tmpdir, extra_args, binary=binary))
        return nodes

    def setup_network(self, split=False):
        super(WalletSaplingTest, self).setup_network(split)
        self._generate_starting_chain()

    def _generate_starting_chain(self):
        # Equivalent of test_framework.util.initialize_chain(): a 200-block chain where each
        # of the 4 nodes mined 25 blocks twice, so every node holds 25 mature and 25 immature
        # coinbases. Block timestamps are 10 minutes apart starting 1 Jan 2014, as there.
        block_time = 1388534400
        for _round in range(2):
            for peer in range(4):
                for _j in range(25):
                    set_node_times(self.nodes, block_time)
                    self.nodes[peer].generate(1)
                    block_time += 10 * 75
                # Must sync before next peer starts generating blocks
                sync_blocks(self.nodes)
        # Drop back to wall-clock time: initialize_chain() stops the cache nodes and the test
        # then runs against freshly started nodes that have no mocktime set.
        set_node_times(self.nodes, 0)
        self.sync_all()

    def run_test(self):
        # Sanity-check the test harness
        assert_equal(self.nodes[0].getblockcount(), 200)

        # Activate Overwinter
        self.nodes[2].generate(1)
        self.sync_all()

        self.nodes[2].generate(2)
        self.sync_all()

        taddr0 = self.nodes[0].getnewaddress()
        # Skip over the address containing node 1's coinbase
        self.nodes[1].getnewaddress()
        taddr1 = self.nodes[1].getnewaddress()
        saplingAddr0 = self.nodes[0].z_getnewaddress('sapling')
        saplingAddr1 = self.nodes[1].z_getnewaddress('sapling')

        # Verify addresses
        assert(saplingAddr0 in self.nodes[0].z_listaddresses())
        assert(saplingAddr1 in self.nodes[1].z_listaddresses())
        assert_equal(self.nodes[0].z_validateaddress(saplingAddr0)['type'], 'sapling')
        assert_equal(self.nodes[0].z_validateaddress(saplingAddr1)['type'], 'sapling')

        # Verify balance
        assert_equal(self.nodes[0].z_getbalance(saplingAddr0), Decimal('0'))
        assert_equal(self.nodes[1].z_getbalance(saplingAddr1), Decimal('0'))
        assert_equal(self.nodes[1].z_getbalance(taddr1), Decimal('0'))

        # Node 0 shields some funds
        # taddr -> Sapling
        #       -> taddr (change)
        recipients = []
        recipients.append({"address": saplingAddr0, "amount": Decimal('20')})
        myopid = self.nodes[0].z_sendmany(taddr0, recipients, 1, 0)
        mytxid = wait_and_assert_operationid_status(self.nodes[0], myopid)

        self.sync_all()

        # Verify priority of tx is MAX_PRIORITY, defined as 1E+16 (10000000000000000)
        mempool = self.nodes[0].getrawmempool(True)
        assert(Decimal(mempool[mytxid]['startingpriority']) == Decimal('1E+16'))

        self.nodes[2].generate(1)
        self.sync_all()

        # Verify balance
        assert_equal(self.nodes[0].z_getbalance(saplingAddr0), Decimal('20'))
        assert_equal(self.nodes[1].z_getbalance(saplingAddr1), Decimal('0'))
        assert_equal(self.nodes[1].z_getbalance(taddr1), Decimal('0'))

        # Node 0 sends some shielded funds to node 1
        # Sapling -> Sapling
        #         -> Sapling (change)
        recipients = []
        recipients.append({"address": saplingAddr1, "amount": Decimal('15')})
        myopid = self.nodes[0].z_sendmany(saplingAddr0, recipients, 1, 0)
        mytxid = wait_and_assert_operationid_status(self.nodes[0], myopid)

        self.sync_all()

        # Verify priority of tx is MAX_PRIORITY, defined as 1E+16 (10000000000000000)
        mempool = self.nodes[0].getrawmempool(True)
        assert(Decimal(mempool[mytxid]['startingpriority']) == Decimal('1E+16'))

        self.nodes[2].generate(1)
        self.sync_all()

        # Verify balance
        assert_equal(self.nodes[0].z_getbalance(saplingAddr0), Decimal('5'))
        assert_equal(self.nodes[1].z_getbalance(saplingAddr1), Decimal('15'))
        assert_equal(self.nodes[1].z_getbalance(taddr1), Decimal('0'))

        # Node 1 sends some shielded funds to node 0, as well as unshielding
        # Sapling -> Sapling
        #         -> taddr
        #         -> Sapling (change)
        recipients = []
        recipients.append({"address": saplingAddr0, "amount": Decimal('5')})
        recipients.append({"address": taddr1, "amount": Decimal('5')})
        myopid = self.nodes[1].z_sendmany(saplingAddr1, recipients, 1, 0)
        mytxid = wait_and_assert_operationid_status(self.nodes[1], myopid)

        self.sync_all()

        # Verify priority of tx is MAX_PRIORITY, defined as 1E+16 (10000000000000000)
        mempool = self.nodes[1].getrawmempool(True)
        assert(Decimal(mempool[mytxid]['startingpriority']) == Decimal('1E+16'))

        self.nodes[2].generate(1)
        self.sync_all()

        # Verify balance
        assert_equal(self.nodes[0].z_getbalance(saplingAddr0), Decimal('10'))
        assert_equal(self.nodes[1].z_getbalance(saplingAddr1), Decimal('5'))
        assert_equal(self.nodes[1].z_getbalance(taddr1), Decimal('5'))

        # Verify existence of Sapling related JSON fields
        resp = self.nodes[0].getrawtransaction(mytxid, 1)
        assert_equal(resp['valueBalance'], Decimal('5'))
        assert(len(resp['vShieldedSpend']) == 1)
        assert(len(resp['vShieldedOutput']) == 2)
        assert('bindingSig' in resp)
        shieldedSpend = resp['vShieldedSpend'][0]
        assert('cv' in shieldedSpend)
        assert('anchor' in shieldedSpend)
        assert('nullifier' in shieldedSpend)
        assert('rk' in shieldedSpend)
        assert('proof' in shieldedSpend)
        assert('spendAuthSig' in shieldedSpend)
        shieldedOutput = resp['vShieldedOutput'][0]
        assert('cv' in shieldedOutput)
        assert('cmu' in shieldedOutput)
        assert('ephemeralKey' in shieldedOutput)
        assert('encCiphertext' in shieldedOutput)
        assert('outCiphertext' in shieldedOutput)
        assert('proof' in shieldedOutput)

        # Verify importing a spending key will update the nullifiers and witnesses correctly
        sk0 = self.nodes[0].z_exportkey(saplingAddr0)
        self.nodes[2].z_importkey(sk0, "yes")
        assert_equal(self.nodes[2].z_getbalance(saplingAddr0), Decimal('10'))
        sk1 = self.nodes[1].z_exportkey(saplingAddr1)
        self.nodes[2].z_importkey(sk1, "yes")
        assert_equal(self.nodes[2].z_getbalance(saplingAddr1), Decimal('5'))

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