import uuid
import ast
from typing import Dict, Optional
from shimoku.utils import IN_BROWSER_PYODIDE
from shimoku.exceptions import ActionError
from shimoku.execution_logger import log_error

import logging

logger = logging.getLogger(__name__)


class Scope:
    def __init__(self, parent_scope: Optional['Scope'] = None):
        self.lines_need_async = set()
        self.inner_scopes: Dict[str, Scope] = {}
        self.needs_async = False
        self.parent_scope: Optional[Scope] = parent_scope
        self.parameters: dict = {}

    def get_function(self, name: str):
        if name in self.inner_scopes:
            return self.inner_scopes[name]
        if self.parent_scope is not None:
            return self.parent_scope.get_function(name)
        return None

    def user_defined_scope(self):
        return self.parent_scope is not None and self.parent_scope.parent_scope is not None

    def enter_scope(self, name: str):
        if name not in self.inner_scopes:
            self.inner_scopes[name] = Scope(self)
        return self.inner_scopes[name]

    def exit_scope(self):
        return self.parent_scope


class ActionsChecker(ast.NodeVisitor):

    shimoku_client_attrs = [
        'universes', 'workspaces', 'boards', 'menu_paths',
        'components', 'data', 'io', 'activities', 'plt',
        'set_workspace', 'set_board', 'set_menu_path', 'run'
    ]

    SHIMOKU_CLIENT_NAME = 'shimoku_client'

    def __init__(self):
        self.current_scope = Scope()
        self.in_class = False
        self.exists_action_function = False
        self.parameters = {}
        self.client_assigned = False
        self.shimoku_imported = False

    def is_shimoku_client(self, node):
        if isinstance(node, ast.Name) and node.id == self.SHIMOKU_CLIENT_NAME:
            return True
        if (isinstance(node, ast.Call) and
                isinstance(node.func, ast.Name) and
                node.func.id == 'Client'):
            return True
        return False

    def visit_Import(self, node):
        if not self.current_scope.parent_scope:
            return
        for alias in node.names:
            if 'asyncio' in alias.name:
                log_error(logger, f"Cannot import asyncio module.", ActionError)

    def visit_ImportFrom(self, node):
        if node.module == 'shimoku':
            if any(asname is not None for asname in [alias.asname for alias in node.names]):
                log_error(logger, f"Cannot import shimoku modules with renaming.", ActionError)
            for alias in node.names:
                if alias.name == 'Client':
                    self.shimoku_imported = True
        if not self.current_scope.parent_scope:
            return
        if node.module == 'asyncio':
            log_error(logger, f"Cannot import from asyncio module.", ActionError)

    def visit_Assign(self, node):
        for target in node.targets:
            if self.is_shimoku_client(node.value):
                if self.client_assigned:
                    log_error(logger, f"Cannot assign Client instance", ActionError)
                self.client_assigned = True
            elif isinstance(target, ast.Name):
                if target.id == self.SHIMOKU_CLIENT_NAME:
                    log_error(logger, f"Cannot assign to {self.SHIMOKU_CLIENT_NAME} variable.", ActionError)

        self.generic_visit(node)

    def visit_FunctionDef(self, node):
        if node.name == 'action':
            if self.in_class:
                log_error(logger, f"Cannot define action function inside a class.", ActionError)
            if self.current_scope.user_defined_scope():
                log_error(logger, f"Cannot define action function inside another function.", ActionError)
            if self.exists_action_function:
                log_error(logger, f"Cannot define more than one action function.", ActionError)
            if len(node.args.args) != 1:
                log_error(logger, f"Action function must have exactly one parameter.", ActionError)
            if node.args.args[0].arg != self.SHIMOKU_CLIENT_NAME:
                log_error(logger, f"Action function parameter must be named 'shimoku_client'.", ActionError)
            if node.args.args[0].annotation and node.args.args[0].annotation.id != 'Client':
                log_error(logger, f"Action function parameter must be annotated with 'Client'.", ActionError)
            self.exists_action_function = True

        self.current_scope = self.current_scope.enter_scope(node.name)
        self.generic_visit(node)

        for arg in node.args.args:
            if self.in_class and arg.arg == 'self':
                continue

            if not arg.annotation:
                log_error(logger, f"Annotations are mandatory for function parameters.", ActionError)
            if isinstance(arg.annotation, ast.Subscript):
                if 'Client' in ast.dump(arg.annotation):
                    log_error(
                        logger,
                        f"Cannot use Client instance as a "
                        f"part of a parameter type, use 'Client' instead.",
                        ActionError
                    )
            elif arg.annotation.id == 'Client':
                if arg.arg != self.SHIMOKU_CLIENT_NAME:
                    log_error(
                        logger,
                        f"Client parameter must be named '{self.SHIMOKU_CLIENT_NAME}' not '{arg.arg}'.",
                        ActionError
                    )
            elif arg.arg == self.SHIMOKU_CLIENT_NAME:
                log_error(
                    logger,
                    f"{self.SHIMOKU_CLIENT_NAME} parameter must be annotated with "
                    f"'Client' not '{arg.annotation.id}'.",
                    ActionError
                )
            self.current_scope.parameters[arg.arg] = arg.annotation.id

        self.current_scope = self.current_scope.exit_scope()

    def visit_AsyncFunctionDef(self, node):
        if not self.current_scope.parent_scope:
            self.current_scope = self.current_scope.enter_scope(node.name)
            self.generic_visit(node)
            self.current_scope = self.current_scope.exit_scope()
        else:
            log_error(logger, f"Cannot use async functions.", ActionError)

    def visit_ClassDef(self, node):
        self.in_class = True
        self.generic_visit(node)
        self.in_class = False

    def visit_Return(self, node):
        if isinstance(node.value, ast.Name):
            if node.value.id == self.SHIMOKU_CLIENT_NAME:
                log_error(logger, f"Cannot return Client instance '{node.value.id}'.", ActionError)
        self.generic_visit(node)

    def visit_Call(self, node):
        self.generic_visit(node)

        # Check if the call is a method call that requires await
        if isinstance(node.func, ast.Attribute):
            aux_node = node.func
            while not isinstance(aux_node.value, ast.Name):
                if not isinstance(aux_node.value, ast.Attribute):
                    return
                aux_node = aux_node.value
            if aux_node.value.id == self.SHIMOKU_CLIENT_NAME and aux_node.attr in self.shimoku_client_attrs:
                print(f"Client method call '{node.func.attr}' used at line {node.lineno} without await.")
                self.current_scope.lines_need_async.add((node.lineno, node.col_offset))
                self.current_scope.needs_async = True

        if not isinstance(node.func, ast.Name) or self.current_scope.get_function(node.func.id) is None:
            return

        function_scope = self.current_scope.get_function(node.func.id)
        if function_scope.needs_async:
            self.current_scope.needs_async = True

        parameters = function_scope.parameters
        parameter_types = list(parameters.values())
        for i, arg in enumerate(node.args):
            if self.is_shimoku_client(arg) and not parameter_types[i] == 'Client':
                log_error(
                    logger,
                    f"The value asigned to the parameter '{arg.id}' is the Client, "
                    f"but the function '{node.func.id}' expects a '{parameter_types[i]}' instance.",
                    ActionError
                )
        for keyword in node.keywords:
            if self.is_shimoku_client(keyword.value) and not parameters[keyword.arg] == 'Client':
                log_error(
                    logger,
                    f"The value asigned to the parameter '{keyword.arg}' is the Client, "
                    f"but the function '{node.func.id}' expects a '{parameters[keyword.arg]}' instance.",
                    ActionError
                )


class AwaitInserter(ast.NodeTransformer):

    def __init__(self, scope: Scope):
        self.current_scope = scope

    def visit_Call(self, node):
        self.generic_visit(node)
        if not self.current_scope.parent_scope:
            return node

        if (node.lineno, node.col_offset) in self.current_scope.lines_need_async:
            self.current_scope.needs_async = True
            return ast.copy_location(ast.Await(value=node), node)

        # Check if the call is a function call that requires await
        func_name = node.func.id if isinstance(node.func, ast.Name) else node.func.attr
        function_scope = self.current_scope.get_function(func_name)
        if function_scope is not None and function_scope.needs_async:
            self.current_scope.needs_async = True
            print(f"Function call '{func_name}' used at line {node.lineno} without await.")
            return ast.copy_location(ast.Await(value=node), node)

        return node

    def visit_FunctionDef(self, node):
        self.current_scope = self.current_scope.enter_scope(node.name)
        node = self.generic_visit(node)
        async_node = None
        if self.current_scope.needs_async:
            # Convert to an async function
            async_node = ast.AsyncFunctionDef(
                name=node.name,
                args=node.args,
                body=node.body,
                decorator_list=node.decorator_list,
                returns=node.returns,
                type_comment=node.type_comment
            )
        self.current_scope = self.current_scope.exit_scope()
        if async_node is not None:
            return ast.copy_location(async_node, node)
        return node

    def visit_AsyncFunctionDef(self, node):
        self.current_scope = self.current_scope.enter_scope(node.name)
        node = self.generic_visit(node)
        self.current_scope = self.current_scope.exit_scope()
        return node


def execute_action(code: str, print_code: bool = False):
    code = code.replace('\r\n', '\n').replace('\\\n', '').replace('\\\r', '')
    code_lines = code.split('\n')
    main_uuid = str(uuid.uuid4()).replace('-', '_')

    script_code_lines = [
        'import asyncio',
        '',
        f'async def main_{main_uuid}():',
        *[f'    {line}' for line in code_lines],
        '    shimoku_client = Client(',
        '        async_execution=True,',
        '        verbosity="INFO",',
        '    )',
        '    action(shimoku_client)',
        '',
    ]
    if IN_BROWSER_PYODIDE:
        script_code_lines.append(f'asyncio.get_event_loop().create_task(main_{main_uuid}())', )
    else:
        script_code_lines.append(f'asyncio.run(main_{main_uuid}())', )

    if print_code:
        for i, line in enumerate(script_code_lines):
            print(f'{i:3} | {line}')

    script_ast = ast.parse('\n'.join(script_code_lines))

    analyzer = ActionsChecker()
    analyzer.visit(script_ast)
    if not analyzer.exists_action_function:
        log_error(logger, f"Action function is not defined.", ActionError)
    if not analyzer.shimoku_imported:
        log_error(logger, f"Shimoku's Client is not imported.", ActionError)
    modifier = AwaitInserter(analyzer.current_scope)
    modifier.visit(script_ast)
    dump = ast.dump(script_ast, indent=4).split('\n')
    # for i, line in enumerate(dump):
    #     print(f'{i:3} | {line}')
    exec(compile(script_ast, filename='shimoku_action', mode='exec'))


def main():
    # with open('../../../tests/unit_tests/test_plot_api.py', 'r') as f:
    with open('../../../tests/personal_test.py', 'r') as f:
        code_lines = f.readlines()
    execute_action('\n'.join(code_lines))


if __name__ == '__main__':
    main()
