import cv2
import numpy as np
import onnxruntime as ort
import os
import urllib.request

class DepthMapper:
    """
    Module: Spatial Engine (ONNX Powered)
    Gives Aelion OS true 3D perception using MiDaS v2.1 Small.
    """
    def __init__(self):
        current_dir = os.path.dirname(os.path.abspath(__file__))
        self.model_path = os.path.join(current_dir, "midas-small.onnx")

        # --- THE AUTO-DOWNLOADER ---
        if not os.path.exists(self.model_path):
            self._download_model()

        print("[Aelion Vision] Booting Spatial Depth Engine...")
        try:
            self.session = ort.InferenceSession(self.model_path)
        except Exception as e:
            print(f"[!] Critical Error: Could not load spatial weights. {e}")
            self.session = None

    def _download_model(self):
        """Grabs the MiDaS v2.1 Small ONNX weights."""
        print("[Aelion Vision] Initializing 3D Perception for the first time...")
        print("[Aelion Vision] Downloading MiDaS Spatial Weights (~45MB). Hold tight...")
        
        # Temporary official URL. Swap this with your Aelion Server URL later!
        url = "https://ai.aelionyxe.top/models/model-small.onnx"
        
        try:
            urllib.request.urlretrieve(url, self.model_path)
            print("[Aelion Vision] Download Complete! Spatial Engine online.")
        except Exception as e:
            print(f"[!] Critical Error: Failed to download the AI model. {e}")

    def get_depth_map(self, frame):
        """Converts a standard 2D camera frame into a full 3D depth map."""
        if self.session is None or frame is None:
            return None

        try:
            img_h, img_w = frame.shape[:2]

            # 1. Formatting: MiDaS small specifically wants 256x256 RGB images
            img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (256, 256))
            img = img.astype(np.float32) / 255.0

            # 2. Normalization: Standard ImageNet values (MiDaS expects this exact math)
            mean = [0.485, 0.456, 0.406]
            std = [0.229, 0.224, 0.225]
            img = (img - mean) / std

            # 3. Reorder dimensions from HWC to NCHW (Batch, Channels, Height, Width)
            img = np.transpose(img, (2, 0, 1))
            img = np.expand_dims(img, axis=0)

            # 4. Run the AI Brain
            inputs = {self.session.get_inputs()[0].name: img}
            outputs = self.session.run(None, inputs)
            
            # 5. Extract the raw map and stretch it back to match the webcam resolution
            depth_map = outputs[0][0]
            depth_map = cv2.resize(depth_map, (img_w, img_h))
            
            # 6. Normalize to 0-255 grayscale so we can actually draw it or read the numbers easily
            depth_min = depth_map.min()
            depth_max = depth_map.max()
            depth_map = (255 * (depth_map - depth_min) / (depth_max - depth_min)).astype(np.uint8)

            return depth_map

        except Exception as e:
            print(f"[!] Spatial mapping failed this frame: {e}")
            return None

    def get_distance(self, depth_map, x, y):
        """
        Returns a relative depth score for a specific pixel.
        """
        if depth_map is None: return 0
        
        # Failsafe: Keep coordinates inside the screen bounds so it doesn't crash
        h, w = depth_map.shape
        x = max(0, min(x, w - 1))
        y = max(0, min(y, h - 1))

        # THE FIX: Don't just check one pixel. Grab a 5x5 patch around the coordinate 
        # and average it. Otherwise, the number twitches uncontrollably.
        x1, y1 = max(0, x - 2), max(0, y - 2)
        x2, y2 = min(w, x + 3), min(h, y + 3)
        
        patch = depth_map[y1:y2, x1:x2]
        
        if patch.size == 0: return 0
        
        # Return a clean integer
        return int(np.mean(patch))