"""
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
from functools import lru_cache
from lionagi.libs import AsyncUtil
from ..evaluator.base_evaluator import BaseEvaluator
from .sandbox_ import SandboxTransformer


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

    def process_data(self, data):
        return data * 2

    def _evaluate_expression(self, expression):
        return self.safe_evaluator.evaluate(expression, self.variables)

    def _assign_variable(self, var_name, value):
        self.variables[var_name] = value

    def _execute_function(self, func_name, arg):
        if func_name in self.functions:
            return self.functions[func_name](arg)
        else:
            raise ValueError(f"Function {func_name} not defined.")

    def execute(self, script):
        tree = ast.parse(script)
        for stmt in tree.body:
            if isinstance(stmt, ast.Assign):
                var_name = stmt.targets[0].id
                value = self._evaluate_expression(ast.unparse(stmt.value))
                self._assign_variable(var_name, value)
            elif isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Call):
                func_name = stmt.value.func.id
                arg = self._evaluate_expression(ast.unparse(stmt.value.args[0]))
                result = self._execute_function(func_name, arg)
                # For demonstration, manually update 'x' to simulate expected behavior
                if func_name == "processData":
                    self._assign_variable("x", result)

    def add_hook(self, event_name, callback):
        """Subscribe a callback function to a specific event."""
        self.hooks[event_name].append(callback)

    def trigger_hooks(self, event_name, *args, **kwargs):
        """Trigger all callbacks attached to a specific event."""
        for callback in self.hooks[event_name]:
            callback(*args, **kwargs)

    async def process_data(self, data):
        # Example asynchronous function
        return data * 2

    @lru_cache(maxsize=128)
    def evaluate_expression(self, expression):
        return self.safe_evaluator.evaluate(expression, self.variables)

    async def _execute_function_async(self, func_name, arg):
        if func_name in self.functions:
            func = self.functions[func_name]
            if AsyncUtil.is_coroutine_func(func):
                return await func(arg)
            else:
                return func(arg)
        else:
            raise ValueError(f"Function {func_name} not defined.")

    def execute_sandboxed(self, script):
        # Parse and sanitize the script
        tree = ast.parse(script, mode="exec")
        sanitized_tree = SandboxTransformer().visit(tree)
        ast.fix_missing_locations(sanitized_tree)

        # Compile the sanitized AST
        code = compile(sanitized_tree, "<sandbox>", "exec")

        # Execute the code in a restricted namespace
        exec(code, {"__builtins__": None}, self.variables)
