import esprima
from typing import Any, Callable, Dict, Optional
import types
import ast
import inspect
from smolagents.tools import Tool

class JavaScriptMethodDispatcher(Tool):
    name = "javascript_as_python_caller"
    description = "turns java script code blocks into callable python functions"
    inputs = {'jsContent':{'type':'string', 'description':'content of a regular java script code block of a function definition'}}
    output_type = "callable"
    
    def __init__(self):
        self._method_cache: Dict[str, Callable] = {}
        
    def register_js_method(self, js_code: str, method_name: Optional[str] = None) -> Callable:
        """
        Convert JavaScript code to a Python callable and register it in the dispatcher.
        
        Args:
            js_code (str): JavaScript code containing the function
            method_name (str, optional): Name to register the method under. 
                                       If None, tries to extract from the JS code.
        
        Returns:
            Callable: Python function that can be called directly
        """
        try:
            # Parse JavaScript code
            parsed = esprima.parseScript(js_code)
            
            # Find the function declaration
            func_decl = None
            for node in parsed.body:
                if node.type == 'FunctionDeclaration':
                    func_decl = node
                    break
            
            if not func_decl:
                raise ValueError("No function declaration found in JavaScript code")
            
            # Get the method name
            if method_name is None:
                method_name = func_decl.id.name
            
            # Convert JavaScript parameters to Python
            params = [param.name for param in func_decl.params]
            
            # Convert the function body
            body_lines = self._convert_js_body(func_decl.body)
            
            # Create the Python function source
            py_source = f"def {method_name}({', '.join(params)}):\n"
            py_source += "\n".join(f"    {line}" for line in body_lines)
            
            # Create function namespace
            namespace = {
                'print': print,  # Common built-ins
                'str': str,
                'int': int,
                'float': float,
                'bool': bool,
                'list': list,
                'dict': dict,
                'None': None,
                'True': True,
                'False': False
            }
            
            # Compile and execute the Python code
            compiled_code = compile(py_source, '<string>', 'exec')
            exec(compiled_code, namespace)
            
            # Get the compiled function
            py_func = namespace[method_name]
            
            # Store in cache
            self._method_cache[method_name] = py_func
            
            return py_func
            
        except Exception as e:
            raise ValueError(f"Failed to convert JavaScript method: {str(e)}")
    
    def get_method(self, method_name: str) -> Optional[Callable]:
        """
        Get a registered method by name.
        
        Args:
            method_name (str): Name of the registered method
            
        Returns:
            Optional[Callable]: The registered Python function, or None if not found
        """
        return self._method_cache.get(method_name)
    
    def call_method(self, method_name: str, *args, **kwargs) -> Any:
        """
        Call a registered method by name with given arguments.
        
        Args:
            method_name (str): Name of the registered method
            *args: Positional arguments to pass to the method
            **kwargs: Keyword arguments to pass to the method
            
        Returns:
            Any: Result of the method call
            
        Raises:
            ValueError: If method is not found or call fails
        """
        method = self.get_method(method_name)
        if method is None:
            raise ValueError(f"Method '{method_name}' not found")
            
        try:
            return method(*args, **kwargs)
        except Exception as e:
            raise ValueError(f"Failed to call method '{method_name}': {str(e)}")
    
    def _convert_js_body(self, body_node: Any) -> list[str]:
        """Convert JavaScript function body to Python code lines."""
        # This is a simplified conversion - you'd want to expand this
        # based on your specific needs
        lines = []
        
        # Handle different types of statements
        if body_node.type == 'BlockStatement':
            for statement in body_node.body:
                if statement.type == 'ReturnStatement':
                    return_value = self._convert_js_expression(statement.argument)
                    lines.append(f"return {return_value}")
                elif statement.type == 'ExpressionStatement':
                    expr = self._convert_js_expression(statement.expression)
                    lines.append(expr)
                elif statement.type == 'IfStatement':
                    condition = self._convert_js_expression(statement.test)
                    lines.append(f"if {condition}:")
                    then_lines = self._convert_js_body(statement.consequent)
                    lines.extend(f"    {line}" for line in then_lines)
                    if statement.alternate:
                        lines.append("else:")
                        else_lines = self._convert_js_body(statement.alternate)
                        lines.extend(f"    {line}" for line in else_lines)
        
        return lines
    
    def _convert_js_expression(self, node: Any) -> str:
        """Convert JavaScript expression to Python code string."""
        if node.type == 'BinaryExpression':
            left = self._convert_js_expression(node.left)
            right = self._convert_js_expression(node.right)
            op = node.operator
            # Convert JavaScript operators to Python
            op_map = {
                '===': '==',
                '!==': '!=',
                '&&': 'and',
                '||': 'or'
            }
            op = op_map.get(op, op)
            return f"({left} {op} {right})"
        elif node.type == 'Literal':
            if isinstance(node.value, str):
                return f"'{node.value}'"
            return str(node.value)
        elif node.type == 'Identifier':
            # Handle JavaScript built-ins
            js_to_py = {
                'undefined': 'None',
                'null': 'None',
                'true': 'True',
                'false': 'False'
            }
            return js_to_py.get(node.name, node.name)
        elif node.type == 'CallExpression':
            func = self._convert_js_expression(node.callee)
            args = [self._convert_js_expression(arg) for arg in node.arguments]
            # Handle special cases
            if func == 'console.log':
                func = 'print'
            return f"{func}({', '.join(args)})"
        
        raise ValueError(f"Unsupported expression type: {node.type}")

    def forward(self, jsContent: str) -> callable:
        return self.register_js_method(jsContent)