#!/usr/bin/env python3
"""
EXIF orientation handling utilities for PunimTag
Handles image orientation correction based on EXIF data
"""

from PIL import Image
from typing import Tuple, Optional, Dict


class EXIFOrientationHandler:
    """Handles EXIF orientation detection and correction"""
    
    # EXIF orientation tag mapping
    ORIENTATION_TAG = 274  # EXIF orientation tag ID
    
    # Orientation values and their corresponding transformations
    ORIENTATION_TRANSFORMS = {
        1: None,  # Normal (no rotation needed)
        2: Image.Transpose.FLIP_LEFT_RIGHT,  # Mirrored horizontally
        3: Image.Transpose.ROTATE_180,  # Rotated 180°
        4: Image.Transpose.FLIP_TOP_BOTTOM,  # Mirrored vertically
        5: [Image.Transpose.FLIP_LEFT_RIGHT, Image.Transpose.ROTATE_90],  # Mirrored horizontally + rotated 90° CCW
        6: Image.Transpose.ROTATE_270,  # Rotated 90° CW
        7: [Image.Transpose.FLIP_LEFT_RIGHT, Image.Transpose.ROTATE_270],  # Mirrored horizontally + rotated 90° CW
        8: Image.Transpose.ROTATE_90,  # Rotated 90° CCW
    }
    
    @staticmethod
    def get_exif_orientation(image_path: str) -> Optional[int]:
        """Get EXIF orientation value from image
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Orientation value (1-8) or None if not found/invalid
        """
        try:
            with Image.open(image_path) as image:
                exif = image.getexif()
                if exif is None:
                    return None
                
                # Get orientation value
                orientation = exif.get(EXIFOrientationHandler.ORIENTATION_TAG)
                return orientation if orientation in EXIFOrientationHandler.ORIENTATION_TRANSFORMS else None
                
        except Exception:
            return None
    
    @staticmethod
    def get_orientation_info(image_path: str) -> Dict[str, any]:
        """Get comprehensive orientation information from image
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary with orientation information
        """
        orientation = EXIFOrientationHandler.get_exif_orientation(image_path)
        
        return {
            'orientation': orientation,
            'needs_correction': orientation is not None and orientation != 1,
            'transforms': EXIFOrientationHandler.ORIENTATION_TRANSFORMS.get(orientation, None),
            'description': EXIFOrientationHandler._get_orientation_description(orientation)
        }
    
    @staticmethod
    def _get_orientation_description(orientation: Optional[int]) -> str:
        """Get human-readable description of orientation"""
        descriptions = {
            1: "Normal (no rotation)",
            2: "Mirrored horizontally",
            3: "Rotated 180°",
            4: "Mirrored vertically", 
            5: "Mirrored horizontally + rotated 90° CCW",
            6: "Rotated 90° CW",
            7: "Mirrored horizontally + rotated 90° CW",
            8: "Rotated 90° CCW",
            None: "Unknown or no EXIF data"
        }
        return descriptions.get(orientation, "Invalid orientation")
    
    @staticmethod
    def correct_image_orientation(image: Image.Image, orientation: Optional[int] = None) -> Image.Image:
        """Correct image orientation based on EXIF data
        
        Args:
            image: PIL Image object
            orientation: EXIF orientation value (if None, will be detected from image)
            
        Returns:
            Corrected PIL Image object
        """
        if orientation is None:
            # Try to get orientation from image's EXIF data
            exif = image.getexif()
            if exif is None:
                return image
            orientation = exif.get(EXIFOrientationHandler.ORIENTATION_TAG)
        
        if orientation is None or orientation == 1:
            return image  # No correction needed
        
        transforms = EXIFOrientationHandler.ORIENTATION_TRANSFORMS.get(orientation)
        if transforms is None:
            return image  # No known transformation
        
        corrected_image = image.copy()
        
        # Apply transformation(s)
        if isinstance(transforms, list):
            # Multiple transformations (e.g., mirror + rotate)
            for transform in transforms:
                corrected_image = corrected_image.transpose(transform)
        else:
            # Single transformation
            corrected_image = corrected_image.transpose(transforms)
        
        return corrected_image
    
    @staticmethod
    def correct_image_orientation_from_path(image_path: str) -> Tuple[Image.Image, Optional[int]]:
        """Load image and correct its orientation
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Tuple of (corrected_image, original_orientation)
        """
        try:
            with Image.open(image_path) as image:
                orientation = EXIFOrientationHandler.get_exif_orientation(image_path)
                corrected_image = EXIFOrientationHandler.correct_image_orientation(image, orientation)
                return corrected_image, orientation
        except Exception:
            # If we can't load the image, return None
            return None, None
    
    @staticmethod
    def save_corrected_image(image: Image.Image, output_path: str, orientation: Optional[int] = None) -> bool:
        """Save image with corrected orientation and remove EXIF orientation tag
        
        Args:
            image: Corrected PIL Image object
            output_path: Path to save the corrected image
            orientation: Original orientation value (for logging)
            
        Returns:
            True if successful, False otherwise
        """
        try:
            # Create a copy to avoid modifying the original
            save_image = image.copy()
            
            # Save without EXIF data to avoid orientation issues
            save_image.save(output_path, quality=95)
            return True
            
        except Exception:
            return False
    
    @staticmethod
    def get_corrected_image_dimensions(original_width: int, original_height: int, orientation: Optional[int]) -> Tuple[int, int]:
        """Get dimensions after orientation correction
        
        Args:
            original_width: Original image width
            original_height: Original image height
            orientation: EXIF orientation value
            
        Returns:
            Tuple of (corrected_width, corrected_height)
        """
        if orientation is None or orientation == 1:
            return original_width, original_height
        
        # For 90° and 270° rotations, width and height are swapped
        if orientation in [5, 6, 7, 8]:
            return original_height, original_width
        
        # For 180° rotation and mirrors, dimensions stay the same
        return original_width, original_height
    
    @staticmethod
    def transform_face_coordinates(face_coords: Dict[str, int], original_width: int, original_height: int, 
                                 orientation: Optional[int]) -> Dict[str, int]:
        """Transform face coordinates based on image orientation correction
        
        Args:
            face_coords: Face coordinates in DeepFace format {x, y, w, h}
            original_width: Original image width
            original_height: Original image height
            orientation: EXIF orientation value
            
        Returns:
            Transformed face coordinates in DeepFace format
        """
        if orientation is None or orientation == 1:
            return face_coords.copy()
        
        x = face_coords['x']
        y = face_coords['y']
        w = face_coords['w']
        h = face_coords['h']
        
        # Calculate face center and corners
        center_x = x + w // 2
        center_y = y + h // 2
        
        if orientation == 2:  # Mirrored horizontally
            new_x = original_width - x - w
            new_y = y
        elif orientation == 3:  # Rotated 180°
            new_x = original_width - x - w
            new_y = original_height - y - h
        elif orientation == 4:  # Mirrored vertically
            new_x = x
            new_y = original_height - y - h
        elif orientation == 5:  # Mirrored horizontally + rotated 90° CCW
            new_x = original_height - y - h
            new_y = x
            new_w = h
            new_h = w
        elif orientation == 6:  # Rotated 90° CW
            new_x = y
            new_y = original_width - x - w
            new_w = h
            new_h = w
        elif orientation == 7:  # Mirrored horizontally + rotated 90° CW
            new_x = y
            new_y = x
            new_w = h
            new_h = w
        elif orientation == 8:  # Rotated 90° CCW
            new_x = original_height - y - h
            new_y = original_width - x - w
            new_w = h
            new_h = w
        else:
            return face_coords.copy()
        
        return {
            'x': new_x,
            'y': new_y,
            'w': new_w if 'new_w' in locals() else w,
            'h': new_h if 'new_h' in locals() else h
        }
