"""Create functions for a Teal contract.

Teal contract with N subroutines is represented using `N + 1` disconnected CFGs. Each subroutine
has its own CFG and the code which is not part of any subroutine is represented using a separate CFG, referred to
as __main__ CFG. The entry block of the __main__ CFG is the entry block of the contract. Execution of the contract
always starts here. Each of the CFG, including the __main__ CFG, is represented using a Subroutine class.

Every contract has abstract divison of function. A function represents an user operation: whenever a contract is executed,
the execution always results in execution of one of the function. The main benefit of functions is the ability to perform
analysis of each function independently which gives more accurate results.

In general, the first thing executed in a contract is the method dispatcher, it passes the execution to the method/function's
entry block based on the method selector. The method dispatcher is often generated by the compiler(PyTeal) or the framework used(Beaker).
Some contracts also contain custom method dispatchers implemented by the developers. And some contracts would not contain any functions,
the entire contract is a single function.

The method dispatcher can be:
    - Compiler or the compilation framework generated
    - Custom method dispatcher.
    - No method dispatcher. Entire contract is a single function.

A Function can be represented using the function entry block. The body of the function is the code reachable from that entry
block. The Function class could just store the entry block and analyze the function using that. However, the contracts with
custom dispatchers might contain the function body as part of its method dispatcher. For example, the dispatcher
could contain validations, variable assignments that are common for multiple functions in it. In order to represent these functions,
Tealer replaces basic blocks that are part of the dispatcher but not part of the function with basic blocks that contain the
error instruction. As a result, the function code that is part of the dispatcher can also be analyzed.

The entry block of a function for contracts with compiler generated method dispatcher is not in the middle of a subroutine.
In practice, the custom method dispatcher also does not have entry block in the middle of the subroutine. It
is possible to implement to the method dispatcher such that some blocks of a subroutine are part of one function and
other blocks are part of another function. However, it is not common in practice.

Tealer assumes the custom method dispatcher does not have a function with entry block in the middle of the subroutine. Main reason
for this assumption is to simplify the design. __main__ CFG is disconnected from other CFGs and if the method dispatcher is only part
of this CFG, it is easy to divide and create copies of the __main__ CFG for each function and keep a single copy of
subroutines CFG.

With M functions and N subroutines, the Teal contract will be represented with `N + M + 1` CFGs: M functions + N subroutines and __main__
CFG.

For functions whose dispatcher contains function body, all blocks from contract entry block to the function entry block will be
preserved and blocks that are part of the dispatcher and not part of the function will be replaced by custom Err instruction.


NOTE: The following documentation is not accurate with the implementation. The implementation should be updated to achieve the following
functionality/design.

Functions whose dispatcher does not contain function body, blocks part of the dispatcher will not be preserved.

Construction of Functions:

Inputs: contract(Teal), user configuration(UserConfig).

    UserConfig.has_functions(bool): True if the contract has more than one function otherwise False
    UserConfig.identified_functions(bool): True if the functions are already identified and if Tealer does not have to use heuristics to
        identify them, otherwise False. If False, Tealer tries to identify the dispatcher and the functions.
    UserConfig.functions(List[Tuple[List[str], bool]]): List of functions. Each function is a Tuple of dispatcher_path(List[str]) and
        dispatcher_has_function_body(bool).
            - The dispatcher path is the list of basic block ids from the __main__ entry block id to the function
            entry block.
            - The dispatcher_has_function_body is True when the dispatcher blocks contain validations, variable declarations, etc..

Output: List of functions of the contract (List[Function]).

Assumptions:
    - The entry block of a function is not part of a subroutine.

* if not UserConfig.has_functions:
    return copy of the __main__ CFG as Function.
* if not UserConfig.identified_functions:
    run heuristics and identify the dispatcher. returned value has same type as UserConfig.functions
* remove duplicates from the new functions and UserConfig.functions
* loop over the new functions and UserConfig.functions
    * if dispatcher_has_function_body:
        * Construct CFG with blocks form __main__ CFG entry to function entry block.
        * Replace unrelated dispatcher blocks with custom err instruction block.
    * else:
        * Construct a copy of __main__ CFG without the dispatcher blocks.
"""

import inspect
import logging
from typing import TYPE_CHECKING, List, Dict, Optional
from collections import defaultdict

from tealer.teal.functions import Function

from tealer.teal.parse_teal import (
    first_pass,
    second_pass,
    create_bb,
    fourth_pass,
    identify_subroutine_blocks,
)
from tealer.teal.instructions.instructions import Callsub, TealerCustomErrInstruction
from tealer.analyses.dataflow.transaction_context import all_constraints
from tealer.analyses.dataflow.transaction_context.generic import DataflowTransactionContext
from tealer.analyses.utils.stack_ast_builder import construct_stack_ast, compute_equations
from tealer.exceptions import TealerException
from tealer.teal.subroutine import Subroutine
from tealer.teal.basic_blocks import BasicBlock

if TYPE_CHECKING:
    from tealer.teal.teal import Teal
    from tealer.teal.instructions.instructions import Instruction, Label


def copy_main_cfg(teal: "Teal") -> List["BasicBlock"]:  # pylint: disable=too-many-locals

    source_code = ""
    original_instructions: List["Instruction"] = []
    original_blocks: List["BasicBlock"] = []

    for bi in sorted(teal.main.blocks, key=lambda bi: bi.idx):
        # sorting groups the blocks which are next to each other in the source code.
        for ins in bi.instructions:
            source_code += "\n".join(ins.comments_before_ins) + "\n" + ins.source_code + "\n"
            original_instructions.append(ins)
        source_code += "\n"
        original_blocks.append(bi)

    # similar to parse_teal.parse_teal
    instructions: List["Instruction"] = []
    labels: Dict[str, "Label"] = {}
    subroutine_callsubs: Dict[str, List[Callsub]] = defaultdict(list)

    lines = source_code.splitlines()

    _, _ = first_pass(lines, labels, subroutine_callsubs, instructions)
    second_pass(instructions, labels)

    all_bbs: List["BasicBlock"] = []
    create_bb(instructions, all_bbs)

    fourth_pass(all_bbs)

    for ins_copy, ins_orig in zip(instructions, original_instructions):
        ins_copy.line = ins_orig.line
        if isinstance(ins_copy, Callsub):
            assert isinstance(ins_orig, Callsub)
            ins_copy.called_subroutine = ins_orig.called_subroutine

    all_bbs = sorted(all_bbs, key=lambda bi: bi.entry_instr.line)
    for bb_copy, bb_orig in zip(all_bbs, original_blocks):
        bb_copy.idx = bb_orig.idx

    for bb in all_bbs:
        bb.teal = teal
        bb.tealer_comments.insert(0, f"block_id = {bb.idx}; cost = {bb.cost}")

    return all_bbs


def _apply_transaction_context_analysis(function: "Function") -> None:
    logger = logging.getLogger("Tealer")
    logger.debug("[+] Running Transaction context analysis")
    group_indices_cls = all_constraints.GroupIndices
    analyses_classes = [getattr(all_constraints, name) for name in dir(all_constraints)]
    analyses_classes = [
        c
        for c in analyses_classes
        if inspect.isclass(c)
        and issubclass(c, DataflowTransactionContext)
        and c != group_indices_cls
    ]
    # Run group indices analysis first as other analysis use them.
    logger.debug(f'[+] Running txn field analysis "{group_indices_cls.__name__}"')
    group_indices_cls(function).run_analysis()
    for cl in analyses_classes:
        logger.debug(f'[+] Running txn field analysis: "{cl.__name__}"')
        cl(function).run_analysis()
    # clear cache
    construct_stack_ast.cache_clear()  # construct stack ast is not used after transaction_context_analysis.
    compute_equations.cache_clear()  # compute_equations is not used after transaction_context_analysis.


# pylint: disable=too-many-locals, too-many-branches
def construct_function(
    teal: "Teal", dispatch_path: List[str], function_name: Optional[str] = None
) -> "Function":
    function_blocks = copy_main_cfg(teal)
    entry = function_blocks[0]
    dispatch_path_blocks: List["BasicBlock"] = []
    current_valid_next_blocks = [entry]
    for bid in dispatch_path:
        block_index = int(bid[1:])  # B{index}
        for bi in current_valid_next_blocks:
            if bi.idx == block_index:
                if bi in dispatch_path_blocks:
                    raise TealerException(f"Dispatch path is a loop: {dispatch_path}")
                dispatch_path_blocks.append(bi)
                current_valid_next_blocks = bi.next
                break
        else:
            raise TealerException(f"Invalid dispatch path: {dispatch_path}, {bid} is not valid")

    for i, bi in enumerate(dispatch_path_blocks[:-1]):
        valid_next = dispatch_path_blocks[i + 1]
        for j, bi_next in enumerate(list(bi.next)):
            if bi_next != valid_next:
                err_block = BasicBlock()
                err_instruction = TealerCustomErrInstruction()
                # TODO: use a more meaningful idx??
                err_instruction.line = (bi_next.entry_instr.line << 16) + bi.exit_instr.line
                err_instruction.source_code = "TealerErr"
                err_block.idx = (bi_next.idx << 16) + (bi_next.idx)
                err_block.add_instruction(err_instruction)
                err_block.teal = teal
                err_block.tealer_comments.insert(0, "Tealer Custom Err Block")

                bi.next[j] = err_block
                err_block.add_prev(bi)
                bi_next.prev.remove(bi)
                # print("bi_next.idx", bi_next.idx, bi_next.prev)

                # instruction edges are not changed here.
                # instruction edges should not be used. Prefer using basic_block .next and .prev.
    entry = dispatch_path_blocks[0]
    function_main_blocks = identify_subroutine_blocks(entry)
    if function_name is None:
        function_name = "_".join(dispatch_path)
    function_main_name = f"__main__.{function_name}"
    # print(function_name)
    # print(function_main_blocks)
    # for block in function_main_blocks:
    #     print(block.idx, block.prev, block.next)
    function_main = Subroutine(function_main_name, entry, function_main_blocks)
    for bi in function_main_blocks:
        bi.subroutine = function_main
    function_main.contract = teal

    used_subroutines: List[Subroutine] = []
    worklist: List[Subroutine] = [function_main]
    while worklist:
        for sub in worklist[0].called_subroutines:
            if sub not in used_subroutines:
                used_subroutines.append(sub)
                if sub not in worklist:
                    worklist.append(sub)
        worklist = worklist[1:]

    function_subroutines: Dict[str, Subroutine] = {sub.name: sub for sub in used_subroutines}

    all_function_blocks: List[BasicBlock] = []
    for sub in [function_main] + used_subroutines:
        all_function_blocks += sub.blocks
    # print("func_blocks =", all_function_blocks)
    function_obj = Function(
        function_name,
        entry,
        all_function_blocks,
        teal,
        function_main,
        function_subroutines,
    )

    # for block in all_function_blocks:
    #     print("block:", repr(block),  block.__hash__())
    #     for prev_b in prev_blocks_global(block):
    #         print("prev_b:", repr(prev_b), prev_b.__hash__())
    #     for next_b in next_blocks_global(block):
    #         print("next_b:", repr(next_b), next_b.__hash__())

    # for block in sorted(all_function_blocks, key=lambda b: b.idx):
    #     print(repr(block))
    #     print(block)
    # exit(1)

    # TODO: Analysis fails with KeyError when one of the basic blocks is part of __main__ and a subroutine CFG. Handle this case.
    _apply_transaction_context_analysis(function_obj)
    return function_obj
