import os
import numpy as np
import time
import io
import base64
import requests
from PIL import Image
from egogym.utils import pixel_to_world

MOONDREAM_API_URL = "https://api.moondream.ai/v1/point"


class UnprivilegedMoondream:
    
    def __init__(self, env, api_key=None):
        self.env = env
        self.unprivileged_T_world_object = None
        
        if api_key is None:
            api_key = os.environ.get("MOONDREAM_API_KEY")
        
        self.api_key = api_key
    
    def _call_api_with_retry(self, image_base64: str, object_name: str, max_retries=3, base_delay=2):
        """Call Moondream API with retry logic for transient errors."""
        headers = {
            "Content-Type": "application/json",
            "X-Moondream-Auth": self.api_key
        }
        
        payload = {
            "image_url": f"data:image/jpeg;base64,{image_base64}",
            "object": object_name
        }
        
        for attempt in range(max_retries):
            try:
                response = requests.post(MOONDREAM_API_URL, headers=headers, json=payload, timeout=30)
                
                if response.status_code == 200:
                    return response.json()
                elif response.status_code in [503, 429] and attempt < max_retries - 1:
                    delay = base_delay * (2 ** attempt)
                    print(f"API error {response.status_code}: {response.text}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
                    time.sleep(delay)
                    continue
                else:
                    response.raise_for_status()
            except requests.exceptions.RequestException as e:
                if attempt < max_retries - 1:
                    delay = base_delay * (2 ** attempt)
                    print(f"Request error: {e}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
                    time.sleep(delay)
                    continue
                else:
                    raise
        
        return None
    
    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
        """Use Moondream API to locate object in image and return normalized coordinates."""
        if rgb.dtype == np.float32 or rgb.dtype == np.float64:
            rgb = (rgb * 255).astype(np.uint8)
        image = Image.fromarray(rgb)
        
        try:
            buf = io.BytesIO()
            image.save(buf, format='JPEG')
            image_bytes = buf.getvalue()
            image_base64 = base64.b64encode(image_bytes).decode('utf-8')
            
            response = self._call_api_with_retry(image_base64, object_name)
            
            if response and "points" in response and len(response["points"]) > 0:
                point = response["points"][0]
                x_norm = float(point["x"])
                y_norm = float(point["y"])
                
                x_norm = max(0.0, min(1.0, x_norm))
                y_norm = max(0.0, min(1.0, y_norm))
                
                return np.array([x_norm, y_norm], dtype=np.float32)
            
            return np.array([0.5, 0.5], dtype=np.float32)
            
        except Exception as e:
            print(f"Moondream API error: {e}")
            return np.array([0.5, 0.5], dtype=np.float32)

    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):
        self.env.close()
    
    def __getattr__(self, name):
        return getattr(self.env, name)
