import argparse
import hashlib
import itertools
import json
import logging
import time

from concurrent.futures import as_completed, ProcessPoolExecutor
from concurrent.futures._base import TimeoutError
from typing import Any, Dict, List, Optional, Literal, Union
from uuid import UUID

from moonstreamdb.blockchain import AvailableBlockchainType
from web3._utils.request import cache_session
from web3.middleware import geth_poa_middleware
from web3 import Web3, AsyncHTTPProvider


from mooncrawl.moonworm_crawler.crawler import _retry_connect_web3

from ..db import PrePing_SessionLocal
from ..settings import (
    INFURA_PROJECT_ID,
    NB_CONTROLLER_ACCESS_ID,
    infura_networks,
    multicall_contracts,
)
from .db import clean_labels, commit_session, view_call_to_label
from .Multicall2_interface import Contract as Multicall2
from .web3_util import FunctionSignature, connect

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def select_provider(
    web3_provider_uri: Optional[str],
    blockchain_type: AvailableBlockchainType,
    access_id: UUID,
) -> Web3:
    if web3_provider_uri is not None:
        try:
            logger.info(
                f"Connecting to blockchain: {blockchain_type} with custom provider!"
            )

            web3_client = connect(web3_provider_uri, request_kwargs={"timeout": 60})

            if blockchain_type != AvailableBlockchainType.ETHEREUM:
                web3_client.middleware_onion.inject(geth_poa_middleware, layer=0)
        except Exception as e:
            logger.error(
                f"Web3 connection to custom provider {web3_provider_uri} failed error: {e}"
            )
            raise (e)
    else:
        logger.info(f"Connecting to blockchain: {blockchain_type} with Node balancer.")
        web3_client = _retry_connect_web3(
            blockchain_type=blockchain_type, access_id=access_id
        )

    return web3_client


def make_multicall(
    web3_provider_uri: Optional[str],
    blockchain_type: AvailableBlockchainType,
    access_id: UUID,
    calls: List[Any],
    block_timestamp: int,
    block_number: Union[
        Literal["latest", "earliest", "pending"],
        int,
    ] = "latest",
    no_block_specifier: bool = False,
) -> Any:
    multicall_calls = []

    web3_client = select_provider(
        web3_provider_uri=web3_provider_uri,
        blockchain_type=blockchain_type,
        access_id=access_id,
    )

    multicall_contract = Multicall2(
        web3_client,
        web3_client.toChecksumAddress(multicall_contracts[blockchain_type]),
    )

    for call in calls:
        try:
            multicall_calls.append(
                (
                    call["address"],
                    call["method"].encode_data(call["inputs"]).hex(),
                )
            )
        except Exception as e:
            logger.error(
                f'Error encoding data for method {call["method"].name} call: {call}'
            )

    start_time = time.time()

    if no_block_specifier:
        """
        tryBlockAndAggregate returns (uint256 blockNumber, bytes32 blockHash, Result[] memory returnData)
        """
        multicall_method = multicall_contract.tryBlockAndAggregate
        multicall_result_with_block = multicall_method(
            False, calls=multicall_calls
        ).call()
        multicall_result = multicall_result_with_block[2]
        block_number = multicall_result_with_block[0]
        block_timestamp = web3_client.eth.get_block(block_number).timestamp  # type: ignore
    else:
        multicall_method = multicall_contract.tryAggregate
        multicall_result = multicall_method(False, calls=multicall_calls).call(
            block_identifier=block_number
        )

    print(f"Multicall2 call time: {time.time() - start_time}")

    results = []

    # Handle the case with not successful calls

    for index, encoded_data in enumerate(multicall_result):
        try:
            if encoded_data[0]:
                results.append(
                    {
                        "result": calls[index]["method"].decode_data(encoded_data[1]),
                        "hash": calls[index]["hash"],
                        "method": calls[index]["method"],
                        "address": calls[index]["address"],
                        "name": calls[index]["method"].name,
                        "inputs": calls[index]["inputs"],
                        "call_data": multicall_calls[index][1],
                        "block_number": block_number,
                        "block_timestamp": block_timestamp,
                        "status": encoded_data[0],
                    }
                )
            else:
                results.append(
                    {
                        "result": calls[index]["method"].decode_data(encoded_data[1]),
                        "hash": calls[index]["hash"],
                        "method": calls[index]["method"],
                        "address": calls[index]["address"],
                        "name": calls[index]["method"].name,
                        "inputs": calls[index]["inputs"],
                        "call_data": multicall_calls[index][1],
                        "block_number": block_number,
                        "block_timestamp": block_timestamp,
                        "status": encoded_data[0],
                    }
                )
        except Exception as e:
            results.append(
                {
                    "result": str(encoded_data[1]),
                    "hash": calls[index]["hash"],
                    "method": calls[index]["method"],
                    "address": calls[index]["address"],
                    "name": calls[index]["method"].name,
                    "inputs": calls[index]["inputs"],
                    "call_data": multicall_calls[index][1],
                    "block_number": block_number,
                    "block_timestamp": block_timestamp,
                    "status": encoded_data[0],
                    "error": str(e),
                }
            )

            logger.error(
                f"Error decoding data for for method {call['method'].name} call {calls[index]}: {e}."
            )
            # data is not decoded, return the encoded data
            logger.error(f"Encoded data: {encoded_data}")

    return results


def crawl_calls_level(
    access_id: UUID,
    web3_provider_uri: Optional[str],
    db_session,
    calls,
    responces,
    contracts_ABIs,
    interfaces,
    batch_size,
    block_number,
    blockchain_type,
    block_timestamp,
    max_workers,
    no_block_specifier,
    max_batch_size=5000,
    min_batch_size=4,
):
    calls_of_level = []

    for call in calls:
        if call["generated_hash"] in responces:
            continue
        parameters = []

        logger.info(f"Call: {json.dumps(call, indent=4)}")

        for input in call["inputs"]:
            if type(input["value"]) in (str, int):
                if input["value"] not in responces:
                    parameters.append([input["value"]])
                else:
                    if (
                        contracts_ABIs[call["address"]][input["value"]]["name"]
                        == "totalSupply"
                    ):  # hack for totalSupply TODO(Andrey): need add propper support for response parsing
                        print(responces[input["value"]][0])
                        parameters.append(
                            list(range(1, responces[input["value"]][0][0] + 1))
                        )
                    else:
                        parameters.append(responces[input["value"]])
            elif type(input["value"]) == list:
                parameters.append(input["value"])
            else:
                raise

        for call_parameters in itertools.product(*parameters):
            # hack for tuples product
            if len(call_parameters) == 1 and type(call_parameters[0]) == tuple:
                call_parameters = call_parameters[0]
            calls_of_level.append(
                {
                    "address": call["address"],
                    "method": FunctionSignature(
                        interfaces[call["address"]].get_function_by_name(call["name"])
                    ),
                    "hash": call["generated_hash"],
                    "inputs": call_parameters,
                }
            )

    retry = 0

    while len(calls_of_level) > 0:
        make_multicall_result = []

        try:
            # make multicall calls in batches with split on max_workers threads

            futures = []

            calls_for_workers = []

            faild_calls = []

            with ProcessPoolExecutor(max_workers=max_workers) as executor:
                for call_chunk in [
                    calls_of_level[i : i + batch_size]
                    for i in range(
                        0,
                        batch_size * max_workers
                        if batch_size * max_workers < len(calls_of_level)
                        else len(calls_of_level),
                        batch_size,
                    )
                ]:
                    calls_for_workers.append(call_chunk)

                for call_chunk in calls_for_workers:
                    logger.info(
                        f"Calling multicall2 with {len(call_chunk)} calls at block {block_number}"
                    )

                    futures.append(
                        executor.submit(
                            make_multicall,
                            web3_provider_uri,
                            blockchain_type,
                            access_id,
                            call_chunk,
                            block_timestamp,
                            block_number,
                            no_block_specifier,
                        )
                    )

                for index, future in enumerate(as_completed(futures)):
                    try:
                        result = future.result(timeout=30)
                        logger.info(
                            f"Multicall2 returned {len(result)} results at block {block_number}"
                        )
                        make_multicall_result.extend(result)

                    except Exception as e:
                        logger.error(f"Exception: {e}")
                        faild_calls.extend(calls_for_workers[index])

            retry = 0
            calls_of_level = calls_of_level[len(make_multicall_result) :]
            calls_of_level.extend(faild_calls)
            logger.info(f"lenght of task left {len(calls_of_level)}.")
            batch_size = min(batch_size * 2, max_batch_size)
        except ValueError as e:  # missing trie node
            logger.error(f"ValueError: {e}, retrying")
            retry += 1
            if "missing trie node" in str(e):
                time.sleep(4)
            if retry > 5:
                raise (e)
            batch_size = max(batch_size // 3, min_batch_size)
        except TimeoutError as e:  # timeout
            logger.error(f"TimeoutError: {e}, retrying")
            retry += 1
            if retry > 5:
                raise (e)
            batch_size = max(batch_size // 3, min_batch_size)
        except Exception as e:
            logger.error(f"Exception: {e}")
            raise (e)
        print(f"retry: {retry}")
        # results parsing and writing to database
        add_to_session_count = 0
        for result in make_multicall_result:
            db_view = view_call_to_label(blockchain_type, result)
            db_session.add(db_view)
            add_to_session_count += 1

            if result["hash"] not in responces:
                responces[result["hash"]] = []
            responces[result["hash"]].append(result["result"])
        start_time = time.time()
        commit_session(db_session)
        print(f"Commit time: {time.time() - start_time}")
        logger.info(f"{add_to_session_count} labels commit to database.")

    return batch_size


def parse_jobs(
    jobs: List[Any],
    blockchain_type: AvailableBlockchainType,
    web3_provider_uri: Optional[str],
    block_number: Optional[int],
    batch_size: int,
    access_id: UUID,
    max_batch_size: int = 5000,
    max_workers: int = 10,
    no_block_specifier: bool = False,
):
    """
    Parse jobs from list and generate web3 interfaces for each contract.
    """

    contracts_ABIs: Dict[str, Any] = {}
    contracts_methods: Dict[str, Any] = {}
    calls: Dict[int, Any] = {0: []}

    web3_client = select_provider(
        web3_provider_uri=web3_provider_uri,
        blockchain_type=blockchain_type,
        access_id=access_id,
    )

    logger.info(f"Crawler started connected to blockchain: {blockchain_type}")

    if block_number is None:
        block_number = web3_client.eth.get_block("latest").number  # type: ignore
    else:
        block_number = int(block_number)

    logger.info(f"Current block number: {block_number}")

    block_timestamp = web3_client.eth.get_block(block_number).timestamp  # type: ignore

    def recursive_unpack(method_abi: Any, level: int = 0) -> Any:
        """
        Generate tree of calls for crawling
        """
        have_subcalls = False

        abi = {
            "inputs": [],
            "outputs": method_abi["outputs"],
            "name": method_abi["name"],
            "type": "function",
            "stateMutability": "view",
        }

        for input in method_abi["inputs"]:
            if type(input["value"]) in (str, int, list):
                abi["inputs"].append(input)

            elif type(input["value"]) == dict:
                if input["value"]["type"] == "function":
                    hash_link = recursive_unpack(input["value"], level + 1)
                    # replace defenition by hash pointing to the result of the recursive_unpack
                    input["value"] = hash_link
                    have_subcalls = True
                abi["inputs"].append(input)
        abi["address"] = method_abi["address"]
        generated_hash = hashlib.md5(
            json.dumps(abi, sort_keys=True, indent=4, separators=(",", ": ")).encode(
                "utf-8"
            )
        ).hexdigest()

        abi["generated_hash"] = generated_hash
        if have_subcalls:
            level += 1
            if not calls.get(level):
                calls[level] = []
            calls[level].append(abi)
        else:
            level = 0

            if not calls.get(level):
                calls[level] = []
            calls[level].append(abi)

        if not contracts_methods.get(job["address"]):
            contracts_methods[job["address"]] = []
        if generated_hash not in contracts_methods[job["address"]]:
            contracts_methods[job["address"]].append(generated_hash)
            if not contracts_ABIs.get(job["address"]):
                contracts_ABIs[job["address"]] = {}
            contracts_ABIs[job["address"]][generated_hash] = abi

        return generated_hash

    for job in jobs:
        if job["address"] not in contracts_ABIs:
            contracts_ABIs[job["address"]] = []

        recursive_unpack(job, 0)

    # generate contracts interfaces

    interfaces = {}

    for contract_address in contracts_ABIs:
        # collect abis for each contract
        abis = []

        for method_hash in contracts_methods[contract_address]:
            abis.append(contracts_ABIs[contract_address][method_hash])

        # generate interface
        interfaces[contract_address] = web3_client.eth.contract(
            address=web3_client.toChecksumAddress(contract_address), abi=abis
        )

    responces: Dict[str, Any] = {}

    # reverse call_tree
    call_tree_levels = sorted(calls.keys(), reverse=True)[:-1]

    db_session = PrePing_SessionLocal()

    # run crawling of levels
    try:
        # initial call of level 0 all call without subcalls directly moved there
        logger.info("Crawl level: 0")
        logger.info(f"Jobs amount: {len(calls[0])}")
        logger.info(f"call_tree_levels: {call_tree_levels}")

        batch_size = crawl_calls_level(
            access_id=access_id,
            web3_provider_uri=web3_provider_uri,
            db_session=db_session,
            calls=calls[0],
            responces=responces,
            contracts_ABIs=contracts_ABIs,
            interfaces=interfaces,
            batch_size=batch_size,
            block_number=block_number,
            blockchain_type=blockchain_type,
            block_timestamp=block_timestamp,
            max_batch_size=max_batch_size,
            max_workers=max_workers,
            no_block_specifier=no_block_specifier,
        )

        for level in call_tree_levels:
            logger.info(f"Crawl level: {level}")
            logger.info(f"Jobs amount: {len(calls[level])}")

            batch_size = crawl_calls_level(
                access_id=access_id,
                web3_provider_uri=web3_provider_uri,
                db_session=db_session,
                calls=calls[level],
                responces=responces,
                contracts_ABIs=contracts_ABIs,
                interfaces=interfaces,
                batch_size=batch_size,
                block_number=block_number,
                blockchain_type=blockchain_type,
                block_timestamp=block_timestamp,
                max_batch_size=max_batch_size,
                max_workers=max_workers,
                no_block_specifier=no_block_specifier,
            )

    finally:
        db_session.close()


def handle_crawl(args: argparse.Namespace) -> None:
    """
    Ability to track states of the contracts.

    Read all view methods of the contracts and crawl
    """

    with open(args.jobs_file, "r") as f:
        jobs = json.load(f)

    blockchain_type = AvailableBlockchainType(args.blockchain)

    custom_web3_provider = args.custom_web3_provider

    if args.infura and INFURA_PROJECT_ID is not None:
        if blockchain_type not in infura_networks:
            raise ValueError(
                f"Infura is not supported for {blockchain_type} blockchain type"
            )
        logger.info(f"Using Infura!")
        custom_web3_provider = infura_networks[blockchain_type]["url"]

    parse_jobs(
        jobs,
        blockchain_type,
        custom_web3_provider,
        args.block_number,
        args.batch_size,
        args.access_id,
        args.max_batch_size,
        args.max_workers,
        args.no_block_specifier,
    )


def parse_abi(args: argparse.Namespace) -> None:
    """
    Parse the abi of the contract and save it to the database.
    """

    with open(args.abi_file, "r") as f:
        abi = json.load(f)

    output_json = []

    for method in abi:
        # read json and parse only stateMutability equal to view
        if method.get("stateMutability") and method["stateMutability"] == "view":
            output_json.append(method)

    with open(f"view+{args.abi_file}", "w") as f:
        json.dump(output_json, f)


def clean_labels_handler(args: argparse.Namespace) -> None:
    blockchain_type = AvailableBlockchainType(args.blockchain)

    web3_client = _retry_connect_web3(
        blockchain_type=blockchain_type, access_id=args.access_id
    )

    logger.info(f"Label cleaner connected to blockchain: {blockchain_type}")

    block_number = web3_client.eth.get_block("latest").number  # type: ignore

    db_session = PrePing_SessionLocal()

    try:
        clean_labels(db_session, blockchain_type, args.blocks_cutoff, block_number)
    finally:
        db_session.close()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.set_defaults(func=lambda _: parser.print_help())

    parser.add_argument(
        "--access-id",
        default=NB_CONTROLLER_ACCESS_ID,
        type=UUID,
        help="User access ID",
    )

    subparsers = parser.add_subparsers()

    view_state_crawler_parser = subparsers.add_parser(
        "crawl-jobs",
        help="continuous crawling the view methods from job structure",  # TODO(ANDREY): move tasks to journal
    )
    view_state_crawler_parser.add_argument(
        "--blockchain",
        "-b",
        type=str,
        help="Type of blovkchain wich writng in database",
        required=True,
    )
    view_state_crawler_parser.add_argument(
        "--infura",
        action="store_true",
        help="Use infura as web3 provider",
    )
    view_state_crawler_parser.add_argument(
        "--custom-web3-provider",
        "-w3",
        type=str,
        help="Type of blovkchain wich writng in database",
    )
    view_state_crawler_parser.add_argument(
        "--block-number", "-N", type=str, help="Block number."
    )
    view_state_crawler_parser.add_argument(
        "--jobs-file",
        "-j",
        type=str,
        help="Path to json file with jobs",
        required=True,
    )
    view_state_crawler_parser.add_argument(
        "--batch-size",
        "-s",
        type=int,
        default=500,
        help="Size of chunks wich send to Multicall2 contract.",
    )
    view_state_crawler_parser.add_argument(
        "--max-batch-size",
        "-m",
        type=int,
        default=4000,
        help="Max size of chunks wich send to Multicall2 contract.",
    )
    view_state_crawler_parser.add_argument(
        "--max-workers",
        "-w",
        type=int,
        default=10,
        help="Max workers for Multicall2 contract.",
    )
    view_state_crawler_parser.add_argument(
        "--no-block-specifier",
        "-nb",
        action="store_true",
        help="Disable block specifier.",
    )
    view_state_crawler_parser.set_defaults(func=handle_crawl)

    view_state_cleaner = subparsers.add_parser(
        "clean-state-labels",
        help="Clean labels from database",
    )
    view_state_cleaner.add_argument(
        "--blockchain",
        "-b",
        type=str,
        help="Type of blovkchain wich writng in database",
        required=True,
    )
    view_state_cleaner.add_argument(
        "--blocks-cutoff",
        "-N",
        required=True,
        type=int,
        help="Amount blocks back, after wich data will be remove.",
    )
    view_state_cleaner.set_defaults(func=clean_labels_handler)

    generate_view_parser = subparsers.add_parser(
        "parse-abi",
        help="Parse view methods from the abi file.",
    )

    generate_view_parser.add_argument(
        "--abi-file",
        "-a",
        type=str,
        help="Path to abi file.",
    )
    generate_view_parser.set_defaults(func=parse_abi)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
