"""
Custom Python Code Executor

Allows users to run custom Python code snippets within the workflow.
"""

from __future__ import annotations

import sys
import io
import traceback
from typing import Any, TYPE_CHECKING

if TYPE_CHECKING:
    from ..schemas import Node
    from ..storage import StorageManager


def custom_python_executor(
    node: "Node",
    input_data: dict[str, Any],
    storage: "StorageManager",
    run_id: str,
) -> Any:
    """
    Execute custom Python code from the node.
    
    The custom code has access to:
    - input_data: Dictionary of inputs from upstream nodes
    - pd: pandas library (if installed)
    - np: numpy library (if installed)
    - storage: StorageManager for saving artifacts
    
    The code should set `output` variable to return data to downstream nodes.
    """
    params = node.data.params if hasattr(node.data, 'params') else {}
    custom_code = params.get("customCode", "")
    
    if not custom_code:
        # If no custom code, pass through the first input (or all inputs)
        if input_data:
            values = list(input_data.values())
            return values[0] if len(values) == 1 else input_data
        return input_data
    
    # Import common libraries for the execution context
    try:
        import pandas as pd
    except ImportError:
        pd = None
    
    try:
        import numpy as np
    except ImportError:
        np = None
    
    try:
        import sklearn
    except ImportError:
        sklearn = None
    
    # Prepare execution context
    local_context = {
        "input_data": input_data,
        "inputs": input_data,  # Alias
        "pd": pd,
        "np": np,
        "sklearn": sklearn,
        "storage": storage,
        "run_id": run_id,
        "node_id": node.id,
        "output": None,
        "print_output": [],
    }
    
    # Capture print statements
    old_stdout = sys.stdout
    sys.stdout = io.StringIO()
    
    try:
        # Execute the custom code
        exec(custom_code, local_context)
        
        # Capture any print output
        print_output = sys.stdout.getvalue()
        if print_output:
            local_context["print_output"] = print_output.strip().split("\n")
        
        # Get the output
        output = local_context.get("output")
        
        # If no explicit output, try to get the last result
        if output is None:
            # Check for common output variable names
            for var_name in ["result", "df", "data", "X", "y", "predictions"]:
                if var_name in local_context and local_context[var_name] is not None:
                    output = local_context[var_name]
                    break
        
        # If still no output, pass through input_data
        if output is None:
            output = input_data
        
        return output
        
    except Exception as e:
        # Restore stdout
        sys.stdout = old_stdout
        
        # Get full traceback
        tb = traceback.format_exc()
        
        raise RuntimeError(
            f"Custom code execution failed in node '{node.id}':\n"
            f"{str(e)}\n\n"
            f"Traceback:\n{tb}"
        )
    finally:
        sys.stdout = old_stdout
