import ast
import functools
from typing import Callable, Dict, List, Tuple, Union

from .inference import get_inferred_v_type, V_WIDTH_RANK


class VTranspilerPlugins:
    def visit_sys_argv(self, node: ast.Attribute) -> str:
        self._usings.add("os")
        return "os.args"

    def visit_range(self, node: ast.Call, vargs: List[str]) -> str:
        if len(node.args) == 1:
            return f"0..{vargs[0]}"
        elif len(node.args) == 2:
            return f"{vargs[0]}..{vargs[1]}"

        raise Exception(
            "encountered range() call with unknown parameters: range({})".format(vargs)
        )

    def visit_print(self, node: ast.Call, vargs: List[str]) -> str:
        args: List[str] = []
        total_args: List[str] = []
        for arg in node.args:
            if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
                args.append(arg.value.replace("'", ""))
            elif get_inferred_v_type(arg) == "string":
                args.append(f"${{{self.visit(arg)}}}")
            else:
                if args:
                    total_args.append(f"'{' '.join(args)}'")
                    args = []
                total_args.append(f"({self.visit(arg)}).str()")
        if args:
            total_args.append(f"'{' '.join(args)}'")
        total_args = " + ' ' + ".join(total_args)

        return f"println({total_args})"

    def visit_bool(self, node: ast.Call, vargs: List[str]) -> str:
        if not vargs:
            return "false"
        if get_inferred_v_type(node.args[0]) in V_WIDTH_RANK:
            return f"{self.visit(node.args[0])} != 0"
        else:
            return f"({self.visit(node.args[0])}).bool()"

    @staticmethod
    def visit_cast(node, vargs, cast_to: str) -> str:
        if not vargs:
            if cast_to == "i32":
                return "0"
            elif cast_to == "f64":
                return "0.0"
        return f"{cast_to}({vargs[0]})"

    def visit_int(self, node: ast.Call, vargs: List[str]) -> str:
        if not vargs:
            return "0"
        return f"int({vargs[0]})"

    def visit_min_max(self, node: ast.Call, vargs: List[str]) -> str:
        self._usings.add("arrays")
        func = "min" if node.func.id == "min" else "max"
        # Since python allows calling min/max on either a container or
        # multiple integers while V only allows containers, the latter case
        # is turned to a container when passing it to the V side.
        if len(vargs) == 1:
            return f"arrays.{func}({vargs[0]}) or {{ panic('!') }}"
        else:
            return f"arrays.{func}([{', '.join(vargs)}]) or {{ panic('!') }}"


SMALL_DISPATCH_MAP: Dict[str, Callable] = {
    "str": lambda n, vargs: f"({vargs[0]}).str()" if vargs else '""',
    "floor": lambda n, vargs: f"int(math.floor({vargs[0]}))",
    "float": functools.partial(VTranspilerPlugins.visit_cast, cast_to="f64"),
    "len": lambda n, vargs: f"{vargs[0]}.len",
    "sys.exit": lambda n, vargs: f"exit({vargs[0] if vargs else '0'})",
    "all": lambda n, vargs: f"{vargs[0]}.all(it)",
    "any": lambda n, vargs: f"{vargs[0]}.any(it)",
}

SMALL_USINGS_MAP: Dict[str, str] = {
    "floor": "math",
}

DISPATCH_MAP: Dict[str, Callable] = {}

MODULE_DISPATCH_TABLE: Dict[str, str] = {}

DECORATOR_DISPATCH_TABLE: Dict[str, Callable] = {}

CLASS_DISPATCH_TABLE: Dict[str, Callable] = {}

AttrType = Union[type, str]

ATTR_DISPATCH_TABLE: Dict[AttrType, Callable] = {
    "sys.argv": VTranspilerPlugins.visit_sys_argv,
}

FuncType = Union[Callable, str]

FUNC_DISPATCH_TABLE: Dict[FuncType, Tuple[Callable, bool]] = {
    range: (VTranspilerPlugins.visit_range, False),
    print: (VTranspilerPlugins.visit_print, False),
    bool: (VTranspilerPlugins.visit_bool, False),
    int: (VTranspilerPlugins.visit_int, False),
    min: (VTranspilerPlugins.visit_min_max, False),
    max: (VTranspilerPlugins.visit_min_max, False),
}
