from typing import Any, Dict, List, Tuple, Optional, Union

from dataclasses import dataclass

from copy import deepcopy

from keyword import iskeyword

from ast import AST, parse as ast_parse
from ast import Call, Import, ImportFrom, NodeVisitor, Name

import json
import json_repair

from openai.types.chat import (
    ChatCompletionToolParam,
    ChatCompletionMessageToolCallParam,
)


ALLOWED_ROLES = [
    "system",
    "tool",
    "user",
    "assistant",
]

ALLOWED_IMPORTS = [
    "typing",
    "itertools",
    "collections",
    "tabulate",
    "dataclasses",
    "requests",
]
DISALLOWED_FUNCTIONS = ["eval", "exec", "setattr", "locals", "globals"]


class SecurityVisitor(NodeVisitor):
    def visit_Call(self, node: Call):
        if isinstance(node.func, Name):
            func_name = node.func.id
            if func_name in DISALLOWED_FUNCTIONS:
                raise ValueError(
                    f"Found dangerous call to {func_name} at line {node.lineno}. Disallowed functions: {DISALLOWED_FUNCTIONS}"
                )

        self.generic_visit(node)

    def visit_Import(self, node: Import) -> Any:
        import_names = [a.name for a in node.names]
        self.check_imports(node, import_names)

        self.generic_visit(node)

    def visit_ImportFrom(self, node: ImportFrom) -> Any:
        import_names = [node.module]
        self.check_imports(node, import_names)

        self.generic_visit(node)

    def check_imports(
        self, node: Union[Import, ImportFrom], import_names: List[str]
    ) -> None:
        disallowed_import_names = [n for n in import_names if n not in ALLOWED_IMPORTS]
        if disallowed_import_names:
            raise ImportError(
                f"Not allowed to import {disallowed_import_names} at line {node.lineno}. Allowed imports include {ALLOWED_IMPORTS}"
            )


class CodeExecutionHelper:
    def clean_input(self, source: str) -> Tuple[str, Optional[AST]]:
        source = source.strip()
        if len(source) == 0:
            return source, None

        tree = ast_parse(source)
        return source, tree

    def check_security(self, tree: AST) -> None:
        security_visitor = SecurityVisitor()
        try:
            security_visitor.visit(tree)
        except ValueError as e:
            raise TypeError from e


"""
A smattering of symbols that are invalid for function and argument names
These are all replaced with "_" downstream.
This is a translation table used in `str.translate`.
"""
INVALID_NAME_CHARS = str.maketrans(dict.fromkeys(" !%&(*+,-/<=>@[^|~", "_"))


def clean_str(s: str) -> str:
    s = s.translate(INVALID_NAME_CHARS)
    if iskeyword(s):
        s = f"_{s}"
    return s


@dataclass
class CleanedFunction:
    dirty_name: str
    dirty_argument_names: List[str]

    def __post_init__(self) -> None:
        self.clean_name = clean_str(self.dirty_name)
        self.clean_argument_names = list(map(clean_str, self.dirty_argument_names))

        self.argument_names_d2c = dict(
            zip(self.dirty_argument_names, self.clean_argument_names)
        )
        self.argument_names_c2d = dict(
            zip(self.clean_argument_names, self.dirty_argument_names)
        )

    def clean_tool_param(
        self, tool: ChatCompletionToolParam
    ) -> ChatCompletionToolParam:
        tool = deepcopy(tool)

        function = tool["function"]
        function["name"] = self.clean_name

        properties = function["parameters"].get("properties")
        if properties:
            properties = {self.argument_names_d2c[k]: v for k, v in properties.items()}
            function["parameters"]["properties"] = properties

        required = function["parameters"].get("required")
        if required:
            function["parameters"]["required"] = [
                self.argument_names_d2c[k] for k in required
            ]

        return tool

    def clean_tool_call_param(
        self, tool_call: ChatCompletionMessageToolCallParam
    ) -> ChatCompletionMessageToolCallParam:
        tool_call = deepcopy(tool_call)

        function = tool_call["function"]
        function["name"] = self.clean_name

        arguments = json_repair.loads(function["arguments"])
        arguments = {self.argument_names_d2c[k]: v for k, v in arguments.items()}
        function["arguments"] = json.dumps(arguments)

        return tool_call

    def dirty_function_call_dict(
        self, clean_name: str, clean_kwargs: Dict[str, Any]
    ) -> Tuple[str, Dict[str, Any]]:
        assert self.clean_name == clean_name
        return self.dirty_name, {
            self.argument_names_c2d[k]: v for k, v in clean_kwargs.items()
        }

    @classmethod
    def from_tool_param(cls, tool: ChatCompletionToolParam) -> "CleanedFunction":
        function = tool["function"]
        function_properties = function["parameters"].get("properties", dict())
        cf = cls(
            dirty_name=function["name"],
            dirty_argument_names=list(function_properties),
        )

        return cf
