"""
Copyright 2024 HaiyangLi

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import ast
import operator


class ASTEvaluator:
    """
    Safely evaluates expressions using AST parsing to prevent unsafe operations.
    """

    def __init__(self):
        self.allowed_operators = {
            ast.Eq: operator.eq,
            ast.NotEq: operator.ne,
            ast.Lt: operator.lt,
            ast.LtE: operator.le,
            ast.Gt: operator.gt,
            ast.GtE: operator.ge,
            # Additional operators can be added here as needed
        }

    def evaluate(self, expression, context):
        """
        Evaluate a condition expression within a given context using AST parsing.
        """
        try:
            tree = ast.parse(expression, mode="eval")
            return self._evaluate_node(tree.body, context)
        except Exception as e:
            raise ValueError(f"Failed to evaluate expression: {expression}. Error: {e}")

    def _evaluate_node(self, node, context):
        if isinstance(node, ast.Compare):
            left = self._evaluate_node(node.left, context)
            for operation, comparator in zip(node.ops, node.comparators):
                op_func = self.allowed_operators.get(type(operation))
                if not op_func:
                    raise ValueError(
                        f"Operation {type(operation).__name__} is not allowed."
                    )
                right = self._evaluate_node(comparator, context)
                if not op_func(left, right):
                    return False
            return True
        elif isinstance(node, ast.Name):
            return context.get(node.id)
        elif isinstance(node, ast.Constant):
            return node.n
        else:
            raise ValueError(
                "Unsupported AST node type encountered in condition evaluation."
            )


class ASTEvaluationEngine:
    """
    Executes scripts safely using the SafeEvaluator for expression evaluation.
    """

    def __init__(self):
        self.variables = {}
        self.safe_evaluator = ASTEvaluator()
        self.functions = {
            "processData": self.process_data,
        }

    def process_data(self, data):
        # Example placeholder function for data processing
        return data * 2

    def _evaluate_expression(self, expression):
        """
        Evaluates expressions within scripts using SafeEvaluator.
        """
        # Here, 'self.variables' serves as the context for the evaluation
        return self.safe_evaluator.evaluate(expression, self.variables)

    def _assign_variable(self, var_name, value):
        """
        Assigns a value to a variable within the script's context.
        """
        self.variables[var_name] = value

    def _execute_function(self, func_name, arg):
        """
        Executes a predefined function with the given argument.
        """
        if func_name in self.functions:
            function = self.functions[func_name]
            return function(arg)
        else:
            raise ValueError(f"Function '{func_name}' is not defined.")

    def execute(self, script):
        """
        Parses and executes a script, handling variable assignments and function calls.
        """
        tree = ast.parse(script, mode="exec")
        for stmt in tree.body:
            if isinstance(stmt, ast.Assign):
                var_name = stmt.targets[
                    0
                ].id  # Assumes single target assignment for simplicity
                # Convert the AST node back to a string for evaluation
                value_expr = ast.unparse(stmt.value)
                value = self._evaluate_expression(value_expr)
                self._assign_variable(var_name, value)
            elif isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
                func_name = stmt.value.func.id
                arg_expr = ast.unparse(stmt.value.args[0])
                arg = self._evaluate_expression(arg_expr)
                self._execute_function(func_name, arg)
            else:
                raise ValueError(
                    "Unsupported statement type encountered in script execution."
                )
