#!/usr/bin/env python3

import click
import subprocess
import threading
import os
from random import *
import time
import json
from web3 import Web3

colors = {
    "red": 91,
    "green": 92,
    "yellow": 93,
    "lightPurple": 94,
    "purple": 95,
    "cyan": 96,
    "lightGray": 97,
}
worker_colors = {}
chain_id = 0
faucet_private_key = ''
faucet_public_address = ''
boot_nodes_public_addresses = {}

data_dir = 'local-net-data'
if not os.path.isdir(f'./{data_dir}'):
    os.mkdir(f'./{data_dir}')

local_net = os.path.realpath(f'./{data_dir}')
taraxa_cli = os.path.dirname(os.path.realpath(__file__)) + '/taraxa'


@click.group()
def cli():
    """Local Net CLI"""
    pass


@cli.command()
@click.option('--boot-nodes', default=3, help='Number of boot nodes')
@click.option('--rpc-nodes', default=3, help='Number of RPC nodes')
@click.option('--template', default='devnet', help='Name of config template to use as base')
@click.option('--tps', default=0, help='Number of transactions per second (if zero the faucet will not start)')
@click.argument('binary')
def start(boot_nodes, rpc_nodes, template, tps, binary):
    """Start a local testnet"""
    binary = os.path.realpath(binary)

    config(template, boot_nodes, rpc_nodes)
    start_workers(binary, boot_nodes, rpc_nodes, tps)


def config(network, boot_nodes, rpc_nodes):
    global boot_nodes_public_addresses

    boot_nodes_flags = []
    private_keys = {}
    for i in range(boot_nodes):
        worker_name = f'boot-node-{i}'

        (private_key, public_key, public_address) = generate_wallet()

        private_keys[worker_name] = private_key[2:]
        boot_nodes_public_addresses[worker_name] = public_address

        boot_nodes_flags.append(
            f'--boot-node 127.0.0.1:10{i+1}02/{public_key[2:]}')

    boot_nodes_flags = " ".join(boot_nodes_flags)

    base_path = os.path.realpath(os.path.dirname(
        os.path.realpath(__file__)) + '/../')
    template = json.loads('{}')
    with open(f'{base_path}/config/networks/{network}.json') as json_file:
        template_data = json.load(json_file)
        template.update(template_data)

    generate_faucet_wallet()

    eligibility_balance_threshold = template['chain_config'][
        'final_chain']['state']['dpos']['eligibility_balance_threshold']
    genesis_state = template['chain_config']['final_chain']['state']['dpos']['genesis_state']
    genesis_balances = template['chain_config']['final_chain']['state']['genesis_balances']

    genesis_balances_keys = list(genesis_balances.keys())
    genesis_balances_faucet_address = genesis_balances_keys[0]
    genesis_balances_faucet_value = genesis_balances[genesis_balances_faucet_address]
    del genesis_balances[genesis_balances_faucet_address]
    genesis_balances[faucet_public_address[2:]] = genesis_balances_faucet_value

    for i in genesis_state:
        genesis_state[i] = {}
        for a in boot_nodes_public_addresses:
            genesis_state[i][boot_nodes_public_addresses[a]
                             ] = eligibility_balance_threshold

    timestamp = hex(int(time.time()))

    for type in ("boot", "rpc"):
        if type == "boot":
            nodes = boot_nodes
        else:
            nodes = rpc_nodes

        for i in range(nodes):
            if type == "boot":
                worker_name = f'boot-node-{i}'
            else:
                worker_name = f'rpc-node-{i}'

            config_name = f'./{data_dir}/conf-{worker_name}.json'

            if not os.path.isdir(f'{local_net}/db-{worker_name}'):
                os.mkdir(f'{local_net}/db-{worker_name}')

            cmd = ''
            cmd += f'{taraxa_cli} config node --network {network} '
            cmd += f'{boot_nodes_flags} '

            if type == "boot":
                cmd += '--as-boot-node '
                cmd += f'--node-secret {private_keys[worker_name]} '

            if type == "boot":
                port_prefix = '0'
            else:
                port_prefix = '1'

            cmd += f'--prop-override network_tcp_port:1{port_prefix}{i+1}02 --prop-override network_udp_port:1{port_prefix}{i+1}02 '
            cmd += f'--prop-override rpc.http_port:7{port_prefix}{i+1}7 --prop-override rpc.ws_port:8{port_prefix}{i+1}7 '
            cmd += f'--prop-override db_path:{local_net}/db-{worker_name} '
            cmd += f'--prop-override chain_config.chain_id:{chain_id} '
            cmd += f'--prop-override chain_config.dag_genesis_block.timestamp:{timestamp} '
            cmd += f'--prop-override chain_config.final_chain.genesis_block_fields.timestamp:{timestamp} '
            cmd += f'--file {config_name}'

            os.system(cmd)

            config = json.loads('{}')
            with open(config_name) as json_file:
                config_data = json.load(json_file)
                config.update(config_data)

            config['chain_config']['final_chain']['state']['dpos']['genesis_state'] = genesis_state
            config['chain_config']['final_chain']['state']['genesis_balances'] = genesis_balances

            with open(config_name, 'w') as json_file:
                json.dump(config, json_file, indent=2)


def generate_faucet_wallet():
    global chain_id
    global faucet_private_key
    global faucet_public_address

    wallet_path = f'./{data_dir}/faucet-wallet.json'

    if os.path.isfile(wallet_path):
        wallet = json.loads('{}')
        with open(wallet_path) as json_file:
            wallet_data = json.load(json_file)
            wallet.update(wallet_data)

        chain_id = wallet['chain_id']
        faucet_private_key = wallet['faucet_private_key']
        faucet_public_address = wallet['faucet_public_address']


    else:
        (faucet_private_key, _, faucet_public_address) = generate_wallet()
        chain_id = hex(randint(500, 1000))

        wallet = {
            "chain_id": chain_id,
            "faucet_private_key": faucet_private_key,
            "faucet_public_address": faucet_public_address,
        }
        with open(f'./{data_dir}/faucet-wallet.json', 'w') as json_file:
            json.dump(wallet, json_file, indent=2)


def generate_wallet():
    private_key = subprocess.check_output(
        f'{taraxa_cli} account | grep private', shell=True).strip().decode('utf-8').split(' ')[2]
    public_key = subprocess.check_output(
        f'{taraxa_cli} account --secret {private_key} | grep public', shell=True).strip().decode('utf-8').split(' ')[2]
    public_address = subprocess.check_output(
        f'{taraxa_cli} account --secret {private_key} | grep address:', shell=True).strip().decode('utf-8').split(' ')[1]

    return (private_key, public_key, public_address)


def start_workers(binary, boot_nodes, rpc_nodes, tps):
    c = list(colors.keys())
    for type in ("boot", "rpc"):
        if type == "boot":
            nodes = boot_nodes
        else:
            nodes = rpc_nodes

        for i in range(nodes):
            if type == "boot":
                worker_name = f'boot-node-{i}'
            else:
                worker_name = f'rpc-node-{i}'

            worker_colors[worker_name] = c.pop()
            thread = threading.Thread(
                target=node_worker, args=(binary, worker_name))
            thread.start()

    if tps > 0:
        worker_colors['faucet'] = c.pop()
        thread = threading.Thread(
            target=faucet_worker, args=(tps,))
        thread.start()


def node_worker(binary, worker_name):
    process = subprocess.Popen(f'{binary} --conf_taraxa ./{data_dir}/conf-{worker_name}.json',
                               shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
    logs(worker_name, process)


def faucet_worker(tps):
    # Wait for the nodes to start
    time.sleep(10)
    web3 = Web3(Web3.HTTPProvider('http://127.0.0.1:7017'))
    nonce = web3.eth.getTransactionCount(
        Web3.toChecksumAddress(faucet_public_address))

    boot_nodes = list(boot_nodes_public_addresses)

    while True:
        time.sleep(1/tps)
        to = boot_nodes_public_addresses[boot_nodes[randint(
            0, len(boot_nodes)-1)]]
        tx = {
            'nonce': nonce,
            'to': Web3.toChecksumAddress(to),
            'value': web3.toWei(100000000, 'gwei'),
            'gas': 21000,
            'gasPrice': web3.toWei(1, 'gwei'),
            'chainId': int(chain_id, 0)
        }
        nonce = nonce + 1
        signed_tx = web3.eth.account.sign_transaction(tx, faucet_private_key)

        current_time = time.localtime()
        t = time.strftime('%Y-%m-%d %H:%M:%S', current_time)

        try:
            tx_hash = web3.eth.sendRawTransaction(signed_tx.rawTransaction)
            log_format(
                'faucet', f'{t} Dripped to {to}, tx_hash: {web3.toHex(tx_hash)}')
        except:
            log_format('faucet', f'{t} Failed to drip to {to}')
            pass


def logs(name, process):
    while process.poll() is None:
        out = process.stdout.readline().strip().decode('utf-8')
        if out:
            log_format(name, out)


def log_format(name, log):
    worker_name = "\033[{}m[{}]\033[00m\t".format(
        colors[worker_colors[name]], name)
    print("{} {}" .format(worker_name, log))


if __name__ == '__main__':
    cli()
