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

MODEL = "gpt-4o"
PROMPT = """Get all points matching the following object: {object_name}. The label returned should be an identifying name for the object detected.
The answer should follow the json format: [{{"point": [y, x], "label": "{object_name}"}}, ...]. The points are in [y, x] format normalized to 0-1000."""

class UnprivilegedChatGPT:
    
    def __init__(self, env, api_key=None):
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError("openai package is not installed. Please install it with `pip install openai`.")

        self.env = env
        self.unprivileged_T_world_object = None
        
        if api_key is None:
            api_key = os.environ.get("OPENAI_API_KEY")
    
        
        self.client = OpenAI(api_key=api_key)
    
    def _call_api_with_retry(self, messages, max_retries=3, base_delay=2):
        """Call OpenAI API with retry logic for transient errors."""
        for attempt in range(max_retries):
            try:
                return self.client.chat.completions.create(
                    model=MODEL,
                    messages=messages,
                    temperature=0.5,
                    max_tokens=1024
                )
            except Exception as e:
                error_code = getattr(e, 'status_code', None)
                if error_code in [503, 429] and attempt < max_retries - 1:
                    delay = base_delay * (2 ** attempt)
                    print(f"API error {error_code}: {str(e)}. Retrying in {delay}s... (attempt {attempt + 1}/{max_retries})")
                    time.sleep(delay)
                    continue
                else:
                    raise
    
    def _parse_json_point_yx(self, text):
        """Parse JSON format with point coordinates in [y, x] format normalized to 0-1000."""
        try:
            json_match = re.search(r'```(?:json)?\s*(\[.*?\])\s*```', text, re.DOTALL)
            if json_match:
                json_str = json_match.group(1)
            else:
                json_match = re.search(r'(\[.*?\])', text, re.DOTALL)
                if json_match:
                    json_str = json_match.group(1)
                else:
                    return None
            
            data = json.loads(json_str)
            if not isinstance(data, list) or len(data) == 0:
                return None
            
            point_data = data[0]
            if 'point' in point_data:
                point = point_data['point']
                if isinstance(point, list) and len(point) >= 2:
                    return (float(point[0]), float(point[1]))
            
            return None
            
        except (json.JSONDecodeError, KeyError, ValueError, IndexError) as e:
            return None
    
    def _infer_point_sync(self, rgb: np.ndarray, object_name: str) -> np.ndarray:
        """Use OpenAI GPT-4o 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)
        img_width, img_height = image.size
        
        prompt = PROMPT.format(object_name=object_name)
        
        try:
            buf = io.BytesIO()
            image.save(buf, format='PNG')
            image_bytes = buf.getvalue()
            base64_image = base64.b64encode(image_bytes).decode('utf-8')
            
            messages = [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{base64_image}"
                            }
                        },
                        {
                            "type": "text",
                            "text": prompt
                        }
                    ]
                }
            ]
            
            response = self._call_api_with_retry(messages)
            generated_text = response.choices[0].message.content
            
            json_point = self._parse_json_point_yx(generated_text)
            if json_point:
                y_pos, x_pos = json_point
                x_norm = float(x_pos) / 1000.0
                y_norm = float(y_pos) / 1000.0
                
                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:
            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)
