import ast
import random
import string

from agentdojo.functions_runtime import FunctionCall, FunctionCallArgTypes


class ASTParsingError(Exception): ...


def create_python_function_from_tool_call(func_call_dict: FunctionCall) -> str:
    try:
        # Extract function name and arguments
        func_name = func_call_dict["function"]
        args_dict = func_call_dict["args"]

        # Create the function name node
        func_name_node = ast.Name(id=func_name, ctx=ast.Load())

        # Create the argument nodes
        keyword_args = [ast.keyword(arg=key, value=ast.Constant(value=value)) for key, value in args_dict.items()]

        # Create the call node
        call_node = ast.Call(func=func_name_node, args=[], keywords=keyword_args)

        # Wrap the call node in an Expression node
        expr_node = ast.Expression(body=call_node)

        return ast.unparse(expr_node)
    except Exception as e:
        raise ValueError(f"Error creating AST from dictionary: {e}")


def parse_arg_value(arg_value: ast.expr | None) -> FunctionCallArgTypes:
    match arg_value:
        case ast.Constant(value=value):
            return value
        case ast.List(elts=elements):
            return [parse_arg_value(element) for element in elements]
        case ast.Dict(keys=keys, values=values):
            return {parse_arg_value(key): parse_arg_value(value) for key, value in zip(keys, values)}
        case None:
            return None
        case _:
            raise ASTParsingError(f"The keyword values should be literals. Got: {arg_value}")


def parse_tool_call_from_ast(call: ast.expr) -> FunctionCall:
    if not isinstance(call, ast.Call):
        raise ASTParsingError("The function calls should be in the form of a call.")
    if not isinstance(call.func, ast.Name):
        raise ASTParsingError("The function calls should be in the form of a function name.")
    function_name = call.func.id
    args = {}
    for keyword in call.keywords:
        if not isinstance(keyword.arg, str):
            raise ASTParsingError(f"The keyword arguments should be strings. Got {keyword.arg}")
        args[keyword.arg] = parse_arg_value(keyword.value)

    random_string = "".join(random.choices(string.ascii_lowercase + string.digits, k=10))

    return FunctionCall(function=function_name, args=args, id=f"tool_call_{random_string}")


def parse_tool_calls_from_python_function(model_output: str) -> list[FunctionCall]:
    try:
        parsed_model_output = ast.parse(model_output)
    except SyntaxError as e:
        raise ASTParsingError(f"Error parsing the model output: {e!s}") from e
    if not isinstance(parsed_model_output.body[0], ast.Expr):
        raise ASTParsingError("Error parsing the model output.")
    if not isinstance(parsed_model_output.body[0].value, ast.List):
        if not isinstance(parsed_model_output.body[0].value, ast.Call):
            raise ASTParsingError(f"The function calls should be given a list. Got {parsed_model_output.body[0].value}")
        calls_list = [parsed_model_output.body[0].value]
    else:
        calls_list = parsed_model_output.body[0].value.elts

    return [parse_tool_call_from_ast(call) for call in calls_list]
