import cv2
import numpy as np
from pathlib import Path
import math
from typing import Optional, Union, List, Tuple, Dict

# --- Basic Power-of-Two Utilities ---

def is_power_of_two(n: int) -> bool:
    """Checks if a number is a power of two."""
    return (n > 0) and (n & (n - 1) == 0)

def get_nearest_pot(value: int) -> int:
    """Finds the nearest power of two to the given value."""
    if value <= 0:
        return 1  # POT must be positive, return 1 as a fallback
    if is_power_of_two(value):
        return value

    lower_pot = 1 << (value.bit_length() - 1)
    upper_pot = 1 << value.bit_length()

    if (value - lower_pot) < (upper_pot - value):
        return lower_pot
    else:
        return upper_pot

def get_nearest_power_of_two_downscale(value: int) -> int:
    """
    Finds the nearest power of two that is less than or equal to the given value.
    If the value is already a power of two, it returns the value itself.
    Returns 1 if the value is less than 1.
    """
    if value &lt; 1:
        return 1
    if is_power_of_two(value):
        return value
    # Find the largest power of two strictly less than value,
    # unless value itself is POT.
    # (1 &lt;&lt; (value.bit_length() - 1)) achieves this.
    # Example: value=7 (0111, bl=3), 1&lt;&lt;2 = 4.
    # Example: value=8 (1000, bl=4), 1&lt;&lt;3 = 8.
    # Example: value=9 (1001, bl=4), 1&lt;&lt;3 = 8.
    return 1 &lt;&lt; (value.bit_length() - 1)
# --- Dimension Calculation ---

def calculate_target_dimensions(
    original_width: int,
    original_height: int,
    target_width: Optional[int] = None,
    target_height: Optional[int] = None,
    resize_mode: str = "fit",  # e.g., "fit", "stretch", "max_dim_pot"
    ensure_pot: bool = False,
    allow_upscale: bool = False,
    target_max_dim_for_pot_mode: Optional[int] = None # Specific for "max_dim_pot"
) -> Tuple[int, int]:
    """
    Calculates target dimensions based on various modes and constraints.

    Args:
        original_width: Original width of the image.
        original_height: Original height of the image.
        target_width: Desired target width.
        target_height: Desired target height.
        resize_mode:
            - "fit": Scales to fit within target_width/target_height, maintaining aspect ratio.
                     Requires at least one of target_width or target_height.
            - "stretch": Scales to exactly target_width and target_height, ignoring aspect ratio.
                         Requires both target_width and target_height.
            - "max_dim_pot": Scales to fit target_max_dim_for_pot_mode while maintaining aspect ratio,
                             then finds nearest POT for each dimension. Requires target_max_dim_for_pot_mode.
        ensure_pot: If True, final dimensions will be adjusted to the nearest power of two.
        allow_upscale: If False, dimensions will not exceed original dimensions unless ensure_pot forces it.
        target_max_dim_for_pot_mode: Max dimension to use when resize_mode is "max_dim_pot".

    Returns:
        A tuple (new_width, new_height).
    """
    if original_width <= 0 or original_height <= 0:
        # Fallback for invalid original dimensions
        fallback_dim = 1
        if ensure_pot:
            if target_width and target_height:
                fallback_dim = get_nearest_pot(max(target_width, target_height, 1))
            elif target_width:
                fallback_dim = get_nearest_pot(target_width)
            elif target_height:
                fallback_dim = get_nearest_pot(target_height)
            elif target_max_dim_for_pot_mode:
                 fallback_dim = get_nearest_pot(target_max_dim_for_pot_mode)
            else: # Default POT if no target given
                fallback_dim = 256
            return (fallback_dim, fallback_dim)
        return (target_width or 1, target_height or 1)


    w, h = original_width, original_height

    if resize_mode == "max_dim_pot":
        if target_max_dim_for_pot_mode is None:
            raise ValueError("target_max_dim_for_pot_mode must be provided for 'max_dim_pot' resize_mode.")
        
        # Logic adapted from old processing_engine.calculate_target_dimensions
        ratio = w / h
        if ratio > 1:  # Width is dominant
            scaled_w = target_max_dim_for_pot_mode
            scaled_h = max(1, round(scaled_w / ratio))
        else:  # Height is dominant or square
            scaled_h = target_max_dim_for_pot_mode
            scaled_w = max(1, round(scaled_h * ratio))
        
        # Upscale check for this mode is implicitly handled by target_max_dim
        # If ensure_pot is true (as it was in the original logic), it's applied here
        # For this mode, ensure_pot is effectively always true for the final step
        w = get_nearest_pot(scaled_w)
        h = get_nearest_pot(scaled_h)
        return int(w), int(h)

    elif resize_mode == "fit":
        if target_width is None and target_height is None:
            raise ValueError("At least one of target_width or target_height must be provided for 'fit' mode.")
        
        if target_width and target_height:
            ratio_orig = w / h
            ratio_target = target_width / target_height
            if ratio_orig > ratio_target: # Original is wider than target aspect
                w_new = target_width
                h_new = max(1, round(w_new / ratio_orig))
            else: # Original is taller or same aspect
                h_new = target_height
                w_new = max(1, round(h_new * ratio_orig))
        elif target_width:
            w_new = target_width
            h_new = max(1, round(w_new / (w / h)))
        else: # target_height is not None
            h_new = target_height
            w_new = max(1, round(h_new * (w / h)))
        w, h = w_new, h_new

    elif resize_mode == "stretch":
        if target_width is None or target_height is None:
            raise ValueError("Both target_width and target_height must be provided for 'stretch' mode.")
        w, h = target_width, target_height
    
    else:
        raise ValueError(f"Unsupported resize_mode: {resize_mode}")

    if not allow_upscale:
        if w > original_width: w = original_width
        if h > original_height: h = original_height
    
    if ensure_pot:
        w = get_nearest_pot(w)
        h = get_nearest_pot(h)
        # Re-check upscale if POT adjustment made it larger than original and not allowed
        if not allow_upscale:
            if w > original_width: w = get_nearest_pot(original_width) # Get closest POT to original
            if h > original_height: h = get_nearest_pot(original_height)


    return int(max(1, w)), int(max(1, h))


# --- Image Statistics ---

def calculate_image_stats(image_data: np.ndarray) -> Optional[Dict]:
    """
    Calculates min, max, mean for a given numpy image array.
    Handles grayscale and multi-channel images. Converts to float64 for calculation.
    Normalizes uint8/uint16 data to 0-1 range before calculating stats.
    """
    if image_data is None:
        return None
    try:
        data_float = image_data.astype(np.float64)

        if image_data.dtype == np.uint16:
            data_float /= 65535.0
        elif image_data.dtype == np.uint8:
            data_float /= 255.0
        
        stats = {}
        if len(data_float.shape) == 2:  # Grayscale (H, W)
            stats["min"] = float(np.min(data_float))
            stats["max"] = float(np.max(data_float))
            stats["mean"] = float(np.mean(data_float))
        elif len(data_float.shape) == 3:  # Color (H, W, C)
            stats["min"] = [float(v) for v in np.min(data_float, axis=(0, 1))]
            stats["max"] = [float(v) for v in np.max(data_float, axis=(0, 1))]
            stats["mean"] = [float(v) for v in np.mean(data_float, axis=(0, 1))]
        else:
            return None # Unsupported shape
        return stats
    except Exception:
        return {"error": "Error calculating image stats"}

# --- Aspect Ratio String ---

def normalize_aspect_ratio_change(original_width: int, original_height: int, resized_width: int, resized_height: int, decimals: int = 2) -> str:
    """
    Calculates the aspect ratio change string (e.g., "EVEN", "X133").
    """
    if original_width <= 0 or original_height <= 0:
        return "InvalidInput"
    if resized_width <= 0 or resized_height <= 0:
        return "InvalidResize"

    width_change_percentage = ((resized_width - original_width) / original_width) * 100
    height_change_percentage = ((resized_height - original_height) / original_height) * 100

    normalized_width_change = width_change_percentage / 100
    normalized_height_change = height_change_percentage / 100

    normalized_width_change = min(max(normalized_width_change + 1, 0), 2)
    normalized_height_change = min(max(normalized_height_change + 1, 0), 2)
    
    epsilon = 1e-9
    if abs(normalized_width_change) < epsilon and abs(normalized_height_change) < epsilon:
         closest_value_to_one = 1.0
    elif abs(normalized_width_change) < epsilon:
         closest_value_to_one = abs(normalized_height_change)
    elif abs(normalized_height_change) < epsilon:
         closest_value_to_one = abs(normalized_width_change)
    else:
         closest_value_to_one = min(abs(normalized_width_change), abs(normalized_height_change))

    scale_factor = 1 / (closest_value_to_one + epsilon) if abs(closest_value_to_one) < epsilon else 1 / closest_value_to_one

    scaled_normalized_width_change = scale_factor * normalized_width_change
    scaled_normalized_height_change = scale_factor * normalized_height_change

    output_width = round(scaled_normalized_width_change, decimals)
    output_height = round(scaled_normalized_height_change, decimals)

    if abs(output_width - 1.0) < epsilon: output_width = 1
    if abs(output_height - 1.0) < epsilon: output_height = 1

    if abs(output_width - output_height) < epsilon: # Handles original square or aspect maintained
        output = "EVEN"
    elif output_width != 1 and abs(output_height - 1.0) < epsilon : # Width changed, height maintained relative to width
        output = f"X{str(output_width).replace('.', '')}"
    elif output_height != 1 and abs(output_width - 1.0) < epsilon: # Height changed, width maintained relative to height
        output = f"Y{str(output_height).replace('.', '')}"
    else: # Both changed relative to each other
        output = f"X{str(output_width).replace('.', '')}Y{str(output_height).replace('.', '')}"
    return output

# --- Image Loading, Conversion, Resizing ---

def load_image(image_path: Union[str, Path], read_flag: int = cv2.IMREAD_UNCHANGED) -> Optional[np.ndarray]:
    """Loads an image from the specified path."""
    try:
        img = cv2.imread(str(image_path), read_flag)
        if img is None:
            # print(f"Warning: Failed to load image: {image_path}") # Optional: for debugging utils
            return None
        return img
    except Exception: # as e:
        # print(f"Error loading image {image_path}: {e}") # Optional: for debugging utils
        return None

def convert_bgr_to_rgb(image: np.ndarray) -> np.ndarray:
    """Converts an image from BGR to RGB color space."""
    if image is None or len(image.shape) < 3:
        return image # Return as is if not a color image or None
    
    if image.shape[2] == 4: # BGRA
        return cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
    elif image.shape[2] == 3: # BGR
        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    return image # Return as is if not 3 or 4 channels

def convert_rgb_to_bgr(image: np.ndarray) -> np.ndarray:
    """Converts an image from RGB to BGR color space."""
    if image is None or len(image.shape) < 3 or image.shape[2] != 3: # Only for 3-channel RGB
        return image # Return as is if not a 3-channel color image or None
    return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)


def resize_image(image: np.ndarray, target_width: int, target_height: int, interpolation: Optional[int] = None) -> np.ndarray:
    """Resizes an image to target_width and target_height."""
    if image is None:
        raise ValueError("Cannot resize a None image.")
    if target_width <= 0 or target_height <= 0:
        raise ValueError("Target width and height must be positive.")

    original_height, original_width = image.shape[:2]

    if interpolation is None:
        # Default interpolation: Lanczos for downscaling, Cubic for upscaling/same
        if (target_width * target_height) < (original_width * original_height):
            interpolation = cv2.INTER_LANCZOS4
        else:
            interpolation = cv2.INTER_CUBIC
            
    return cv2.resize(image, (target_width, target_height), interpolation=interpolation)

# --- Image Saving ---

def save_image(
    image_path: Union[str, Path],
    image_data: np.ndarray,
    output_format: Optional[str] = None, # e.g. "png", "jpg", "exr"
    output_dtype_target: Optional[np.dtype] = None, # e.g. np.uint8, np.uint16, np.float16
    params: Optional[List[int]] = None,
    convert_to_bgr_before_save: bool = True # True for most formats except EXR
) -> bool:
    """
    Saves image data to a file. Handles data type and color space conversions.

    Args:
        image_path: Path to save the image.
        image_data: NumPy array of the image.
        output_format: Desired output format (e.g., 'png', 'jpg'). If None, derived from extension.
        output_dtype_target: Target NumPy dtype for saving (e.g., np.uint8, np.uint16).
                             If None, tries to use image_data.dtype or a sensible default.
        params: OpenCV imwrite parameters (e.g., [cv2.IMWRITE_JPEG_QUALITY, 90]).
        convert_to_bgr_before_save: If True and image is 3-channel, converts RGB to BGR.
                                   Set to False for formats like EXR that expect RGB.

    Returns:
        True if saving was successful, False otherwise.
    """
    if image_data is None:
        return False
    
    img_to_save = image_data.copy()
    path_obj = Path(image_path)
    path_obj.parent.mkdir(parents=True, exist_ok=True)

    # 1. Data Type Conversion
    if output_dtype_target is not None:
        if output_dtype_target == np.uint8 and img_to_save.dtype != np.uint8:
            if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0 * 255.0).astype(np.uint8)
            elif img_to_save.dtype in [np.float16, np.float32, np.float64]: img_to_save = (np.clip(img_to_save, 0.0, 1.0) * 255.0).astype(np.uint8)
            else: img_to_save = img_to_save.astype(np.uint8)
        elif output_dtype_target == np.uint16 and img_to_save.dtype != np.uint16:
            if img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0 * 65535.0).astype(np.uint16) # More accurate
            elif img_to_save.dtype in [np.float16, np.float32, np.float64]: img_to_save = (np.clip(img_to_save, 0.0, 1.0) * 65535.0).astype(np.uint16)
            else: img_to_save = img_to_save.astype(np.uint16)
        elif output_dtype_target == np.float16 and img_to_save.dtype != np.float16:
            if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0).astype(np.float16)
            elif img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0).astype(np.float16)
            elif img_to_save.dtype in [np.float32, np.float64]: img_to_save = img_to_save.astype(np.float16)
            # else: cannot convert to float16 easily
        elif output_dtype_target == np.float32 and img_to_save.dtype != np.float32:
             if img_to_save.dtype == np.uint16: img_to_save = (img_to_save.astype(np.float32) / 65535.0)
             elif img_to_save.dtype == np.uint8: img_to_save = (img_to_save.astype(np.float32) / 255.0)
             elif img_to_save.dtype == np.float16: img_to_save = img_to_save.astype(np.float32)


    # 2. Color Space Conversion (RGB -> BGR)
    # Typically, OpenCV expects BGR for formats like PNG, JPG. EXR usually expects RGB.
    # The `convert_to_bgr_before_save` flag controls this.
    # If output_format is exr, this should generally be False.
    current_format = output_format if output_format else path_obj.suffix.lower().lstrip('.')
    
    if convert_to_bgr_before_save and current_format != 'exr':
        if len(img_to_save.shape) == 3 and img_to_save.shape[2] == 3:
            img_to_save = convert_rgb_to_bgr(img_to_save)
        # BGRA is handled by OpenCV imwrite for PNGs, no explicit conversion needed if saving as RGBA.
        # If it's 4-channel and not PNG/TIFF with alpha, it might need stripping or specific handling.
        # For simplicity, this function assumes 3-channel RGB input if BGR conversion is active.

    # 3. Save Image
    try:
        if params:
            cv2.imwrite(str(path_obj), img_to_save, params)
        else:
            cv2.imwrite(str(path_obj), img_to_save)
        return True
    except Exception: # as e:
        # print(f"Error saving image {path_obj}: {e}") # Optional: for debugging utils
        return False