#!/usr/bin/python3
# -*- encoding: Utf-8 -*-
from typing import List, Set, Dict, Tuple, Optional, Sequence, Union, Any
from os.path import dirname, realpath
from dataclasses import dataclass
from re import sub
import sys

SCRIPT_DIR = dirname(realpath(__file__))
MODULE_DIR = dirname(realpath(SCRIPT_DIR))
SRC_DIR = dirname(realpath(MODULE_DIR))

sys.path.insert(0, SRC_DIR)

from hermes_dec.parsers.hbc_bytecode_parser import (
    ParsedInstruction,
    get_builtin_functions,
)
from hermes_dec.parsers.hbc_file_parser import HBCReader


class HermesDecompiler:
    input_file: str
    output_file: Optional[str]

    calldirect_function_ids: Set[int]

    hbc_reader: HBCReader

    function_header: (
        object  # For the function being transformed to dehydrated state
    )
    indent_level: int = 0  # Used while producing decompilation output


# This should encompass the scope for a HBC assembly-level basic
# block (retranscribed as a for(;;) switch { case X: } case in
# the decompiled code, if not matching a particular NestingFrame
# object as defined below)
# This will provide the ability to construct a graph of
# the basic assembly blocks in the disassembled/decompiled
# output
class BasicBlock:
    start_address: int
    end_address: int

    # These flags should indicate whether we have
    # encountered cycling in
    # "graph_traversers/step2_visite_code_paths.py"
    # (which indicates the presence of a loop):
    may_be_cycling_anchor: bool = False
    may_be_cycling_target: bool = False

    # These flags delimite blocks that will mark
    # the code not continue into something
    # whichsoever
    is_unconditional_throw_anchor: bool = False
    is_unconditional_return_end: bool = False

    # These flags delimite blocks that have the
    # code to continue somewhere else
    is_unconditional_jump_anchor: bool = False
    is_yield_action_anchor: bool = False

    # These flags delimite blocks that
    # should have more than one branching
    # end
    is_conditional_jump_anchor: bool = False
    if_switch_action_anchor: bool = False

    anchor_instruction: Optional[ParsedInstruction] = None
    jump_targets_for_anchor: Optional[List[int]] = None

    # These attributes will be used to build
    # the decompiler's graph structure

    child_nodes: List['BasicBlock']
    parent_nodes: List['BasicBlock']

    error_handling_child_nodes: List['BasicBlock']
    error_handling_parent_nodes: List['BasicBlock']

    # These attributes should be added at steps 3+
    # of the graph traversal process (using
    # instruction pattern maching, etc)

    do_expr_at_begin: Optional['XXX_TODO']
    if_expr_at_begin: Optional['XXX_TODO']
    else_if_expr_at_begin: Optional['XXX_TODO']
    else_expr_at_begin: Optional['XXX_TODO']
    switch_expr_at_begin: Optional['XXX_TODO']
    while_expr_at_begin: Optional['XXX_TODO']
    for_expr_at_begin: Optional['XXX_TODO']
    while_expr_at_end_matching_block: Optional['XXX_TODO']
    closing_brace_at_end_matching_block: Optional['XXX_TODO']

    # Whether this basic block should still be visible in the
    # decompiled code (switch this attribute to False when
    # a NestingFrame totally overlaps the concerned block)
    stay_visible: bool = True


# This should encompass the scope for a decompiled code-level
# nesting frame (switch, if, else if, for, while, do... block)
@dataclass
class NestedFrame:
    start_address: int
    end_address: int

    pass  # WIP


@dataclass
class Environment:
    parent_environment: Optional['Environment'] = None
    nesting_quantity: int = 0
    slot_index_to_varname: Dict[int, str] = None


class DecompiledFunctionBody:
    is_global: bool  # Is this the global function?
    function_name: str
    function_id: int
    function_object: object
    exc_handlers: 'HBCExceptionHandlerInfo'

    try_starts: Dict[int, List[str]]
    try_ends: Dict[int, List[str]]
    catch_targets: Dict[int, List[str]]

    jump_anchors: Dict[
        int, ParsedInstruction
    ]  # Non-jump address associated with a jump/switch/yield instruction
    ret_anchors: Dict[
        int, ParsedInstruction
    ]  # Unreacheable address associated with a return instruction
    throw_anchors: Dict[
        int, ParsedInstruction
    ]  # Unreacheable address associated with a jump instruction
    jump_targets: Set[int]  # Jump target address

    is_closure: bool = False
    is_async: bool = False
    is_generator: bool = False

    # Closure variable naming-related attributes:
    parent_environment: Optional[Environment] = None
    environment_id: Optional[int] = None
    local_items: Dict[int, Environment] = {}  # {env_register: Environment}

    basic_blocks: List[BasicBlock]
    nested_frames: List[NestedFrame]

    statements: List['TokenString']

    # Output decompiled code for this function to stdout
    def output_code(self, state: HermesDecompiler):
        output = ''

        if not self.is_global:  # Don't prototype the global function
            if self.is_async:
                output += 'async '
            output += 'function'
            if self.is_generator:
                output += '*'
            if not (self.is_closure or self.is_generator):
                assert self.function_name
                output += ' ' + self.function_name
            elif self.is_generator:
                output += ' '
            output += '('
            output += ', '.join(
                'a' + str(index)
                for index in range(self.function_object.paramCount - 1)
            )
            # TODO: Handle function arguments properly otherwise
            output += ') {'
            if self.is_closure or self.is_generator:
                if self.function_name:
                    output += ' // Original name: ' + self.function_name
                    if self.environment_id is not None:
                        output += ', environment: r' + str(self.environment_id)
                elif self.environment_id is not None:
                    output += ' // Environment: r' + str(self.environment_id)
            output += '\n'
            state.indent_level += 1

        nested_frame_starts = [
            nested_frame.start_address for nested_frame in self.nested_frames
        ]
        nested_frame_ends = [
            nested_frame.end_address for nested_frame in self.nested_frames
        ]

        basic_block_starts = [
            basic_block.start_address for basic_block in self.basic_blocks
        ]
        basic_block_ends = [
            basic_block.end_address for basic_block in self.basic_blocks
        ]

        if len(basic_block_starts) > 1:
            output += ' ' * (state.indent_level * 4)
            output += (
                '_fun%d: for(var _fun%d_ip = 0; ; ) switch(_fun%d_ip) {\n'
                % (
                    self.function_id,
                    self.function_id,
                    self.function_id,
                )
            )
            state.indent_level += 1

        sys.stdout.write(output)
        output = ''

        # Process each statement, including nested functions

        for statement in self.statements:
            if statement.assembly:
                pos = statement.assembly[0].original_pos

                while pos in nested_frame_ends:
                    nested_frame_ends.pop(nested_frame_ends.index(pos))
                    state.indent_level -= 1
                    output += (' ' * (state.indent_level * 4)) + '}\n'

                if len(basic_block_starts) > 1 and pos in basic_block_starts:
                    output += 'case %d:' % pos

                    if pos in self.try_starts:
                        for label in self.try_starts[pos]:
                            output += ' // %s' % label
                    if pos in self.try_ends:
                        for label in self.try_ends[pos]:
                            output += ' // %s' % label
                    if pos in self.catch_targets:
                        for label in self.catch_targets[pos]:
                            output += ' // %s' % label

                    output += '\n'

                while pos in nested_frame_starts:
                    nested_frame_starts.pop(nested_frame_starts.index(pos))
                    output += (' ' * (state.indent_level * 4)) + '{\n'
                    state.indent_level += 1

                sys.stdout.write(output)
                output = ''

            if statement.tokens:
                sys.stdout.write(' ' * (state.indent_level * 4))
                is_block = statement.tokens[:2] == [
                    RawToken('for'),
                    LeftParenthesisToken(),
                ]
                while statement.tokens:
                    op = statement.tokens.pop(0)
                    if isinstance(op, FunctionTableIndex):
                        op.closure_decompile(self)
                    elif isinstance(op, JumpNotCondition):
                        conditions = ''.join(
                            str(token) for token in statement.tokens
                        )
                        statement.tokens = []
                        if conditions == 'false':
                            sys.stdout.write(
                                '_fun%d_ip = %d; continue _fun%d'
                                % (
                                    self.function_id,
                                    op.target_address,
                                    self.function_id,
                                )
                            )
                        else:
                            is_block = True
                            sys.stdout.write('if(')
                            if '(' in conditions or ' ' in conditions:
                                sys.stdout.write('!(%s)' % conditions)
                            elif conditions[0] == '!':
                                sys.stdout.write(conditions[1:])
                            else:
                                sys.stdout.write('!' + conditions)
                            sys.stdout.write(
                                ') { _fun%d_ip = %d; continue _fun%d }'
                                % (
                                    self.function_id,
                                    op.target_address,
                                    self.function_id,
                                )
                            )
                    elif isinstance(op, JumpCondition):
                        conditions = ''.join(
                            str(token) for token in statement.tokens
                        )
                        statement.tokens = []
                        if conditions == 'true':
                            sys.stdout.write(
                                '_fun%d_ip = %d; continue _fun%d'
                                % (
                                    self.function_id,
                                    op.target_address,
                                    self.function_id,
                                )
                            )
                        else:
                            is_block = True
                            sys.stdout.write('if(')
                            sys.stdout.write(conditions)
                            sys.stdout.write(
                                ') { _fun%d_ip = %d; continue _fun%d }'
                                % (
                                    self.function_id,
                                    op.target_address,
                                    self.function_id,
                                )
                            )
                    else:
                        s = str(op)
                        sys.stdout.write(s)
                if is_block:
                    sys.stdout.write('\n')
                else:
                    sys.stdout.write(';\n')

        if len(basic_block_starts) > 1:
            state.indent_level -= 1
            sys.stdout.write(' ' * (state.indent_level * 4) + '}\n')

        if not self.is_global:
            state.indent_level -= 1
            sys.stdout.write(' ' * (state.indent_level * 4) + '}')


@dataclass
class TokenString:
    tokens: List['Token']
    assembly: ParsedInstruction


class Token:
    pass

    # def __str__(self):
    #     return repr(self).replace(self.__class__.__name__,
    #         sub(r'[^A-Z]', '', self.__class__.__name__), 1)


@dataclass
class ResumeGenerator(Token):
    result_out_reg: int
    return_bool_out_reg: int


@dataclass
class SaveGenerator(Token):
    address: int


@dataclass
class StartGenerator(Token):
    pass


@dataclass
class ReturnDirective(Token):
    def __str__(self):
        return 'return '


@dataclass
class ThrowDirective(Token):
    def __str__(self):
        return 'throw '


@dataclass
class LeftHandRegToken(Token):
    register: int

    def __str__(self):
        return 'r%d' % self.register


@dataclass
class AssignmentToken(Token):
    def __str__(self):
        return ' = '


@dataclass
class LeftParenthesisToken(Token):
    def __str__(self):
        return '('


@dataclass
class RightParenthesisToken(Token):
    def __str__(self):
        return ')'


@dataclass
class CatchBlockStart(Token):
    arg_register: int


@dataclass
class DotAccessorToken(Token):
    def __str__(self):
        return '.'


@dataclass
class IndexStringToken(Token):
    index: str

    def __str__(self):
        return '["' + self.index + '"]'


@dataclass
class BindToken(Token):
    register: int

    def __str__(self):
        return '.bind(r%d)' % self.register


@dataclass
class RightHandRegToken(Token):
    register: int

    def __str__(self):
        return 'r%d' % self.register


@dataclass
class GetEnvironmentToken(Token):
    register: int
    nesting_level: int


@dataclass
class LoadFromEnvironmentToken(Token):
    register: int
    slot_index: int


@dataclass
class NewEnvironmentToken(Token):
    register: int


@dataclass
class NewInnerEnvironmentToken(Token):
    dest_register: int
    parent_register: int
    number_of_slots: int


@dataclass
class SwitchImm(Token):
    value_reg: int
    jump_table_address: int
    default_jump_address: int
    unsigned_min_value: int
    unsigned_max_value: int


@dataclass
class StoreToEnvironment(Token):
    env_register: int
    slot_index: int
    value_register: int


@dataclass
class ForInLoopInit(Token):
    obj_props_register: int
    obj_register: int
    iter_index_register: int
    iter_size_register: int


@dataclass
class ForInLoopNextIter(Token):
    next_value_register: int
    obj_props_register: int
    obj_register: int
    iter_index_register: int
    iter_size_register: int


# Used to mark backward jumps (used in loops)
@dataclass
class JumpCondition(Token):
    target_address: int


# Used to mark forward jumps
@dataclass
class JumpNotCondition(Token):
    target_address: int


@dataclass
class FunctionTableIndex(Token):
    function_id: int
    state: HermesDecompiler
    environment_id: Optional[int] = None

    is_closure: bool = False
    is_builtin: bool = False
    is_generator: bool = False
    is_async: bool = False

    parent_environment: Optional[Environment] = None

    def closure_decompile(self, parent_closure: DecompiledFunctionBody):

        # Is this a nested closure?
        # If yes, print its decompiled code
        if (self.is_closure or self.is_generator) and not self.is_builtin:
            import hermes_dec.decompilation.hbc_decompiler

            hermes_dec.decompilation.hbc_decompiler.decompile_function(
                self.state,
                self.function_id,
                parent_environment=self.parent_environment,
                environment_id=self.environment_id,
                is_closure=self.is_closure,
                is_generator=self.is_generator,
                is_async=self.is_async,
            )
            return

        # Is this a know builtin or named function?
        name = self._name_if_any()
        if name:
            sys.stdout.write(name)
            return

        # Is this anything else?
        sys.stdout.write(
            '<Function #%d: %s>'
            % (
                self.function_id,
                ', '.join(
                    '%s: %s' % (key, value)
                    for key, value in self.__dict__.items()
                    if key != 'state' and value
                ),
            )
        )

    def _name_if_any(self) -> Optional[str]:

        if self.is_builtin:
            builtin_functions = get_builtin_functions(
                self.state.hbc_reader.parser_module
            )
            return builtin_functions[self.function_id]

        else:
            return (
                self.state.hbc_reader.strings[
                    self.state.hbc_reader.function_headers[
                        self.function_id
                    ].functionName
                ]
                or None
            )


@dataclass
class RawToken(Token):
    token: str

    def __str__(self):
        return self.token


# WIP .. .
