import uuid
import time
import random
import concurrent.futures
from typing import Any, Dict, Callable
from dataclasses import dataclass
from multiprocessing import Manager

@dataclass
class ObjectRef:
    """A handle to a result that may not yet be computed (a Future)."""
    id: str

class RayRuntime:
    """Simplified simulation of the Ray Distributed Execution Engine."""
    def __init__(self, num_workers: int = 4):
        self._manager = Manager()
        self._object_store = self._manager.dict() 
        self._actor_registry = self._manager.dict()
        self._executor = concurrent.futures.ProcessPoolExecutor(max_workers=num_workers)
        print(f"[RayRuntime] Initialized with {num_workers} workers.")

    def remote_task(self, func: Callable, *args, **kwargs) -> ObjectRef:
        """Schedules a stateless TASK."""
        ref_id = str(uuid.uuid4())
        ref = ObjectRef(ref_id)
        future = self._executor.submit(self._execute_task, func, *args, **kwargs)
        future.add_done_callback(lambda f: self._object_store.update({ref_id: f.result()}))
        return ref

    def _execute_task(self, func, *args, **kwargs):
        return func(*args, **kwargs)

    def create_actor(self, actor_class, *args, **kwargs):
        """Instantiates a stateful ACTOR."""
        actor_id = str(uuid.uuid4())
        actor_instance = actor_class(*args, **kwargs)
        self._actor_registry[actor_id] = actor_instance
        return RayActorHandle(actor_id, self)

    def get(self, ref: ObjectRef) -> Any:
        """Blocking call to retrieve value from the object store."""
        while ref.id not in self._object_store:
            time.sleep(0.01)
        return self._object_store[ref.id]

class RayActorHandle:
    """Proxy handle to route calls to a specific remote Actor."""
    def __init__(self, actor_id: str, runtime: RayRuntime):
        self.actor_id = actor_id
        self.runtime = runtime

    def __getattr__(self, name):
        def wrapper(*args, **kwargs):
            actor = self.runtime._actor_registry[self.actor_id]
            method = getattr(actor, name)
            return method(*args, **kwargs)
        return wrapper

# --- RL Pipeline Demonstration ---

class ParameterServer:
    """Stateful Actor: Maintains global model weights."""
    def __init__(self, lr=0.01):
        self.weights = 0.0 
        self.lr = lr

    def apply_gradients(self, grad):
        self.weights -= self.lr * grad
        return self.weights

    def get_weights(self):
        return self.weights

def simulate_environment(policy_weight: float, env_id: int):
    """Stateless Task: Simulates an environment rollout."""
    time.sleep(random.uniform(0.1, 0.3))
    gradient = (policy_weight - 10.0) * 1.0 # Target weight is 10.0
    print(f"[Task] Env {env_id} computed gradient: {gradient:.4f}")
    return gradient

if __name__ == '__main__':
    ray = RayRuntime(num_workers=4)
    ps = ray.create_actor(ParameterServer, lr=0.1)

    for epoch in range(3):
        print(f"\n--- Epoch {epoch} ---")
        current_weights = ps.get_weights()
        
        # Parallel Simulation (Tasks)
        task_refs = [ray.remote_task(simulate_environment, current_weights, i) for i in range(4)]
        
        # Collect and Update (Actor)
        gradients = [ray.get(ref) for ref in task_refs]
        avg_grad = sum(gradients) / len(gradients)
        new_weight = ps.apply_gradients(avg_grad)
        print(f"Updated Weight: {new_weight:.4f}")