from __future__ import annotations
import multiprocessing as mp
import time
from typing import Any


# Helper that will run in the child process
def _child_process_worker(params, body, closure_vars, source_dir, channel, args):
    """
    Worker for multiprocessing.
    Reconstructs an interpreter and runs the given function AST.
    """
    from .interpreter import Interpreter, Environment
    
    # Create an interpreter
    interp = Interpreter(source_dir=source_dir)
    
    # We recreate a minimal environment from the closure vars
    # Since pickling full closures is dangerous, we only pass primitive values.
    # In a real system you'd want to serialize functions too.
    env = Environment(parent=interp.global_env)
    for k, v in closure_vars.items():
        env.set(k, v)
        
    # Bind the builtins channel into the environment so the function can access it
    # We'll just pass it as an argument if needed, or set it globally.
    if channel:
        env.set("__channel__", channel)
        
    # Execute the function body
    from .ast_nodes import ReturnStatement
    from .interpreter import ReturnException
    
    # Simple mockup of applying args to params
    for param_name, arg_val in zip(params, args):
        env.set(param_name, arg_val)
        
    try:
        interp._exec_block(body, env)
    except ReturnException as e:
        return e.value
    except Exception as e:
        print(f"Child process error: {e}")
        return None

class Channel:
    """True multiprocessing queue wrapper."""
    def __init__(self):
        self._q = mp.Queue()
        
    def send(self, value):
        self._q.put(value)
        
    def recv(self):
        return self._q.get()
        
    def empty(self):
        return self._q.empty()

class ConcurrencyModule:
    """Actor model / multiprocessing concurrency for Nexa."""
    
    def spawn(self, func: Any, *args) -> mp.Process:
        """
        Spawn a new OS process to run the given Nexa function bypassing the GIL.
        """
        if hasattr(func, "_nexa_fn"):
            func = func._nexa_fn
            
        if func.__class__.__name__ != "NexaFunction":
            raise Exception(f"spawn() requires a Nexa function as the first argument, got {type(func)}")
            
        # Extract pickleable state
        # We cannot easily pickle the whole global environment.
        # So we extract just the immediate closure vars that are primitives.
        safe_closure = {}
        if func.closure:
            for k, v in func.closure._vars.items():
                if isinstance(v, (int, float, str, bool, type(None), list, dict)):
                    safe_closure[k] = v
        
        # We need the source_dir for imports inside the child process
        # Wait, how to grab it without passing the interpreter?
        # A bit hacky, but we can default to the current cwd if not available.
        process = mp.Process(
            target=_child_process_worker,
            args=(func.params, func.body, safe_closure, ".", None, args)
        )
        process.start()
        return process

    def Channel(self):
        return Channel()

    def __repr__(self):
        return "<concurrency module>"

