"""
Camera capture module for Unitree G1 robot
Handles OpenCV camera interface and image processing for AI analysis
"""

import cv2
import numpy as np
import base64
import io
import os
import time
from pathlib import Path
from PIL import Image
from typing import Optional, Union, List
import logging
from Core.settings import CAMERA_INDEX, FRAME_WIDTH, FRAME_HEIGHT, FPS

class CameraCapture:
    """
    Camera capture class for Unitree G1 robot
    Provides methods to capture frames and prepare them for AI analysis
    """
    
    def __init__(self, camera_index: Optional[Union[int, str]] = None):
        """
        Initialize camera capture
        
        Args:
            camera_index: int index, /dev/videoX path, "auto", or None
        """
        self.camera_source = self._resolve_camera_source(camera_index)
        # Keep legacy field name for backward compatibility with older modules/tests
        self.camera_index = self.camera_source
        self.cap = None
        self.is_initialized = False
        self.logger = logging.getLogger(__name__)
        self._last_read_warn_ts = 0.0

    def _resolve_camera_source(self, camera_source: Optional[Union[int, str]]) -> Union[int, str]:
        """
        Resolve requested camera source from argument/env/defaults.
        Priority:
          1) explicit argument
          2) CAMERA_DEVICE env (supports /dev/videoX or "auto")
          3) CAMERA_INDEX env/config int
        """
        if camera_source is None:
            env_device = os.environ.get("CAMERA_DEVICE", "").strip()
            if env_device:
                camera_source = env_device
            else:
                camera_source = os.environ.get("CAMERA_INDEX", str(CAMERA_INDEX))

        if isinstance(camera_source, str):
            s = camera_source.strip()
            if not s:
                return int(CAMERA_INDEX)
            if s.lower() == "auto":
                return "auto"
            try:
                return int(s)
            except ValueError:
                return s

        return int(camera_source)

    def _candidate_sources(self) -> List[Union[int, str]]:
        """
        Build candidate capture sources. When CAMERA_FALLBACK_SCAN=1 (default),
        we also try available /dev/video* nodes if the primary source fails.
        """
        src = self.camera_source
        candidates: List[Union[int, str]] = [src]

        # If primary source is /dev/videoX, also try numeric index X.
        if isinstance(src, str) and src.startswith("/dev/video"):
            try:
                idx = int(src.replace("/dev/video", ""))
                candidates.append(idx)
            except Exception:
                pass

        do_scan = os.environ.get("CAMERA_FALLBACK_SCAN", "1").strip().lower() not in (
            "0", "false", "no", "off"
        )
        if src == "auto" or do_scan:
            for p in sorted(Path("/dev").glob("video*")):
                sp = str(p)
                if sp not in candidates:
                    candidates.append(sp)
                try:
                    idx = int(sp.replace("/dev/video", ""))
                    if idx not in candidates:
                        candidates.append(idx)
                except Exception:
                    pass

        return candidates

    def _open_capture(self, source: Union[int, str]):
        if isinstance(source, str) and source.startswith("/dev/video"):
            # V4L2 backend is more reliable for /dev/videoX on Linux.
            return cv2.VideoCapture(source, cv2.CAP_V4L2)
        return cv2.VideoCapture(source)
        
    def initialize(self) -> bool:
        """
        Initialize camera connection
        
        Returns:
            bool: True if successful, False otherwise
        """
        try:
            self.release()
            tried: List[str] = []

            for source in self._candidate_sources():
                tried.append(str(source))
                cap = None
                try:
                    cap = self._open_capture(source)
                except Exception:
                    cap = None

                if cap is None or not cap.isOpened():
                    try:
                        if cap is not None:
                            cap.release()
                    except Exception:
                        pass
                    continue

                # Set camera properties
                cap.set(cv2.CAP_PROP_FRAME_WIDTH, FRAME_WIDTH)
                cap.set(cv2.CAP_PROP_FRAME_HEIGHT, FRAME_HEIGHT)
                cap.set(cv2.CAP_PROP_FPS, FPS)

                # Test capture
                ret, _frame = cap.read()
                if not ret:
                    cap.release()
                    continue

                self.cap = cap
                self.camera_source = source
                self.camera_index = source
                self.is_initialized = True
                self.logger.info(f"Camera initialized successfully: {source}")
                return True

            self.is_initialized = False
            self.cap = None
            self.logger.error(f"Failed to open camera. Tried sources: {', '.join(tried)}")
            return False
            
        except Exception as e:
            self.logger.error(f"Camera initialization failed: {e}")
            return False
    
    def capture_frame(self) -> Optional[np.ndarray]:
        """
        Capture a single frame from camera
        
        Returns:
            numpy.ndarray: Captured frame or None if failed
        """
        if not self.is_initialized or self.cap is None:
            return None
            
        try:
            ret, frame = self.cap.read()
            if not ret:
                now = time.time()
                # Throttle noisy warnings when camera stream is unstable/unavailable.
                if now - self._last_read_warn_ts >= 5.0:
                    self.logger.warning("Failed to capture frame from camera")
                    self._last_read_warn_ts = now
                return None
                
            return frame
            
        except Exception as e:
            self.logger.error(f"Frame capture failed: {e}")
            return None
    
    def frame_to_base64(self, frame: np.ndarray, format: str = 'JPEG') -> Optional[str]:
        """
        Convert OpenCV frame to base64 string for API transmission
        
        Args:
            frame: OpenCV frame (BGR format)
            format: Image format ('JPEG' or 'PNG')
            
        Returns:
            str: Base64 encoded image or None if failed
        """
        try:
            # Convert BGR to RGB
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # Convert to PIL Image
            pil_image = Image.fromarray(rgb_frame)
            
            # Convert to bytes
            buffer = io.BytesIO()
            pil_image.save(buffer, format=format, quality=85)
            buffer.seek(0)
            
            # Encode to base64
            img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            
            return img_base64
            
        except Exception as e:
            self.logger.error(f"Base64 conversion failed: {e}")
            return None
    
    def frame_to_pil(self, frame: np.ndarray) -> Optional[Image.Image]:
        """
        Convert OpenCV frame to PIL Image for Gemini API
        
        Args:
            frame: OpenCV frame (BGR format)
            
        Returns:
            PIL.Image: Converted image or None if failed
        """
        try:
            # Convert BGR to RGB
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # Convert to PIL Image
            pil_image = Image.fromarray(rgb_frame)
            
            return pil_image
            
        except Exception as e:
            self.logger.error(f"PIL conversion failed: {e}")
            return None
    
    def get_frame_info(self) -> dict:
        """
        Get current camera frame information
        
        Returns:
            dict: Frame information including resolution, FPS
        """
        if not self.is_initialized or self.cap is None:
            return {}
            
        try:
            info = {
                'width': int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
                'height': int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
                'fps': self.cap.get(cv2.CAP_PROP_FPS),
                'format': int(self.cap.get(cv2.CAP_PROP_FORMAT)),
                'brightness': self.cap.get(cv2.CAP_PROP_BRIGHTNESS),
                'contrast': self.cap.get(cv2.CAP_PROP_CONTRAST),
                'saturation': self.cap.get(cv2.CAP_PROP_SATURATION)
            }
            return info
            
        except Exception as e:
            self.logger.error(f"Failed to get frame info: {e}")
            return {}
    
    def release(self):
        """
        Release camera resources
        """
        try:
            if self.cap is not None:
                self.cap.release()
                self.cap = None
                self.logger.info("Camera released successfully")
                
            self.is_initialized = False
            
        except Exception as e:
            self.logger.error(f"Camera release failed: {e}")
        finally:
            # Headless systems often ship OpenCV without HighGUI window backends.
            # Do not fail camera lifecycle because window teardown is unavailable.
            try:
                import cv2
                cv2.destroyAllWindows()
            except Exception:
                pass
    
    def __enter__(self):
        """Context manager entry"""
        self.initialize()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit"""
        self.release()
        return False  # Don't suppress exceptions

# Utility functions for camera testing
def test_camera_capture():
    """
    Test camera capture functionality
    """
    logging.basicConfig(level=logging.INFO)
    
    with CameraCapture() as camera:
        if not camera.is_initialized:
            print("Failed to initialize camera")
            return
            
        print("Camera info:", camera.get_frame_info())
        
        # Capture test frame
        frame = camera.capture_frame()
        if frame is not None:
            print(f"Captured frame shape: {frame.shape}")
            
            # Test base64 conversion
            base64_str = camera.frame_to_base64(frame)
            if base64_str:
                print(f"Base64 conversion successful, length: {len(base64_str)}")
            
            # Test PIL conversion
            pil_img = camera.frame_to_pil(frame)
            if pil_img:
                print(f"PIL conversion successful, size: {pil_img.size}")
                
            # Save test image
            cv2.imwrite('test_capture.jpg', frame)
            print("Test image saved as test_capture.jpg")
        else:
            print("Failed to capture frame")
