#!/usr/bin/env python3

# Copyright 2025-2026 Dimensional Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import threading
import time

from dimos_lcm.sensor_msgs import CameraInfo
import numpy as np

from dimos.core import In, Module, Out, rpc
from dimos.core.global_config import GlobalConfig
from dimos.msgs.sensor_msgs import Image, ImageFormat
from dimos.utils.logging_config import setup_logger

logger = setup_logger()


class DepthModule(Module):
    """
    Depth module for Unitree Go2 that processes RGB images to generate depth using Metric3D.

    Subscribes to:
        - /go2/color_image: RGB camera images from Unitree
        - /go2/camera_info: Camera calibration information

    Publishes:
        - /go2/depth_image: Depth images generated by Metric3D
    """

    # LCM inputs
    color_image: In[Image]
    camera_info: In[CameraInfo]

    # LCM outputs
    depth_image: Out[Image]

    def __init__(  # type: ignore[no-untyped-def]
        self,
        gt_depth_scale: float = 0.5,
        global_config: GlobalConfig | None = None,
        **kwargs,
    ) -> None:
        """
        Initialize Depth Module.

        Args:
            gt_depth_scale: Ground truth depth scaling factor
        """
        super().__init__(**kwargs)

        self.camera_intrinsics = None
        self.gt_depth_scale = gt_depth_scale
        self.metric3d = None
        self._camera_info_received = False

        # Processing state
        self._running = False
        self._latest_frame = None
        self._last_image = None
        self._last_timestamp = None
        self._last_depth = None
        self._cannot_process_depth = False

        # Threading
        self._processing_thread: threading.Thread | None = None
        self._stop_processing = threading.Event()

        if global_config:
            if global_config.simulation:
                self.gt_depth_scale = 1.0

    @rpc
    def start(self) -> None:
        super().start()

        if self._running:
            logger.warning("Camera module already running")
            return

        # Set running flag before starting
        self._running = True

        # Subscribe to video and camera info inputs
        self.color_image.subscribe(self._on_video)
        self.camera_info.subscribe(self._on_camera_info)

        # Start processing thread
        self._start_processing_thread()

        logger.info("Depth module started")

    @rpc
    def stop(self) -> None:
        if not self._running:
            return

        self._running = False
        self._stop_processing.set()

        # Wait for thread to finish
        if self._processing_thread and self._processing_thread.is_alive():
            self._processing_thread.join(timeout=2.0)

        super().stop()

    def _on_camera_info(self, msg: CameraInfo) -> None:
        """Process camera info to extract intrinsics."""
        if self.metric3d is not None:
            return  # Already initialized

        try:
            # Extract intrinsics from camera matrix K
            K = msg.K
            fx = K[0]
            fy = K[4]
            cx = K[2]
            cy = K[5]

            self.camera_intrinsics = [fx, fy, cx, cy]  # type: ignore[assignment]

            # Initialize Metric3D with camera intrinsics
            from dimos.models.depth.metric3d import Metric3D

            self.metric3d = Metric3D(camera_intrinsics=self.camera_intrinsics)  # type: ignore[assignment]
            self._camera_info_received = True

            logger.info(
                f"Initialized Metric3D with intrinsics from camera_info: {self.camera_intrinsics}"
            )

        except Exception as e:
            logger.error(f"Error processing camera info: {e}")

    def _on_video(self, msg: Image) -> None:
        """Store latest video frame for processing."""
        if not self._running:
            return

        # Simply store the latest frame - processing happens in main loop
        self._latest_frame = msg  # type: ignore[assignment]
        logger.debug(
            f"Received video frame: format={msg.format}, shape={msg.data.shape if hasattr(msg.data, 'shape') else 'unknown'}"
        )

    def _start_processing_thread(self) -> None:
        """Start the processing thread."""
        self._stop_processing.clear()
        self._processing_thread = threading.Thread(target=self._main_processing_loop, daemon=True)
        self._processing_thread.start()
        logger.info("Started depth processing thread")

    def _main_processing_loop(self) -> None:
        """Main processing loop that continuously processes latest frames."""
        logger.info("Starting main processing loop")

        while not self._stop_processing.is_set():
            # Process latest frame if available
            if self._latest_frame is not None:
                try:
                    msg = self._latest_frame
                    self._latest_frame = None  # Clear to avoid reprocessing
                    # Store for publishing
                    self._last_image = msg.data
                    self._last_timestamp = msg.ts if msg.ts else time.time()
                    # Process depth
                    self._process_depth(self._last_image)

                except Exception as e:
                    logger.error(f"Error in main processing loop: {e}", exc_info=True)
            else:
                # Small sleep to avoid busy waiting
                time.sleep(0.001)

        logger.info("Main processing loop stopped")

    def _process_depth(self, img_array: np.ndarray) -> None:  # type: ignore[type-arg]
        """Process depth estimation using Metric3D."""
        if self._cannot_process_depth:
            self._last_depth = None
            return

        # Wait for camera info to initialize Metric3D
        if self.metric3d is None:
            logger.debug("Waiting for camera_info to initialize Metric3D")
            return

        try:
            logger.debug(f"Processing depth for image shape: {img_array.shape}")

            # Generate depth map
            depth_array = self.metric3d.infer_depth(img_array) * self.gt_depth_scale

            self._last_depth = depth_array
            logger.debug(f"Generated depth map shape: {depth_array.shape}")

            self._publish_depth()

        except Exception as e:
            logger.error(f"Error processing depth: {e}")
            self._cannot_process_depth = True

    def _publish_depth(self) -> None:
        """Publish depth image."""
        if not self._running:
            return

        try:
            # Publish depth image
            if self._last_depth is not None:
                # Convert depth to uint16 (millimeters) for more efficient storage
                # Clamp to valid range [0, 65.535] meters before converting
                depth_clamped = np.clip(self._last_depth, 0, 65.535)
                depth_uint16 = (depth_clamped * 1000).astype(np.uint16)
                depth_msg = Image(
                    data=depth_uint16,
                    format=ImageFormat.DEPTH16,  # Use DEPTH16 format for uint16 depth
                    frame_id="camera_link",
                    ts=self._last_timestamp,
                )
                self.depth_image.publish(depth_msg)
                logger.debug(f"Published depth image (uint16): shape={depth_uint16.shape}")

        except Exception as e:
            logger.error(f"Error publishing depth data: {e}", exc_info=True)


depth_module = DepthModule.blueprint


__all__ = ["DepthModule", "depth_module"]
