#!/usr/bin/env python3
"""Verify an isolated SAAVECORD chain with two nodes and two temporary wallets."""
import argparse
import base64
from decimal import Decimal
import hashlib
import http.client
import ipaddress
import json
import os
from pathlib import Path
import signal
import socket
import struct
import subprocess
import tempfile
import time

GENESIS = {
    'saavecord': '000000d5834ff77f43e6d8b2e5f1811bd47e0ed5ac486e49b57dbf9803eb0550',
    'saavecord-dev': '0b8ef792aba3b324482eb6ddc693866412281e155a02578152b668f7b655219a',
}

FOUNDER_ACTIVATION_HEIGHT = 200
FOUNDER_SCRIPT = bytes.fromhex('00143cee3282e1da374d6d0e70c9bd5d5ed03e6e3199')
FOUNDER_RULE = 'saavecord-founder-v1'
MAX_MONEY = 21000000 * 100000000


class RpcError(RuntimeError):
    def __init__(self, value):
        self.value = value
        super().__init__(str(value))


def free_port():
    with socket.socket() as sock:
        sock.bind(('127.0.0.1', 0))
        return sock.getsockname()[1]


def hash256(data):
    return hashlib.sha256(hashlib.sha256(data).digest()).digest()


def compact_size(value):
    if value < 253:
        return bytes([value])
    if value <= 0xffff:
        return b'\xfd' + struct.pack('<H', value)
    if value <= 0xffffffff:
        return b'\xfe' + struct.pack('<I', value)
    return b'\xff' + struct.pack('<Q', value)


def height_script(height):
    if 1 <= height <= 16:
        return bytes([0x50 + height])
    data = height.to_bytes((height.bit_length() + 7) // 8, 'little')
    if data[-1] & 0x80:
        data += b'\0'
    return bytes([len(data)]) + data


def scheduled_subsidy(height):
    if type(height) is not int or not 1 <= height <= 2147483647:
        raise ValueError('Invalid template height')
    halvings = height // 210000
    return (50 * 100000000) >> halvings if halvings < 64 else 0


def founder_payment(template, chain):
    """Validate the fixed SAAVECORD policy and return its active payment.

    GBT coinbasevalue remains the total subsidy plus fees. The founder amount
    is derived solely from the subsidy; fees belong entirely to the miner.
    Legacy public nodes are usable only before the activation height.
    """
    if chain != 'saavecord':
        return 0, None
    height = template.get('height')
    subsidy = scheduled_subsidy(height)
    total = template.get('coinbasevalue')
    if type(total) is not int or not subsidy <= total <= MAX_MONEY:
        raise ValueError('Invalid total coinbase value')
    active = height >= FOUNDER_ACTIVATION_HEIGHT
    metadata = template.get('founder_payment')
    if metadata is None:
        if active:
            raise ValueError('Active founder payment metadata is required')
        return 0, None
    if not isinstance(metadata, dict):
        raise ValueError('Invalid founder payment metadata')
    if type(metadata.get('active')) is not bool or metadata['active'] != active:
        raise ValueError('Unexpected founder payment activation state')
    if type(metadata.get('activation_height')) is not int or metadata['activation_height'] != FOUNDER_ACTIVATION_HEIGHT:
        raise ValueError('Unexpected founder payment activation height')
    if metadata.get('script') != FOUNDER_SCRIPT.hex():
        raise ValueError('Unexpected founder payment recipient')
    if type(metadata.get('subsidy_sat')) is not int or metadata['subsidy_sat'] != subsidy:
        raise ValueError('Unexpected founder payment subsidy')
    expected = subsidy // 5 if active else 0
    if type(metadata.get('amount_sat')) is not int or metadata['amount_sat'] != expected:
        raise ValueError('Unexpected founder payment amount')
    rules = template.get('rules', [])
    if active and (not isinstance(rules, list) or ('!' + FOUNDER_RULE) not in rules):
        raise ValueError('Active founder payment rule is required')
    return expected, FOUNDER_SCRIPT if expected else None


def mining_template(node, output_script, extra_nonce):
    chain = getattr(node, 'chain', None)
    if not isinstance(chain, str) or not chain:
        raise ValueError('Mining requires an explicit node.chain')
    requested_rules = ['segwit'] + ([FOUNDER_RULE] if chain == 'saavecord' else [])
    template = node.call('getblocktemplate', [{'rules': requested_rules}])
    owner_amount, owner_script = founder_payment(template, chain)
    script = height_script(template['height']) + b'\x08' + struct.pack('<Q', extra_nonce) + b'/SAAVECORD/'
    vin = b'\x01' + b'\0' * 32 + b'\xff' * 4 + compact_size(len(script)) + script + b'\xff' * 4
    outputs = [(template['coinbasevalue'] - owner_amount, output_script)]
    if owner_amount:
        # Keep the allocation visible as a separate output even if both
        # scripts are identical. Subtraction above prevents extra issuance.
        outputs.append((owner_amount, owner_script))
    commitment = template.get('default_witness_commitment')
    if commitment:
        outputs.append((0, bytes.fromhex(commitment)))
    vout = compact_size(len(outputs)) + b''.join(struct.pack('<Q', value) + compact_size(len(data)) + data for value, data in outputs)
    version = struct.pack('<I', 2)
    locktime = b'\0' * 4
    coinbase_base = version + vin + vout + locktime
    coinbase = (version + b'\0\x01' + vin + vout + b'\x01\x20' + b'\0' * 32 + locktime) if commitment else coinbase_base
    transactions = template['transactions']
    layer = [hash256(coinbase_base)] + [bytes.fromhex(tx['txid'])[::-1] for tx in transactions]
    while len(layer) > 1:
        if len(layer) % 2:
            layer.append(layer[-1])
        layer = [hash256(layer[i] + layer[i + 1]) for i in range(0, len(layer), 2)]
    header = (struct.pack('<I', template['version'] & 0xffffffff)
              + bytes.fromhex(template['previousblockhash'])[::-1] + layer[0]
              + struct.pack('<III', template['curtime'], int(template['bits'], 16), 0))
    body = compact_size(len(transactions) + 1) + coinbase + b''.join(bytes.fromhex(tx['data']) for tx in transactions)
    return header, body, template


def wait_for(check, timeout=30):
    end = time.monotonic() + timeout
    while time.monotonic() < end:
        try:
            value = check()
            if value:
                return value
        except (OSError, RpcError, FileNotFoundError, http.client.HTTPException):
            pass
        time.sleep(0.1)
    raise TimeoutError('Condition did not become true within the time limit')


class Node:
    def __init__(self, binary, chain, directory, index):
        self.binary = binary
        self.chain = chain
        self.directory = directory
        self.index = index
        self.rpc_port = free_port()
        self.p2p_port = free_port()
        self.process = None
        self.stdout_file = None
        self.exit_codes = []
        self.stop_errors = []
        self.shutdown_log_checks = []
        self.stopped = True
        self.log_offset = 0

    def call(self, method, params=None, wallet=None):
        cookie = (self.directory / self.chain / '.cookie').read_text().strip()
        auth = base64.b64encode(cookie.encode()).decode()
        connection = http.client.HTTPConnection('127.0.0.1', self.rpc_port, timeout=45)
        try:
            path = '/' if wallet is None else '/wallet/' + wallet
            body = json.dumps({'jsonrpc': '2.0', 'id': 'private-check', 'method': method, 'params': params or []})
            connection.request('POST', path, body, {'Authorization': 'Basic ' + auth, 'Content-Type': 'application/json'})
            response = connection.getresponse()
            data = json.loads(response.read(), parse_float=Decimal)
            if data.get('error'):
                raise RpcError(data['error'])
            return data['result']
        finally:
            connection.close()

    def start(self):
        self.directory.mkdir(mode=0o700, parents=True, exist_ok=True)
        debug_log = self.directory / self.chain / 'debug.log'
        self.log_offset = debug_log.stat().st_size if debug_log.exists() else 0
        self.stdout_file = (self.directory / 'process.log').open('a')
        command = [str(self.binary), '-chain=' + self.chain, '-datadir=' + str(self.directory),
                   '-server=1', '-daemon=0', '-listen=1', '-connect=0', '-dnsseed=0',
                   '-fixedseeds=0', '-discover=0', '-listenonion=0', '-natpmp=0', '-dbcache=64', '-par=1',
                   '-bind=127.0.0.1:' + str(self.p2p_port), '-port=' + str(self.p2p_port),
                   '-rpcbind=127.0.0.1', '-rpcallowip=127.0.0.1', '-rpcport=' + str(self.rpc_port),
                   '-fallbackfee=0.0001', '-maxconnections=32']
        self.process = subprocess.Popen(command, stdout=self.stdout_file, stderr=subprocess.STDOUT,
                                        preexec_fn=lambda: os.nice(10))
        self.stopped = False

        def ready():
            if self.process.poll() is not None:
                raise RuntimeError('Node exited before RPC became available; see ' + str(self.directory / 'process.log'))
            return self.call('getblockchaininfo')
        wait_for(ready)

    def stop(self):
        if self.process is None or self.stopped:
            return
        try:
            if self.process.poll() is None:
                try:
                    self.call('stop')
                except Exception as exc:
                    self.stop_errors.append('RPC stop failed: ' + str(exc))
                    self.process.terminate()
            try:
                code = self.process.wait(timeout=30)
            except subprocess.TimeoutExpired:
                self.stop_errors.append('Graceful shutdown timed out')
                self.process.terminate()
                try:
                    code = self.process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    self.stop_errors.append('Termination timed out; stopped own child process')
                    self.process.kill()
                    code = self.process.wait(timeout=10)
            self.exit_codes.append(code)
        except Exception as exc:
            self.stop_errors.append(str(exc))
        finally:
            try:
                if self.stdout_file:
                    self.stdout_file.close()
                debug_log = self.directory / self.chain / 'debug.log'
                if debug_log.exists():
                    with debug_log.open('rb') as handle:
                        handle.seek(self.log_offset)
                        self.shutdown_log_checks.append(b'Shutdown done' in handle.read())
                else:
                    self.shutdown_log_checks.append(False)
            except OSError as exc:
                self.stop_errors.append('Shutdown log check failed: ' + str(exc))
                self.shutdown_log_checks.append(False)
            self.stopped = self.process.poll() is not None

    def listeners(self):
        inodes = set()
        for entry in Path('/proc', str(self.process.pid), 'fd').iterdir():
            try:
                target = os.readlink(entry)
            except OSError:
                continue
            if target.startswith('socket:['):
                inodes.add(target[8:-1])
        found = []
        for table in ('tcp', 'tcp6'):
            for line in Path('/proc/net', table).read_text().splitlines()[1:]:
                fields = line.split()
                if fields[3] != '0A' or fields[9] not in inodes:
                    continue
                host, port = fields[1].split(':')
                raw = bytes.fromhex(host)
                raw = raw[::-1] if table == 'tcp' else b''.join(raw[i:i + 4][::-1] for i in range(0, 16, 4))
                address = ipaddress.ip_address(raw)
                assert address.is_loopback, 'Node has a non-loopback listener'
                found.append({'address': str(address), 'port': int(port, 16)})
        assert len(found) >= 2, 'Expected local RPC and P2P listeners'
        return found


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bin-dir', type=Path, required=True)
    parser.add_argument('--chain', choices=GENESIS, default='saavecord-dev')
    parser.add_argument('--out-dir', type=Path, required=True)
    parser.add_argument('--mining-seconds', type=int, default=300)
    parser.add_argument('--max-hash-attempts', type=int, default=2400000000)
    parser.add_argument('--benchmark', action='store_true')
    parser.add_argument('--grinder', type=Path, help='Optional bounded local header grinder')
    parser.add_argument('--threads', type=int, choices=(1, 2), default=2)
    args = parser.parse_args()
    args.out_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
    binary = (args.bin_dir / 'bitcoind').resolve()
    assert binary.is_file(), binary
    report = {'chain': args.chain, 'binary': str(binary), 'started_utc': time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime()),
              'binary_sha256': hashlib.sha256(binary.read_bytes()).hexdigest(), 'ok': False}
    report['mining_method'] = 'getblocktemplate/submitblock' if args.grinder else 'generatetoaddress'
    if args.grinder:
        report['grinder_sha256'] = hashlib.sha256(args.grinder.read_bytes()).hexdigest()
        report['mining_threads'] = args.threads
    root = Path(tempfile.mkdtemp(prefix='saavecord-two-node-'))
    report['temporary_data'] = str(root)
    nodes = [Node(binary, args.chain, root / ('node' + str(i)), i) for i in range(2)]
    def interrupted(signum, _frame):
        raise InterruptedError('Verifier interrupted by signal ' + str(signum))
    signal.signal(signal.SIGTERM, interrupted)
    signal.signal(signal.SIGINT, interrupted)
    started = time.monotonic()
    try:
        for node in nodes:
            node.start()
            assert node.call('getblockchaininfo')['chain'] == args.chain
            assert node.call('getblockhash', [0]) == GENESIS[args.chain]
        report['listeners'] = [node.listeners() for node in nodes]
        nodes[0].call('addnode', ['127.0.0.1:' + str(nodes[1].p2p_port), 'onetry'])
        wait_for(lambda: nodes[0].call('getconnectioncount') > 0 and nodes[1].call('getconnectioncount') > 0)
        for node in nodes:
            assert all(peer['addr'].startswith('127.0.0.1:') for peer in node.call('getpeerinfo'))
        report['local_peer_connection'] = True
        print(json.dumps({'stage': 'two_local_nodes_connected', 'chain': args.chain}), flush=True)
        nodes[0].call('createwallet', {'wallet_name': 'sender', 'load_on_startup': True})
        nodes[1].call('createwallet', {'wallet_name': 'receiver', 'load_on_startup': True})
        mining_address = nodes[0].call('getnewaddress', ['', 'bech32'], wallet='sender')
        output_script = bytes.fromhex(nodes[0].call('getaddressinfo', [mining_address], wallet='sender')['scriptPubKey'])
        receiving_address = nodes[1].call('getnewaddress', ['', 'bech32'], wallet='receiver')
        hrp = 'scdev1' if args.chain == 'saavecord-dev' else 'sc1'
        assert mining_address.startswith(hrp) and receiving_address.startswith(hrp)
        assert nodes[1].call('getaddressinfo', [receiving_address], wallet='receiver')['iswitness']
        report['address_hrp'] = hrp
        attempts = 0
        mine_start = time.monotonic()
        reported_heights = set()
        template_counter = 0

        def mine_one(limit):
            nonlocal template_counter
            if args.grinder:
                template_counter += 1
                header, body, template = mining_template(nodes[0], output_script, template_counter)
                remaining = args.mining_seconds - (time.monotonic() - mine_start)
                if remaining <= 0:
                    raise TimeoutError('Mining exceeded the configured time budget')
                result = subprocess.run([str(args.grinder.resolve()), '--solve', header.hex(), template['target'],
                                         str(limit), str(args.threads)], capture_output=True, text=True,
                                        timeout=remaining, preexec_fn=lambda: os.nice(10))
                if result.returncode not in (0, 2):
                    raise RuntimeError('Header grinder failed: ' + result.stderr)
                solved = json.loads(result.stdout)
                if solved['found']:
                    solved_header = bytes.fromhex(solved['header'])
                    assert len(solved_header) == 80 and solved_header[:76] == header[:76]
                    actual_hash = hash256(solved_header)[::-1].hex()
                    assert actual_hash == solved['hash']
                    assert int(actual_hash, 16) <= int(template['target'], 16)
                    rejection = nodes[0].call('submitblock', [(solved_header + body).hex()])
                    assert rejection is None, rejection
                    assert nodes[0].call('getblockhash', [template['height']]) == actual_hash
                return solved['hashes']
            destination = (nodes[0].call('getnewaddress', ['', 'bech32'], wallet='sender')
                           if args.chain == 'saavecord' else mining_address)
            nodes[0].call('generatetoaddress', [1, destination, limit])
            return limit

        def mine_to(height):
            nonlocal attempts
            while nodes[0].call('getblockcount') < height:
                if time.monotonic() - mine_start > args.mining_seconds:
                    raise TimeoutError('Mining exceeded the configured time budget')
                remaining_attempts = args.max_hash_attempts - attempts
                if remaining_attempts <= 0:
                    raise TimeoutError('Mining exhausted the configured hash-attempt budget')
                limit = min(remaining_attempts, 2**32 if args.grinder else 1000000)
                attempts += mine_one(limit)
                current = nodes[0].call('getblockcount')
                progress = {'height': current, 'mining_seconds': time.monotonic() - mine_start,
                            'hash_attempts': attempts, 'chain': args.chain, 'temporary_data': str(root)}
                temporary_progress = args.out_dir / 'progress.json.tmp'
                temporary_progress.write_text(json.dumps(progress))
                temporary_progress.replace(args.out_dir / 'progress.json')
                if current and current % 10 == 0 and current not in reported_heights:
                    reported_heights.add(current)
                    print(json.dumps({'height': current, 'mining_seconds': round(time.monotonic() - mine_start, 2)}), flush=True)

        if args.benchmark:
            before = time.monotonic()
            hashes = mine_one(1000000)
            report['benchmark'] = {'requested_max_hashes': 1000000, 'elapsed_seconds': time.monotonic() - before,
                                   'height': nodes[0].call('getblockcount'), 'hashes': hashes}
            report['ok'] = True
        else:
            mine_to(100)
            boundary = nodes[0].call('getbalances', wallet='sender')['mine']
            assert boundary['trusted'] == Decimal('0') and boundary['immature'] == Decimal('5000'), boundary
            mine_to(101)
            wait_for(lambda: nodes[1].call('getblockcount') == 101)
            before = nodes[0].call('getbalances', wallet='sender')['mine']
            assert before['trusted'] == Decimal('50'), before
            assert before['immature'] == Decimal('5000'), before
            assert nodes[1].call('getbalance', wallet='receiver') == 0
            txid = nodes[0].call('sendtoaddress', [receiving_address, 10], wallet='sender')
            wait_for(lambda: txid in nodes[1].call('getrawmempool'))
            sender_tx = nodes[0].call('gettransaction', [txid], wallet='sender')
            receiver_tx = wait_for(lambda: nodes[1].call('gettransaction', [txid], wallet='receiver'))
            assert sender_tx['confirmations'] == receiver_tx['confirmations'] == 0
            fee = -sender_tx['fee']
            assert fee > 0
            mine_to(102)
            wait_for(lambda: nodes[1].call('getblockcount') == 102)
            sender_tx = nodes[0].call('gettransaction', [txid], wallet='sender')
            receiver_tx = nodes[1].call('gettransaction', [txid], wallet='receiver')
            assert sender_tx['confirmations'] == receiver_tx['confirmations'] == 1
            assert sender_tx['blockhash'] == receiver_tx['blockhash']
            balances_a = nodes[0].call('getbalances', wallet='sender')['mine']
            balances_b = nodes[1].call('getbalances', wallet='receiver')['mine']
            assert balances_a['trusted'] == Decimal('90') - fee, balances_a
            assert balances_a['immature'] == Decimal('5000') + fee, balances_a
            assert balances_b['trusted'] == Decimal('10'), balances_b
            report.update({'height': 102, 'transaction_id': txid, 'confirmation_block': sender_tx['blockhash'],
                           'confirmations': 1, 'sent_sc': '10', 'fee_sc': str(fee),
                           'sender_balances': balances_a, 'receiver_balances': balances_b,
                           'attempt_budget_upper_bound': attempts, 'mining_seconds': time.monotonic() - mine_start})
            for node in reversed(nodes):
                node.stop()
                assert node.exit_codes[-1] == 0
            for node in nodes:
                node.start()
                assert node.call('getblockcount') == 102
                assert node.call('getblockhash', [102]) == report['confirmation_block']
            report['restart_listeners'] = [node.listeners() for node in nodes]
            assert nodes[1].call('getbalance', wallet='receiver') == Decimal('10')
            assert nodes[0].call('getbalance', wallet='sender') == Decimal('90') - fee
            report['wallet_and_chain_restart_verified'] = True
            report['ok'] = True
    except Exception as exc:
        report['error'] = str(exc)
    finally:
        for node in reversed(nodes):
            node.stop()
        report['exit_codes'] = [node.exit_codes for node in nodes]
        report['stop_errors'] = [node.stop_errors for node in nodes]
        report['clean_shutdown'] = (all(codes and all(code == 0 for code in codes) for codes in report['exit_codes'])
                                    and not any(report['stop_errors']))
        report['processes_stopped'] = all(node.process is not None and node.process.poll() is not None for node in nodes)
        report['shutdown_log_confirmed'] = all(node.shutdown_log_checks and all(node.shutdown_log_checks) for node in nodes)
        report['shutdown_log_checks'] = [node.shutdown_log_checks for node in nodes]
        report['elapsed_seconds'] = time.monotonic() - started
        if not (report['clean_shutdown'] and report['processes_stopped'] and report['shutdown_log_confirmed']):
            report['ok'] = False
        target = args.out_dir / ('benchmark.json' if args.benchmark else 'verification.json')
        target.write_text(json.dumps(report, indent=2, default=str) + '\n')
        print(json.dumps({'ok': report['ok'], 'report': str(target), 'error': report.get('error')}), flush=True)
    return 0 if report['ok'] else 1


if __name__ == '__main__':
    raise SystemExit(main())