import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
import io
from scipy.signal import wiener
from skimage.restoration import denoise_nl_means, estimate_sigma

def auto_white_balance(image_bytes: bytes) -> bytes:
    """Gray-world white balance"""
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    avg_a = np.average(result[:, :, 1])
    avg_b = np.average(result[:, :, 2])
    result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 1] / 255.0))
    result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 2] / 255.0))
    result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
    _, encoded = cv2.imencode('.png', result)
    return encoded.tobytes()

def adjust_exposure(image_bytes: bytes, exposure: float = 1.0) -> bytes:
    """exposure factor >1 brightens, <1 darkens"""
    img = Image.open(io.BytesIO(image_bytes))
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(exposure)
    out = io.BytesIO()
    img.save(out, format='PNG')
    return out.getvalue()

def recover_shadows_highlights(image_bytes: bytes, shadow_boost: float = 1.2, highlight_reduce: float = 0.8) -> bytes:
    """Simple gamma correction in shadows/highlights"""
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
    # Adaptive gamma: shadows ^0.7, highlights ^1.3
    gamma = np.ones_like(img)
    gamma[img < 0.3] = shadow_boost
    gamma[img > 0.7] = highlight_reduce
    img = np.power(img, gamma)
    img = (img * 255).astype(np.uint8)
    _, encoded = cv2.imencode('.png', img)
    return encoded.tobytes()

def advanced_denoise(image_bytes: bytes, patch_size: int = 5, h: float = 1.0) -> bytes:
    """Non-local means denoising with auto sigma estimation"""
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    # Estimate noise sigma
    sigma_est = estimate_sigma(img, channel_axis=-1, average_sigmas=True)
    denoised = denoise_nl_means(img, h=h * sigma_est, patch_size=patch_size, patch_distance=6, channel_axis=-1)
    denoised = (denoised * 255).astype(np.uint8)
    _, encoded = cv2.imencode('.png', denoised)
    return encoded.tobytes()

def deblur(image_bytes: bytes, kernel_size: int = 5, noise_var: float = 0.01) -> bytes:
    """Wiener deconvolution – assumes Gaussian blur"""
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
    # Create a simple Gaussian kernel
    kernel = cv2.getGaussianKernel(kernel_size, -1)
    kernel = kernel @ kernel.T
    # Apply Wiener filter to each channel
    from scipy.signal import convolve2d
    restored = np.zeros_like(img)
    for c in range(3):
        channel = img[:,:,c]
        # Wiener in frequency domain (simplified)
        from scipy.fft import fft2, ifft2, fftshift
        H = fft2(kernel, s=channel.shape)
        Y = fft2(channel)
        H_conj = np.conj(H)
        Wiener = H_conj / (np.abs(H)**2 + noise_var)
        restored_channel = np.real(ifft2(Y * Wiener))
        restored[:,:,c] = np.clip(restored_channel, 0, 1)
    restored = (restored * 255).astype(np.uint8)
    _, encoded = cv2.imencode('.png', restored)
    return encoded.tobytes()