#!/usr/bin/env python3
# Copyright (c) 2024-present The Limenka developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test utxo-to-sqlite conversion tool"""
import os
try:
    import sqlite3
except ImportError:
    pass
import platform
import subprocess
import sys

from test_framework.key import ECKey
from test_framework.messages import (
    COutPoint,
    CTxOut,
)
from test_framework.crypto.muhash import MuHash3072
from test_framework.script import (
    CScript,
    CScriptOp,
)
from test_framework.script_util import (
    PAY_TO_ANCHOR,
    key_to_p2pk_script,
    key_to_p2pkh_script,
    key_to_p2wpkh_script,
    keys_to_multisig_script,
    output_key_to_p2tr_script,
    script_to_p2sh_script,
    script_to_p2wsh_script,
)
from test_framework.test_framework import LimenkaTestFramework
from test_framework.util import (
    assert_equal,
)
from test_framework.wallet import MiniWallet


def calculate_muhash_from_sqlite_utxos(filename):
    muhash = MuHash3072()
    con = sqlite3.connect(filename)
    cur = con.cursor()
    for (txid_hex, vout, value, coinbase, height, spk_hex) in cur.execute("SELECT * FROM utxos"):
        # serialize UTXO for MuHash (see function `TxOutSer` in the  coinstats module)
        utxo_ser = COutPoint(int(txid_hex, 16), vout).serialize()
        utxo_ser += (height * 2 + coinbase).to_bytes(4, 'little')
        utxo_ser += CTxOut(value, bytes.fromhex(spk_hex)).serialize()
        muhash.insert(utxo_ser)
    con.close()
    return muhash.digest()[::-1].hex()


class UtxoToSqliteTest(LimenkaTestFramework):
    def set_test_params(self):
        self.num_nodes = 1
        # we want to create some UTXOs with non-standard output scripts
        self.extra_args = [['-acceptnonstdtxn=1']]

    def skip_test_if_missing_module(self):
        self.skip_if_no_py_sqlite3()

    def run_test(self):
        node = self.nodes[0]
        wallet = MiniWallet(node)
        key = ECKey()

        self.log.info('Test that oversized output scripts are rejected')
        key.generate(compressed=False)
        uncompressed_pubkey = key.get_pubkey().get_bytes()
        key.generate(compressed=True)
        pubkey = key.get_pubkey().get_bytes()

        # Test that scripts exceeding MAX_OUTPUT_SCRIPT_SIZE=34 are rejected
        invalid_scripts = [
            (key_to_p2pk_script(pubkey), "P2PK compressed (35 bytes)"),
            (key_to_p2pk_script(uncompressed_pubkey), "P2PK uncompressed (67 bytes)"),
            (keys_to_multisig_script([pubkey]), "Bare multisig 1-of-1 (37 bytes)"),
            (keys_to_multisig_script([uncompressed_pubkey]*2), "Bare multisig 2-of-2 uncompressed"),
            (CScript([CScriptOp.encode_op_n(1)]*1000), "Large script (1000 bytes)"),
        ]

        for script, description in invalid_scripts:
            try:
                wallet.send_to(from_node=node, scriptPubKey=script, amount=1, fee=20000)
                raise AssertionError(f"{description} should have been rejected")
            except Exception as e:
                assert 'bad-txns-vout-script-toolarge' in str(e), \
                    f"{description} rejected with wrong error: {e}"
                self.log.info(f"  ✓ {description} correctly rejected")

        self.log.info('Create UTXOs with valid output script types (≤34 bytes)')
        for i in range(1, 10+1):
            key.generate(compressed=True)
            pubkey = key.get_pubkey().get_bytes()

            # Only include output scripts that comply with MAX_OUTPUT_SCRIPT_SIZE=34
            output_scripts = (
                key_to_p2pkh_script(pubkey),                        # 25 bytes
                script_to_p2sh_script(key_to_p2pkh_script(pubkey)), # 23 bytes
                key_to_p2wpkh_script(pubkey),                       # 22 bytes
                script_to_p2wsh_script(key_to_p2pkh_script(pubkey)),# 34 bytes
                output_key_to_p2tr_script(pubkey[1:]),              # 34 bytes
                PAY_TO_ANCHOR,                                      # 4 bytes
            )

            # create outputs and mine them in a block
            for output_script in output_scripts:
                wallet.send_to(from_node=node, scriptPubKey=output_script, amount=i, fee=20000)
            self.generate(wallet, 1)

        self.log.info('Dump UTXO set via `dumptxoutset` RPC')
        input_filename = os.path.join(self.options.tmpdir, "utxos.dat")
        node.dumptxoutset(input_filename, "latest")

        self.log.info('Convert UTXO set from compact-serialized format to sqlite format')
        output_filename = os.path.join(self.options.tmpdir, "utxos.sqlite")
        base_dir = self.config["environment"]["SRCDIR"]
        utxo_to_sqlite_path = os.path.join(base_dir, "contrib", "utxo-tools", "utxo_to_sqlite.py")
        subprocess.run([sys.executable, utxo_to_sqlite_path, input_filename, output_filename],
                       check=True, stderr=subprocess.STDOUT)

        self.log.info('Verify that both UTXO sets match by comparing their MuHash')
        muhash_sqlite = calculate_muhash_from_sqlite_utxos(output_filename)
        muhash_compact_serialized = node.gettxoutsetinfo('muhash')['muhash']
        assert_equal(muhash_sqlite, muhash_compact_serialized)

        if platform.system() != "Windows":  # FIFOs are not available on Windows
            self.log.info('Convert UTXO set directly (without intermediate dump) via named pipe')
            fifo_filename = os.path.join(self.options.tmpdir, "utxos.fifo")
            os.mkfifo(fifo_filename)
            output_direct_filename = os.path.join(self.options.tmpdir, "utxos_direct.sqlite")
            p = subprocess.Popen([sys.executable, utxo_to_sqlite_path, fifo_filename, output_direct_filename],
                                 stderr=subprocess.STDOUT)
            node.dumptxoutset(fifo_filename, "latest")
            p.wait(timeout=10)
            muhash_direct_sqlite = calculate_muhash_from_sqlite_utxos(output_direct_filename)
            assert_equal(muhash_sqlite, muhash_direct_sqlite)
            os.remove(fifo_filename)


if __name__ == "__main__":
    UtxoToSqliteTest(__file__).main()