# 📁 File: tools/image_enhancer/tools.py

import os
import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
from utils import validate_path, ensure_path_from_img, temp_output

# === TOOL IMPLEMENTATIONS ===

def apply_color_correction(input_path: str,
                           brightness: float = 1.0,
                           contrast: float = 1.0,
                           saturation: float = 1.0,
                           output_path: str = None):
    """
    Applies color correction to an image including brightness, contrast and saturation adjustments.

    Args:
        input_path (str): Path to the input image file
        output_path (str): Path where the corrected image will be saved
        brightness (float): Brightness factor (1.0 is original, <1.0 darker, >1.0 brighter)
        contrast (float): Contrast factor (1.0 is original, <1.0 less contrast, >1.0 more contrast)
        saturation (float): Saturation factor (1.0 is original, <1.0 less saturated, >1.0 more saturated)

    Returns:
        None: The function saves the output file to the specified path

    Raises:
        FileNotFoundError: If the input file does not exist
        ValueError: If the input is not a valid image
        OSError: If there's an error creating the output directory or writing the file
    """
    try:
        if output_path is None:
            output_path = temp_output()
        img = ensure_path_from_img(input_path)

        # Open the image
        img = Image.open(img)

        # Apply brightness adjustment
        if brightness != 1.0:
            enhancer = ImageEnhance.Brightness(img)
            img = enhancer.enhance(brightness)

        # Apply contrast adjustment
        if contrast != 1.0:
            enhancer = ImageEnhance.Contrast(img)
            img = enhancer.enhance(contrast)

        # Apply saturation adjustment
        if saturation != 1.0:
            enhancer = ImageEnhance.Color(img)
            img = enhancer.enhance(saturation)

        # Ensure the output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Save the result
        img.save(output_path)
    except Exception as e:
        # Re-raise with more context
        raise type(e)(f"Error in color correction: {str(e)}")
    return str(output_path)


def apply_sharpening(input_path: str, strength: float = 1.0, output_path: str = None)-> str:
    """
    Applies sharpening to an image to enhance edges and details.

    Args:
        input_path (str): Path to the input image file
        output_path (str): Path where the sharpened image will be saved
        strength (float): Sharpening strength (1.0 is standard, >1.0 is stronger)

    Returns:
        output_path (str): Path where the sharpened image will be saved

    Raises:
        FileNotFoundError: If the input file does not exist
        ValueError: If the input is not a valid image
        OSError: If there's an error creating the output directory or writing the file
    """
    try:
        if output_path is None:
            output_path = temp_output()
        # Validate input file exists
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")
        img = ensure_path_from_img(input_path)

        # Open the image
        img = Image.open(img)

        # Apply sharpening
        if strength <= 1.0:
            # Use PIL's built-in SHARPEN filter for moderate sharpening
            sharpened = img.filter(ImageFilter.SHARPEN)
        else:
            # For stronger sharpening, use UnsharpMask with adjusted parameters
            radius = 2.0
            percent = int(150 * strength)
            threshold = 3
            sharpened = img.filter(ImageFilter.UnsharpMask(radius, percent, threshold))

        # Ensure the output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Save the result
        sharpened.save(output_path)
    except Exception as e:
        # Re-raise with more context
        raise type(e)(f"Error in sharpening: {str(e)}")
    return str(output_path)

# --- Constants for Denoising Parameters ---
# These values control the filtering strength.
# h: Luminance component filter strength. Larger h means more noise removal.
# hColor: Color component filter strength.
# templateWindowSize: Should be odd. Recommended: 7.
# searchWindowSize: Should be odd. Recommended: 21.
DENOISE_PARAMS = {
    "low": {
        "h": 5,
        "hColor": 5,
        "templateWindowSize": 7,
        "searchWindowSize": 21,
    },
    "medium": {
        "h": 10,
        "hColor": 10,
        "templateWindowSize": 7,
        "searchWindowSize": 21,
    },
    "high": {
        "h": 15,
        "hColor": 15,
        "templateWindowSize": 7,
        "searchWindowSize": 21,
    }
}

def reduce_noise(input_path: str, strength: str = "medium", output_path: str = None) -> str:
    """
    Reduces noise in an image using cv2.fastNlMeansDenoisingColored.

    Args:
        input_path (str): Path to the input image file.
        output_path (str): Path where the denoised image will be saved.
        strength (str): Noise reduction strength ("low", "medium", or "high").
    Returns:
    output_path (str): Path where the sharpened image will be saved

    Raises:
        FileNotFoundError: If the input file does not exist.
        ValueError: If the strength parameter is invalid or the image cannot be loaded.
        OSError: If there's an error creating the output directory or writing the file.
    """
    if output_path is None:
        output_path = temp_output()
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input file not found: {input_path}")

    if strength not in DENOISE_PARAMS:
        raise ValueError(f"Invalid strength value: {strength}. Must be one of {list(DENOISE_PARAMS.keys())}")

    try:
        img = cv2.imread(input_path)

        if img is None:
            # This can happen if the file exists but is corrupted or not a valid image format.
            raise ValueError(f"Failed to load image from path: {input_path}. The file may be corrupted or in an unsupported format.")

        # Get parameters for the chosen strength
        params = DENOISE_PARAMS[strength]

        # Apply noise reduction
        denoised = cv2.fastNlMeansDenoisingColored(
            img,
            None,
            params["h"],
            params["hColor"],
            params["templateWindowSize"],
            params["searchWindowSize"]
        )

        # Ensure the output directory exists
        output_dir = os.path.dirname(output_path)
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)

        # Save the result
        success = cv2.imwrite(output_path, denoised)
        if not success:
            raise OSError(f"Failed to save the denoised image to {output_path}")

    except (cv2.error, OSError) as e:
        # Catch specific errors from OpenCV or file system and re-raise
        # them to provide context without losing the original error.
        raise OSError(f"An error occurred during image processing or file saving for {input_path}") from e
    return str(output_path)

def upscale_resolution(input_path: str, scale_factor: int = 2, method: str = "bicubic", output_path: str = None) -> str:
    """
    Increases the resolution of an image using various upscaling methods.

    Args:
        input_path (str): Path to the input image file
        scale_factor (int): How much to increase resolution (2 = 2x, 3 = 3x, 4 = 4x)
        method (str): Upscaling method ("nearest", "bilinear", "bicubic", or "lanczos")
        output_path (str): Path where the upscaled image will be saved

    Returns:
        output_path (str): Path where the sharpened image will be saved

    Raises:
        FileNotFoundError: If the input file does not exist
        ValueError: If the scale_factor or method parameters are invalid
        OSError: If there's an error creating the output directory or writing the file
    """
    try:
        if output_path is None:
            output_path = temp_output()

        # Validate input file exists
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")

        # Validate scale factor
        if scale_factor not in [2, 3, 4]:
            raise ValueError(f"Invalid scale factor: {scale_factor}. Must be 2, 3, or 4")

        # Map method names to PIL resampling filters
        method_map = {
            "nearest": Image.NEAREST,
            "bilinear": Image.BILINEAR,
            "bicubic": Image.BICUBIC,
            "lanczos": Image.LANCZOS
        }

        # Validate method name
        if method.lower() not in method_map:
            raise ValueError(f"Invalid method: {method}. Must be one of {list(method_map.keys())}")

        # Use bicubic as default if method is not recognized
        resampling_filter = method_map.get(method.lower(), Image.BICUBIC)

        # Open the image
        img = Image.open(input_path)

        # Get original dimensions
        width, height = img.size

        # Calculate new dimensions
        new_width = width * scale_factor
        new_height = height * scale_factor

        # Resize the image
        upscaled = img.resize((new_width, new_height), resampling_filter)

        # Ensure the output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Save the result
        upscaled.save(output_path)
    except Exception as e:
        # Re-raise with more context
        raise type(e)(f"Error in upscaling: {str(e)}")
    return str(output_path)

def auto_enhance(input_path: str, output_path: str = None) -> str:
    """
    Automatically enhances an image by applying balanced corrections to
    brightness, contrast, color, and sharpness.

    Args:
        input_path (str): Path to the input image file
        output_path (str): Path where the enhanced image will be saved

    Returns:
        output_path (str): Path where the sharpened image will be saved

    Raises:
        FileNotFoundError: If the input file does not exist
        ValueError: If the input is not a valid image
        OSError: If there's an error creating the output directory or writing the file
    """
    try:
        if output_path is None:
            output_path = temp_output()
        # Validate input file exists
        if not os.path.exists(input_path):
            raise FileNotFoundError(f"Input file not found: {input_path}")

        # Open the image
        img = Image.open(input_path)

        # Auto-enhance: apply a series of balanced enhancements
        # Step 1: Auto-contrast
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(1.2)

        # Step 2: Moderate brightness adjustment
        enhancer = ImageEnhance.Brightness(img)
        img = enhancer.enhance(1.1)

        # Step 3: Moderate color enhancement
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(1.1)

        # Step 4: Moderate sharpening
        enhancer = ImageEnhance.Sharpness(img)
        img = enhancer.enhance(1.3)

        # Ensure the output directory exists
        os.makedirs(os.path.dirname(output_path), exist_ok=True)

        # Save the result
        img.save(output_path)
    except Exception as e:
        # Re-raise with more context
        raise type(e)(f"Error in auto enhancement: {str(e)}")
    return str(output_path)