import json
from collections import defaultdict
from typing import NamedTuple

from slither.analyses.data_dependency.data_dependency import is_dependent
from slither.core.cfg.node import Node
from slither.core.declarations import Enum, Function, Contract
from slither.core.declarations.solidity_variables import (
    SolidityVariableComposed,
    SolidityFunction,
    SolidityVariable,
)
from slither.core.expressions import NewContract
from slither.core.slither_core import SlitherCore
from slither.core.solidity_types import TypeAlias
from slither.core.source_mapping.source_mapping import SourceMapping
from slither.core.variables.state_variable import StateVariable
from slither.core.variables.variable import Variable
from slither.printers.abstract_printer import AbstractPrinter
from slither.slithir.operations import (
    Member,
    Operation,
    SolidityCall,
    LowLevelCall,
    HighLevelCall,
    EventCall,
    Send,
    Transfer,
    InternalDynamicCall,
    InternalCall,
    TypeConversion,
)
from slither.slithir.operations.binary import Binary
from slither.slithir.variables import Constant, ReferenceVariable
from slither.utils.output import Output
from slither.visitors.expression.constants_folding import ConstantFolding, NotConstant


def _get_name(f: Function | Variable) -> str:
    # Return the name of the function or variable
    if isinstance(f, Function):
        if f.is_fallback or f.is_receive:
            return "()"
    return f.solidity_signature


def _extract_payable(contracts: list[Contract]) -> dict[str, list[str]]:
    ret: dict[str, list[str]] = {}
    for contract in contracts:
        payable_functions = [_get_name(f) for f in contract.functions_entry_points if f.payable]
        if payable_functions:
            ret[contract.name] = payable_functions
    return ret


def _extract_solidity_variable_usage(
    contracts: list[Contract], sol_var: SolidityVariable
) -> dict[str, list[str]]:
    ret: dict[str, list[str]] = {}
    for contract in contracts:
        functions_using_sol_var = []
        for f in contract.functions_entry_points:
            for v in f.all_solidity_variables_read():
                if v == sol_var:
                    functions_using_sol_var.append(_get_name(f))
                    break
        if functions_using_sol_var:
            ret[contract.name] = functions_using_sol_var
    return ret


def _is_constant(f: Function) -> bool:
    """
    Heuristic:
    - If view/pure with Solidity >= 0.4 -> Return true
    - If it contains assembly -> Return false (SlitherCore doesn't analyze asm)
    - Otherwise check for the rules from
    https://solidity.readthedocs.io/en/v0.5.0/contracts.html?highlight=pure#view-functions
    with an exception: internal dynamic call are not correctly handled, so we consider them as non-constant
    :param f:
    :return:
    """
    if f.view or f.pure:
        if not f.compilation_unit.solc_version.startswith("0.4"):
            return True
    if f.payable:
        return False
    if not f.is_implemented:
        return False
    if f.contains_assembly:
        return False
    if f.all_state_variables_written():
        return False
    for ir in f.all_slithir_operations():
        if isinstance(ir, InternalDynamicCall):
            return False
        if isinstance(ir, (EventCall, NewContract, LowLevelCall, Send, Transfer)):
            return False
        if isinstance(ir, SolidityCall) and ir.function in [
            SolidityFunction("selfdestruct(address)"),
            SolidityFunction("suicide(address)"),
        ]:
            return False
        if isinstance(ir, HighLevelCall):
            if isinstance(ir.function, Variable) or ir.function.view or ir.function.pure:
                # External call to constant functions are ensured to be constant only for solidity >= 0.5
                if f.compilation_unit.solc_version.startswith("0.4"):
                    return False
            else:
                return False
        if isinstance(ir, InternalCall) and ir.function:
            # Storage write are not properly handled by all_state_variables_written
            if any(parameter.is_storage for parameter in ir.function.parameters):
                return False
    return True


def _extract_constant_functions(contracts: list[Contract]) -> dict[str, list[str]]:
    ret: dict[str, list[str]] = {}
    for contract in contracts:
        cst_functions = [_get_name(f) for f in contract.functions_entry_points if _is_constant(f)]
        cst_functions += [
            v.solidity_signature for v in contract.state_variables if v.visibility in ["public"]
        ]
        if cst_functions:
            ret[contract.name] = cst_functions
    return ret


def _extract_assert(contracts: list[Contract]) -> dict[str, dict[str, list[dict]]]:
    """
    Return the list of contract -> function name -> List(source mapping of the assert))

    Args:
        contracts: list of contracts

    Returns:

    """
    ret: dict[str, dict[str, list[dict]]] = {}
    for contract in contracts:
        functions_using_assert: dict[str, list[dict]] = defaultdict(list)
        for f in contract.functions_entry_points:
            for ir in f.all_solidity_calls():
                if ir.function == SolidityFunction("assert(bool)") and ir.node.source_mapping:
                    func_name = _get_name(f)
                    functions_using_assert[func_name].append(ir.node.source_mapping.to_json())
        if functions_using_assert:
            ret[contract.name] = functions_using_assert
    return ret


# Create a named tuple that is serialization in json
def json_serializable(cls):
    my_super = super

    def as_dict(self):
        yield dict(zip(self._fields, my_super(cls, self).__iter__(), strict=False))

    cls.__iter__ = as_dict
    return cls


@json_serializable
class ConstantValue(NamedTuple):
    # Here value should be  Union[str, int, bool]
    # But the json lib in Echidna does not handle large integer in json
    # So we convert everything to string
    value: str
    type: str


def _extract_constant_from_read(
    ir: Operation,
    r: SourceMapping,
    all_cst_used: list[ConstantValue],
    all_cst_used_in_binary: dict[str, list[ConstantValue]],
    context_explored: set[Node],
) -> None:
    var_read = r.points_to_origin if isinstance(r, ReferenceVariable) else r
    # Do not report struct_name in a.struct_name
    if isinstance(ir, Member):
        return
    if isinstance(var_read, Variable) and var_read.is_constant:
        # In case of type conversion we use the destination type
        if isinstance(ir, TypeConversion):
            if isinstance(ir.type, TypeAlias):
                value_type = ir.type.type
            else:
                value_type = ir.type
        else:
            value_type = var_read.type
        try:
            value = ConstantFolding(var_read.expression, value_type).result()
            all_cst_used.append(ConstantValue(str(value), str(value_type)))
        except NotConstant:
            pass
    if isinstance(var_read, Constant):
        all_cst_used.append(ConstantValue(str(var_read.value), str(var_read.type)))
    if isinstance(var_read, StateVariable):
        if var_read.node_initialization:
            if var_read.node_initialization.irs:
                if var_read.node_initialization in context_explored:
                    return
                context_explored.add(var_read.node_initialization)
                _extract_constants_from_irs(
                    var_read.node_initialization.irs,
                    all_cst_used,
                    all_cst_used_in_binary,
                    context_explored,
                )


def _extract_constant_from_binary(
    ir: Binary,
    all_cst_used: list[ConstantValue],
    all_cst_used_in_binary: dict[str, list[ConstantValue]],
):
    for r in ir.read:
        if isinstance(r, Constant):
            all_cst_used_in_binary[str(ir.type)].append(ConstantValue(str(r.value), str(r.type)))
        if isinstance(ir.variable_left, Constant) or isinstance(ir.variable_right, Constant):
            if ir.lvalue:
                try:
                    type_ = ir.lvalue.type
                    cst = ConstantFolding(ir.expression, type_).result()
                    all_cst_used.append(ConstantValue(str(cst.value), str(type_)))
                except NotConstant:
                    pass


def _extract_constants_from_irs(
    irs: list[Operation],
    all_cst_used: list[ConstantValue],
    all_cst_used_in_binary: dict[str, list[ConstantValue]],
    context_explored: set[Node],
) -> None:
    for ir in irs:
        if isinstance(ir, Binary):
            _extract_constant_from_binary(ir, all_cst_used, all_cst_used_in_binary)
        if isinstance(ir, TypeConversion):
            if isinstance(ir.variable, Constant):
                if isinstance(ir.type, TypeAlias):
                    value_type = ir.type.type
                else:
                    value_type = ir.type
                all_cst_used.append(ConstantValue(str(ir.variable.value), str(value_type)))
                continue
        if (
            isinstance(ir, Member)
            and isinstance(ir.variable_left, Enum)
            and isinstance(ir.variable_right, Constant)
        ):
            # enums are constant values
            try:
                internal_num = ir.variable_left.values.index(ir.variable_right.value)
                all_cst_used.append(ConstantValue(str(internal_num), "uint256"))
            except ValueError:  # index could fail; should never happen in working solidity code
                pass
        for r in ir.read:
            _extract_constant_from_read(
                ir, r, all_cst_used, all_cst_used_in_binary, context_explored
            )


def _extract_constants(
    contracts: list[Contract],
) -> tuple[dict[str, dict[str, list]], dict[str, dict[str, dict]]]:
    # contract -> function -> [ {"value": value, "type": type} ]
    ret_cst_used: dict[str, dict[str, list[ConstantValue]]] = defaultdict(dict)
    # contract -> function -> binary_operand -> [ {"value": value, "type": type ]
    ret_cst_used_in_binary: dict[str, dict[str, dict[str, list[ConstantValue]]]] = defaultdict(dict)
    for contract in contracts:
        for function in contract.functions_entry_points:
            all_cst_used: list = []
            all_cst_used_in_binary: dict = defaultdict(list)

            context_explored = set()
            context_explored.add(function)
            _extract_constants_from_irs(
                function.all_slithir_operations(),
                all_cst_used,
                all_cst_used_in_binary,
                context_explored,
            )

            # Note: use list(set()) instead of set
            # As this is meant to be serialized in JSON, and JSON does not support set
            if all_cst_used:
                ret_cst_used[contract.name][_get_name(function)] = list(set(all_cst_used))
            if all_cst_used_in_binary:
                ret_cst_used_in_binary[contract.name][_get_name(function)] = {
                    k: list(set(v)) for k, v in all_cst_used_in_binary.items()
                }
    return ret_cst_used, ret_cst_used_in_binary


def _extract_function_relations(
    contracts: list[Contract],
) -> dict[str, dict[str, dict[str, list[str]]]]:
    # contract -> function -> [functions]
    ret: dict[str, dict[str, dict[str, list[str]]]] = defaultdict(dict)
    for contract in contracts:
        ret[contract.name] = defaultdict(dict)
        written = {
            _get_name(function): function.all_state_variables_written()
            for function in contract.functions_entry_points
        }
        read = {
            _get_name(function): function.all_state_variables_read()
            for function in contract.functions_entry_points
        }
        for function in contract.functions_entry_points:
            ret[contract.name][_get_name(function)] = {
                "impacts": [],
                "is_impacted_by": [],
            }
            for candidate, varsWritten in written.items():
                if any(r in varsWritten for r in function.all_state_variables_read()):
                    ret[contract.name][_get_name(function)]["is_impacted_by"].append(candidate)
            for candidate, varsRead in read.items():
                if any(r in varsRead for r in function.all_state_variables_written()):
                    ret[contract.name][_get_name(function)]["impacts"].append(candidate)
    return ret


def _have_external_calls(contracts: list[Contract]) -> dict[str, list[str]]:
    """
    Detect the functions with external calls
    :param slither:
    :return:
    """
    ret: dict[str, list[str]] = defaultdict(list)
    for contract in contracts:
        for function in contract.functions_entry_points:
            if function.all_high_level_calls() or function.all_low_level_calls():
                ret[contract.name].append(_get_name(function))
        if contract.name in ret:
            ret[contract.name] = list(set(ret[contract.name]))
    return ret


def _use_balance(contracts: list[Contract]) -> dict[str, list[str]]:
    """
    Detect the functions with external calls
    :param slither:
    :return:
    """
    ret: dict[str, list[str]] = defaultdict(list)
    for contract in contracts:
        for function in contract.functions_entry_points:
            for ir in function.all_slithir_operations():
                if isinstance(ir, SolidityCall) and ir.function == SolidityFunction(
                    "balance(address)"
                ):
                    ret[contract.name].append(_get_name(function))
        if contract.name in ret:
            ret[contract.name] = list(set(ret[contract.name]))
    return ret


def _with_fallback(contracts: list[Contract]) -> set[str]:
    ret: set[str] = set()
    for contract in contracts:
        for function in contract.functions_entry_points:
            if function.is_fallback:
                ret.add(contract.name)
    return ret


def _with_receive(contracts: list[Contract]) -> set[str]:
    ret: set[str] = set()
    for contract in contracts:
        for function in contract.functions_entry_points:
            if function.is_receive:
                ret.add(contract.name)
    return ret


def _call_a_parameter(slither: SlitherCore, contracts: list[Contract]) -> dict[str, list[dict]]:
    """
    Detect the functions with external calls
    :param slither:
    :return:
    """
    # contract -> [ (function, idx, interface_called) ]
    ret: dict[str, list[dict]] = defaultdict(list)
    for contract in contracts:
        for function in contract.functions_entry_points:
            try:
                for ir in function.all_slithir_operations():
                    if isinstance(ir, HighLevelCall):
                        for idx, parameter in enumerate(function.parameters):
                            if is_dependent(ir.destination, parameter, function):
                                ret[contract.name].append(
                                    {
                                        "function": _get_name(function),
                                        "parameter_idx": idx,
                                        "signature": _get_name(ir.function),
                                    }
                                )
                    if isinstance(ir, LowLevelCall):
                        for idx, parameter in enumerate(function.parameters):
                            if is_dependent(ir.destination, parameter, function):
                                ret[contract.name].append(
                                    {
                                        "function": _get_name(function),
                                        "parameter_idx": idx,
                                        "signature": None,
                                    }
                                )
            except Exception as e:
                if slither.no_fail:
                    continue
                raise e
    return ret


class Echidna(AbstractPrinter):
    ARGUMENT = "echidna"
    HELP = "Export Echidna guiding information"

    WIKI = "https://github.com/trailofbits/slither/wiki/Printer-documentation#echidna"

    def output(self, filename: str) -> Output:
        """
        Output the inheritance relation

        _filename is not used
        Args:
            _filename(string)
        """

        contracts = self.slither.contracts

        payable = _extract_payable(contracts)
        timestamp = _extract_solidity_variable_usage(
            contracts, SolidityVariableComposed("block.timestamp")
        )
        block_number = _extract_solidity_variable_usage(
            contracts, SolidityVariableComposed("block.number")
        )
        msg_sender = _extract_solidity_variable_usage(
            contracts, SolidityVariableComposed("msg.sender")
        )
        msg_gas = _extract_solidity_variable_usage(contracts, SolidityVariableComposed("msg.gas"))
        assert_usage = _extract_assert(contracts)
        cst_functions = _extract_constant_functions(contracts)
        (cst_used, cst_used_in_binary) = _extract_constants(contracts)

        functions_relations = _extract_function_relations(contracts)

        constructors = {
            contract.name: contract.constructor.full_name
            for contract in contracts
            if contract.constructor
        }

        external_calls = _have_external_calls(contracts)

        # call_parameters = _call_a_parameter(self.slither, contracts)

        use_balance = _use_balance(contracts)

        with_fallback = list(_with_fallback(contracts))

        with_receive = list(_with_receive(contracts))

        d = {
            "payable": payable,
            "timestamp": timestamp,
            "block_number": block_number,
            "msg_sender": msg_sender,
            "msg_gas": msg_gas,
            "assert": assert_usage,
            "constant_functions": cst_functions,
            "constants_used": cst_used,
            "constants_used_in_binary": cst_used_in_binary,
            "functions_relations": functions_relations,
            "constructors": constructors,
            "have_external_calls": external_calls,
            # "call_a_parameter": call_parameters,
            "use_balance": use_balance,
            "solc_versions": [unit.solc_version for unit in self.slither.compilation_units],
            "with_fallback": with_fallback,
            "with_receive": with_receive,
        }

        self.info(json.dumps(d, indent=4))

        res = self.generate_output(json.dumps(d, indent=4))

        return res
