"""
Image decompression module for Sempress-compressed images.
"""

from __future__ import annotations
from typing import Tuple
import numpy as np
from .container import unpack_container
from .image_encoder import _ycbcr_to_rgb, _reconstruct_from_patches
from PIL import Image
import io


def _decode_channel(channel_data: dict, n_patches: int, patch_size: int) -> np.ndarray:
    """
    Decode a single channel from compressed data.

    Args:
        channel_data: Compressed channel information
        n_patches: Total number of patches
        patch_size: Size of each patch

    Returns:
        Decoded patches (n_patches, patch_size, patch_size, 1)
    """
    # Load codebook
    codebook_bytes = channel_data["codebook"]
    patch_dim = patch_size * patch_size
    n_centroids = len(codebook_bytes) // (4 * patch_dim)  # float32 = 4 bytes
    codebook = np.frombuffer(codebook_bytes, dtype=np.float32).reshape(n_centroids, patch_dim)

    # Load indices (backward compatible: check for index_dtype metadata)
    indices_bytes = channel_data["indices"]
    idx_dtype_str = channel_data.get("index_dtype", "uint16")
    idx_dtype = np.uint8 if idx_dtype_str == "uint8" else np.uint16
    indices = np.frombuffer(indices_bytes, dtype=idx_dtype)

    # Reconstruct patches from codebook
    recon_flat = codebook[indices]
    recon_patches = recon_flat.reshape(n_patches, patch_size, patch_size, 1)

    # Apply residuals if present
    residuals = channel_data.get("residuals", {})
    if residuals and "indices" in residuals:
        residual_indices = np.frombuffer(residuals["indices"], dtype=np.uint32)
        # Backward compatible: check residual dtype
        residual_dtype_str = residuals.get("residual_dtype", "float32")
        residual_dtype = np.float16 if residual_dtype_str == "float16" else np.float32
        residual_data = np.frombuffer(residuals["data"], dtype=residual_dtype).astype(np.float32)
        residual_data = residual_data.reshape(len(residual_indices), patch_dim)

        # Add residuals to reconstruction
        recon_flat_view = recon_patches.reshape(n_patches, patch_dim)
        recon_flat_view[residual_indices] += residual_data
        recon_patches = recon_flat_view.reshape(n_patches, patch_size, patch_size, 1)

    # Clip to valid range
    recon_patches = np.clip(recon_patches, 0, 255).astype(np.uint8)

    return recon_patches


def decode_image(compressed_blob: bytes) -> np.ndarray:
    """
    Decode a compressed image blob to numpy array.

    Args:
        compressed_blob: Compressed image data

    Returns:
        Reconstructed image as numpy array
    """
    payload = unpack_container(compressed_blob)

    assert payload["domain"] == "image", "Not an image payload"

    # Extract metadata
    original_shape = tuple(payload["original_shape"])
    patch_size = payload["patch_size"]
    grid_shape = tuple(payload["grid_shape"])
    color_space = payload["color_space"]
    channels = payload["channels"]
    channel_data = payload["channel_data"]

    # Calculate total patches
    n_patches = grid_shape[0] * grid_shape[1]

    # Decode each channel
    decoded_channels = []
    for ch in range(channels):
        ch_key = f"ch{ch}"
        ch_patches = _decode_channel(channel_data[ch_key], n_patches, patch_size)
        decoded_channels.append(ch_patches)

    # Combine channels
    if channels > 1:
        patches_combined = np.concatenate(decoded_channels, axis=3)
    else:
        patches_combined = decoded_channels[0].squeeze(axis=3)

    # Reconstruct image from patches
    img_array = _reconstruct_from_patches(patches_combined, grid_shape, original_shape)

    # Convert color space back if needed
    if color_space == "YCbCr" and img_array.ndim == 3:
        img_array = _ycbcr_to_rgb(img_array)

    return img_array


def decode_to_image_file(compressed_blob: bytes, output_path: str, format: str = None):
    """
    Decode compressed blob and save to image file.

    Args:
        compressed_blob: Compressed image data
        output_path: Output file path
        format: Image format (PNG, JPEG, etc.). Auto-detected from path if None.
    """
    img_array = decode_image(compressed_blob)

    # Convert to PIL Image
    if img_array.ndim == 2:
        img = Image.fromarray(img_array, mode='L')
    else:
        img = Image.fromarray(img_array, mode='RGB')

    # Save
    if format:
        img.save(output_path, format=format)
    else:
        img.save(output_path)


def decode_to_pil(compressed_blob: bytes) -> Image.Image:
    """
    Decode compressed blob to PIL Image.

    Args:
        compressed_blob: Compressed image data

    Returns:
        PIL Image object
    """
    img_array = decode_image(compressed_blob)

    if img_array.ndim == 2:
        return Image.fromarray(img_array, mode='L')
    else:
        return Image.fromarray(img_array, mode='RGB')
