import asyncio
import numpy as np
import threading
from egogym.utils import pixel_to_world
from egogym.misc.molmo_client import MolmoClient

class UnprivilegedMolmo:
    
    def __init__(self, env, host="localhost", port=8765):
        self.env = env
        self.unprivileged_T_world_object = None
        self.molmo_client = MolmoClient(host=host, port=port)
        self._loop = None
        self._loop_thread = None
        self._connected = False
        
    def _start_event_loop(self):
        if self._loop is None:
            self._loop = asyncio.new_event_loop()
            self._loop_thread = threading.Thread(target=self._run_event_loop, daemon=True)
            self._loop_thread.start()
    
    def _run_event_loop(self):
        asyncio.set_event_loop(self._loop)
        self._loop.run_forever()
        
    def _ensure_connected(self):
        if not self._connected:
            try:
                self._start_event_loop()
                future = asyncio.run_coroutine_threadsafe(self.molmo_client.connect(), self._loop)
                future.result(timeout=10)
                self._connected = True
            except Exception as e:
                raise RuntimeError(
                    f"Failed to connect to Molmo server at {self.molmo_client.uri}. "
                    f"Make sure the Molmo server is running with: "
                    f"'python baselines/rum/molmo_server.py'. "
                    f"Original error: {e}"
                ) from e
    
    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
        self._ensure_connected()
        future = asyncio.run_coroutine_threadsafe(
            self.molmo_client.infer_point(rgb, object_name=object_name), 
            self._loop
        )
        return future.result(timeout=30)
        
    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        
        robot = self.env.unwrapped.robot
        robot.rgb_renderers[robot.camera_names[0]].enable_depth_rendering()
        robot.rgb_renderers[robot.camera_names[0]].update_scene(robot.data, robot.camera_names[0])
        depth = robot.rgb_renderers[robot.camera_names[0]].render()
        robot.rgb_renderers[robot.camera_names[0]].disable_depth_rendering()

        point_norm = self._infer_point_sync(obs["rgb_ego"], obs["object_name"])
        x_norm, y_norm = point_norm
        x = int(x_norm * self.env.unwrapped.render_width)
        y = int(y_norm * self.env.unwrapped.render_height)
        depth_value = depth[y, x] + 0.03
        self.unprivileged_T_world_object = pixel_to_world(x, y, depth_value, self.env.unwrapped.model, self.env.unwrapped.data, robot.camera_names[0], self.env.unwrapped.render_width, self.env.unwrapped.render_height)

        return obs, info
    
    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        
        object_pose = obs["object_pose"].reshape(4,4)
        object_pose[:3, 3] = self.unprivileged_T_world_object
        obs["object_pose"] = object_pose.flatten()

        return obs, reward, terminated, truncated, info
    
    def close(self):
        if self._connected:
            future = asyncio.run_coroutine_threadsafe(self.molmo_client.disconnect(), self._loop)
            future.result(timeout=5)
            self._connected = False
        if self._loop is not None:
            self._loop.call_soon_threadsafe(self._loop.stop)
            self._loop_thread.join(timeout=2)
        self.env.close()
    
    def __getattr__(self, name):
        return getattr(self.env, name)
