"""
Stegasoo Steganography Functions (v3.2.0)

This is the core embedding/extraction module. Two modes available:

LSB (Least Significant Bit) Mode:
- Classic steganography technique - hide bits in the least significant bit of pixel values
- Works on any image, outputs lossless PNG/BMP
- Higher capacity than DCT, but destroyed by JPEG compression
- Great for: high-capacity needs, lossless workflows

DCT Mode (see dct_steganography.py):
- Hides data in frequency-domain coefficients
- Survives some image processing, works with JPEG
- Lower capacity but more robust
- Great for: JPEG images, robustness needs

Both modes use pseudo-random pixel/coefficient selection based on a key.
Without the key, you don't know where to look - security through obscurity
PLUS actual encryption of the payload.

v3.0: Added DCT mode with scipy
v3.0.1: DCT output format options (PNG/JPEG, grayscale/color)
v3.2.0: Fixed overhead calculations after removing date field
"""

import io
import struct
from typing import TYPE_CHECKING, Union

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms
from PIL import Image

if TYPE_CHECKING:
    from .dct_steganography import DCTEmbedStats

from .constants import (
    EMBED_MODE_AUTO,
    EMBED_MODE_DCT,
    EMBED_MODE_LSB,
    VALID_EMBED_MODES,
)
from .debug import debug
from .exceptions import CapacityError, EmbeddingError
from .models import EmbedStats, FilePayload

# Progress reporting interval
PROGRESS_INTERVAL = 1000  # Write every N pixels for LSB


def _write_progress(progress_file: str | None, current: int, total: int, phase: str = "embedding"):
    """Write progress to file for frontend polling."""
    if progress_file is None:
        return
    try:
        import json

        with open(progress_file, "w") as f:
            json.dump(
                {
                    "current": current,
                    "total": total,
                    "percent": round((current / total) * 100, 1) if total > 0 else 0,
                    "phase": phase,
                },
                f,
            )
    except Exception:
        pass  # Don't let progress writing break encoding


# Lossless formats that preserve LSB data
LOSSLESS_FORMATS = {"PNG", "BMP", "TIFF"}

# Format to extension mapping
FORMAT_TO_EXT = {
    "PNG": "png",
    "BMP": "bmp",
    "TIFF": "tiff",
}

# Extension to PIL format mapping
EXT_TO_FORMAT = {
    "png": "PNG",
    "bmp": "BMP",
    "tiff": "TIFF",
    "tif": "TIFF",
}

# =============================================================================
# OVERHEAD CONSTANTS
# =============================================================================
#
# Every stego image has some overhead before the actual payload:
#
# The encrypted message format (v4.0.0):
# ┌─────────────────────────────────────────────────────────────────┐
# │ \x89ST3 │ v5 │ flags │  salt (32)  │  iv (12)  │  tag (16)  │ ... │
# │ magic  │ ver│       │             │           │            │ data│
# └─────────────────────────────────────────────────────────────────┘
#   4 bytes  1    1         32            12           16         var
#
# Plus LSB embedding adds a 4-byte length prefix so we know where to stop.
#
# History of overhead sizes (in case you're debugging old images):
# - v3.1.0: 76 bytes (had date field - 10+1 bytes)
# - v3.2.0: 65 bytes (removed date, simpler)
# - v4.0.0: 66 bytes (added flags byte for channel key)
# - v4.4.0: 82 bytes (added 16-byte message nonce for HKDF)

HEADER_OVERHEAD = 82  # What the crypto layer adds to any message (v6 format)
LENGTH_PREFIX = 4  # We prepend the payload length for LSB extraction
ENCRYPTION_OVERHEAD = HEADER_OVERHEAD + LENGTH_PREFIX  # Total: 86 bytes

# That 86 bytes is your minimum image capacity requirement.
# A tiny 100x100 image gives you ~3750 bytes capacity, minus 86 = ~3664 usable.

# DCT output format options (v3.0.1)
DCT_OUTPUT_PNG = "png"
DCT_OUTPUT_JPEG = "jpeg"

# DCT color mode options (v3.0.1)
DCT_COLOR_GRAYSCALE = "grayscale"
DCT_COLOR_COLOR = "color"


# =============================================================================
# DCT MODULE LAZY LOADING
# =============================================================================

_dct_module = None


def _get_dct_module():
    """Lazy load DCT module to avoid scipy import if not needed."""
    global _dct_module
    if _dct_module is None:
        from . import dct_steganography

        _dct_module = dct_steganography
    return _dct_module


def has_dct_support() -> bool:
    """
    Check if DCT steganography mode is available.

    Returns:
        True if scipy is installed and DCT functions work

    Example:
        >>> if has_dct_support():
        ...     result = encode(..., embed_mode='dct')
    """
    try:
        dct_mod = _get_dct_module()
        return dct_mod.has_dct_support()
    except (ImportError, ValueError):
        # ValueError: numpy binary incompatibility (e.g., jpeglib built against numpy 2.x)
        return False


# =============================================================================
# FORMAT UTILITIES
# =============================================================================


def get_output_format(input_format: str | None) -> tuple[str, str]:
    """
    Determine the output format based on input format.

    Args:
        input_format: PIL format string of input image (e.g., 'JPEG', 'PNG')

    Returns:
        Tuple of (PIL format string, file extension) for output
        Falls back to PNG for lossy or unknown formats.
    """
    debug.validate(
        input_format is None or isinstance(input_format, str), "Input format must be string or None"
    )

    if input_format and input_format.upper() in LOSSLESS_FORMATS:
        fmt = input_format.upper()
        ext = FORMAT_TO_EXT.get(fmt, "png")
        debug.print(f"Using lossless format: {fmt} -> .{ext}")
        return fmt, ext

    debug.print(f"Input format {input_format} is lossy or unknown, defaulting to PNG")
    return "PNG", "png"


# =============================================================================
# CAPACITY FUNCTIONS
# =============================================================================


def will_fit(
    payload: str | bytes | FilePayload | int,
    carrier_image: bytes,
    bits_per_channel: int = 1,
    include_compression_estimate: bool = True,
) -> dict:
    """
    Check if a payload will fit in a carrier image (LSB mode).

    Args:
        payload: Message string, raw bytes, FilePayload, or size in bytes
        carrier_image: Carrier image bytes
        bits_per_channel: Bits to use per color channel (1-2)
        include_compression_estimate: Estimate compressed size

    Returns:
        Dict with fits, capacity, usage info
    """
    # Determine payload size
    if isinstance(payload, int):
        payload_size = payload
        payload_data = None
    elif isinstance(payload, str):
        payload_data = payload.encode("utf-8")
        payload_size = len(payload_data)
    elif isinstance(payload, FilePayload):
        payload_data = payload.data
        filename_overhead = len(payload.filename.encode("utf-8")) if payload.filename else 0
        mime_overhead = len(payload.mime_type.encode("utf-8")) if payload.mime_type else 0
        payload_size = len(payload.data) + filename_overhead + mime_overhead + 5
    else:
        payload_data = payload
        payload_size = len(payload)

    capacity = calculate_capacity(carrier_image, bits_per_channel)

    # Estimate encrypted size with padding
    # Padding adds 64-319 bytes, rounded up to 256-byte boundary
    # Average case: ~190 bytes padding
    estimated_padding = 190
    estimated_encrypted_size = payload_size + estimated_padding + ENCRYPTION_OVERHEAD

    compressed_estimate = None
    if include_compression_estimate and payload_data is not None and len(payload_data) >= 64:
        try:
            import zlib

            compressed = zlib.compress(payload_data, level=6)
            compressed_size = len(compressed) + 9  # Compression header
            if compressed_size < payload_size:
                compressed_estimate = compressed_size
                estimated_encrypted_size = compressed_size + estimated_padding + ENCRYPTION_OVERHEAD
        except Exception:
            pass

    headroom = capacity - estimated_encrypted_size
    fits = headroom >= 0
    usage_percent = (estimated_encrypted_size / capacity * 100) if capacity > 0 else 100.0

    return {
        "fits": fits,
        "payload_size": payload_size,
        "estimated_encrypted_size": estimated_encrypted_size,
        "capacity": capacity,
        "usage_percent": min(usage_percent, 100.0),
        "headroom": headroom,
        "compressed_estimate": compressed_estimate,
        "mode": EMBED_MODE_LSB,
    }


def calculate_capacity(image_data: bytes, bits_per_channel: int = 1) -> int:
    """
    Calculate the maximum message capacity of an image (LSB mode).

    Args:
        image_data: Image bytes
        bits_per_channel: Bits to use per color channel

    Returns:
        Maximum bytes that can be embedded (minus overhead)
    """
    debug.validate(
        bits_per_channel in (1, 2), f"bits_per_channel must be 1 or 2, got {bits_per_channel}"
    )

    img_file = Image.open(io.BytesIO(image_data))
    try:
        num_pixels = img_file.size[0] * img_file.size[1]
        bits_per_pixel = 3 * bits_per_channel
        max_bytes = (num_pixels * bits_per_pixel) // 8

        capacity = max(0, max_bytes - ENCRYPTION_OVERHEAD)
        debug.print(f"LSB capacity: {capacity} bytes at {bits_per_channel} bit(s)/channel")
        return capacity
    finally:
        img_file.close()


def calculate_capacity_by_mode(
    image_data: bytes,
    embed_mode: str = EMBED_MODE_LSB,
    bits_per_channel: int = 1,
) -> dict:
    """
    Calculate capacity for specified embedding mode.

    Args:
        image_data: Carrier image bytes
        embed_mode: 'lsb' or 'dct'
        bits_per_channel: Bits per channel for LSB mode

    Returns:
        Dict with capacity information
    """
    if embed_mode == EMBED_MODE_DCT:
        if not has_dct_support():
            raise ImportError("scipy required for DCT mode. Install: pip install scipy")

        dct_mod = _get_dct_module()
        dct_info = dct_mod.calculate_dct_capacity(image_data)

        return {
            "mode": EMBED_MODE_DCT,
            "capacity_bytes": dct_info.usable_capacity_bytes,
            "capacity_bits": dct_info.total_capacity_bits,
            "width": dct_info.width,
            "height": dct_info.height,
            "total_blocks": dct_info.total_blocks,
        }
    else:
        capacity = calculate_capacity(image_data, bits_per_channel)
        img = Image.open(io.BytesIO(image_data))
        try:
            width, height = img.size
        finally:
            img.close()

        return {
            "mode": EMBED_MODE_LSB,
            "capacity_bytes": capacity,
            "capacity_bits": capacity * 8,
            "width": width,
            "height": height,
            "bits_per_channel": bits_per_channel,
        }


def will_fit_by_mode(
    payload: str | bytes | FilePayload | int,
    carrier_image: bytes,
    embed_mode: str = EMBED_MODE_LSB,
    bits_per_channel: int = 1,
) -> dict:
    """
    Check if payload fits in specified mode.

    Args:
        payload: Message, bytes, FilePayload, or size in bytes
        carrier_image: Carrier image bytes
        embed_mode: 'lsb' or 'dct'
        bits_per_channel: For LSB mode

    Returns:
        Dict with fits, capacity, usage info
    """
    if embed_mode == EMBED_MODE_DCT:
        if not has_dct_support():
            return {"fits": False, "error": "scipy not available", "mode": EMBED_MODE_DCT}

        if isinstance(payload, int):
            payload_size = payload
        elif isinstance(payload, str):
            payload_size = len(payload.encode("utf-8"))
        elif hasattr(payload, "data"):
            payload_size = len(payload.data)
        else:
            payload_size = len(payload)

        estimated_size = payload_size + ENCRYPTION_OVERHEAD + 190  # padding estimate

        dct_mod = _get_dct_module()
        fits = dct_mod.will_fit_dct(estimated_size, carrier_image)
        capacity_info = dct_mod.calculate_dct_capacity(carrier_image)
        capacity = capacity_info.usable_capacity_bytes

        usage_percent = (estimated_size / capacity * 100) if capacity > 0 else 100.0

        return {
            "fits": fits,
            "payload_size": payload_size,
            "capacity": capacity,
            "usage_percent": min(usage_percent, 100.0),
            "headroom": capacity - estimated_size,
            "mode": EMBED_MODE_DCT,
        }
    else:
        return will_fit(payload, carrier_image, bits_per_channel)


def get_available_modes() -> dict:
    """
    Get available embedding modes and their status.

    Returns:
        Dict mapping mode name to availability info
    """
    return {
        EMBED_MODE_LSB: {
            "available": True,
            "name": "Spatial LSB",
            "description": "Embed in pixel LSBs, outputs PNG/BMP",
            "output_format": "PNG (color)",
        },
        EMBED_MODE_DCT: {
            "available": has_dct_support(),
            "name": "DCT Domain",
            "description": "Embed in DCT coefficients, outputs grayscale PNG or JPEG",
            "output_formats": ["PNG (grayscale)", "JPEG (grayscale)"],
            "requires": "scipy",
        },
    }


def compare_modes(image_data: bytes) -> dict:
    """
    Compare embedding modes for a carrier image.

    Args:
        image_data: Carrier image bytes

    Returns:
        Dict with comparison of LSB vs DCT modes
    """
    img = Image.open(io.BytesIO(image_data))
    try:
        width, height = img.size
    finally:
        img.close()

    lsb_bytes = calculate_capacity(image_data, 1)

    if has_dct_support():
        dct_mod = _get_dct_module()
        dct_info = dct_mod.calculate_dct_capacity(image_data)
        dct_bytes = dct_info.usable_capacity_bytes
        dct_available = True
    else:
        safe_blocks = (height // 8) * (width // 8)
        dct_bytes = (safe_blocks * 16) // 8  # Estimated
        dct_available = False

    return {
        "width": width,
        "height": height,
        "lsb": {
            "capacity_bytes": lsb_bytes,
            "capacity_kb": lsb_bytes / 1024,
            "available": True,
            "output": "PNG (color)",
        },
        "dct": {
            "capacity_bytes": dct_bytes,
            "capacity_kb": dct_bytes / 1024,
            "available": dct_available,
            "output": "PNG or JPEG (grayscale)",
            "ratio_vs_lsb": (dct_bytes / lsb_bytes * 100) if lsb_bytes > 0 else 0,
        },
    }


# =============================================================================
# PIXEL INDEX GENERATION
# =============================================================================
#
# The key insight: we don't hide data in sequential pixels (that's easy to find).
# Instead, we scatter the data across pseudo-random pixel locations.
#
# The pixel selection key (derived from passphrase + photo + pin) determines
# WHICH pixels get modified. Without the key, an attacker would have to:
# 1. Know we're using LSB steganography
# 2. Try every possible subset of pixels
# 3. Decrypt the result (which they also can't do without the key)
#
# We use ChaCha20 as a CSPRNG (Cryptographically Secure PRNG). It's:
# - Fast (faster than AES-CTR on most CPUs)
# - Deterministic (same key = same sequence, needed for extraction)
# - Secure (can't predict the sequence without the key)


@debug.time
def generate_pixel_indices(key: bytes, num_pixels: int, num_needed: int) -> list[int]:
    """
    Generate pseudo-random pixel indices for embedding.

    This is the "where do we hide the bits?" function. We use ChaCha20
    to generate a deterministic sequence of pixel indices that only
    someone with the same key can reproduce.

    Two strategies based on how much of the image we're using:
    - >= 50% capacity: Full Fisher-Yates shuffle (sample without replacement)
    - < 50% capacity: Direct random sampling (faster, same result)
    """
    debug.validate(len(key) == 32, f"Pixel key must be 32 bytes, got {len(key)}")
    debug.validate(num_pixels > 0, f"Number of pixels must be positive, got {num_pixels}")
    debug.validate(num_needed > 0, f"Number needed must be positive, got {num_needed}")
    debug.validate(
        num_needed <= num_pixels, f"Cannot select {num_needed} pixels from {num_pixels} available"
    )

    debug.print(f"Generating {num_needed} pixel indices from {num_pixels} total pixels")

    # Strategy 1: Full shuffle when we need a lot of pixels
    # Fisher-Yates shuffle is O(n) and gives us perfect random sampling
    if num_needed >= num_pixels // 2:
        debug.print(f"Using full shuffle (needed {num_needed}/{num_pixels} pixels)")
        nonce = b"\x00" * 16
        cipher = Cipher(algorithms.ChaCha20(key, nonce), mode=None, backend=default_backend())
        encryptor = cipher.encryptor()

        indices = list(range(num_pixels))
        # Get enough random bytes to do the shuffle
        random_bytes = encryptor.update(b"\x00" * (num_pixels * 4))

        # Fisher-Yates shuffle - swap each element with a random earlier element
        for i in range(num_pixels - 1, 0, -1):
            j_bytes = random_bytes[(num_pixels - 1 - i) * 4 : (num_pixels - i) * 4]
            j = int.from_bytes(j_bytes, "big") % (i + 1)
            indices[i], indices[j] = indices[j], indices[i]

        selected = indices[:num_needed]
        debug.print(f"Generated {len(selected)} indices via shuffle")
        return selected

    # Strategy 2: Direct sampling when we need fewer pixels
    # Generate random indices until we have enough unique ones
    debug.print(f"Using optimized selection (needed {num_needed}/{num_pixels} pixels)")
    selected = []
    used = set()  # Track which pixels we've already picked

    nonce = b"\x00" * 16
    cipher = Cipher(algorithms.ChaCha20(key, nonce), mode=None, backend=default_backend())
    encryptor = cipher.encryptor()

    # Pre-generate 2x the bytes we think we'll need (for collision handling)
    bytes_needed = (num_needed * 2) * 4
    random_bytes = encryptor.update(b"\x00" * bytes_needed)

    byte_offset = 0
    collisions = 0
    while len(selected) < num_needed and byte_offset < len(random_bytes) - 4:
        idx = int.from_bytes(random_bytes[byte_offset : byte_offset + 4], "big") % num_pixels
        byte_offset += 4

        if idx not in used:
            used.add(idx)
            selected.append(idx)
        else:
            collisions += 1  # Birthday paradox in action

    # Edge case: ran out of pre-generated bytes (very high collision rate)
    if len(selected) < num_needed:
        debug.print(f"Need {num_needed - len(selected)} more indices, generating...")
        extra_needed = num_needed - len(selected)
        for _ in range(extra_needed * 2):
            extra_bytes = encryptor.update(b"\x00" * 4)
            idx = int.from_bytes(extra_bytes, "big") % num_pixels
            if idx not in used:
                used.add(idx)
                selected.append(idx)
                if len(selected) == num_needed:
                    break

    debug.print(f"Generated {len(selected)} indices with {collisions} collisions")
    debug.validate(
        len(selected) == num_needed,
        f"Failed to generate enough indices: {len(selected)}/{num_needed}",
    )
    return selected


# =============================================================================
# EMBEDDING FUNCTIONS
# =============================================================================
#
# The actual bit-hiding magic happens here. LSB embedding is conceptually simple:
#
# Original pixel RGB: (142, 87, 201)
# In binary:          (10001110, 01010111, 11001001)
#                                      ^       ^       ^
#                      These are the LSBs (least significant bits)
#
# To hide the bits [1, 0, 1]:
# Modified pixel RGB: (10001111, 01010110, 11001001) = (143, 86, 201)
#                                      ^       ^       ^
#                      Changed!     Changed!  Already 1, no change needed
#
# The human eye can't see the difference between 142 and 143.
# But we've hidden 3 bits of secret data in one pixel.
#
# With a 1000x1000 image: 1 million pixels * 3 channels = 3 million bits = 375 KB!


@debug.time
def embed_in_image(
    data: bytes,
    image_data: bytes,
    pixel_key: bytes,
    bits_per_channel: int = 1,
    output_format: str | None = None,
    embed_mode: str = EMBED_MODE_LSB,
    dct_output_format: str = DCT_OUTPUT_PNG,
    dct_color_mode: str = "color",
    progress_file: str | None = None,
    quant_step: int | None = None,
    jpeg_quality: int | None = None,
    max_dimension: int | None = None,
) -> tuple[bytes, Union[EmbedStats, "DCTEmbedStats"], str]:
    """
    Embed data into an image using specified mode.

    Args:
        data: Data to embed (encrypted payload)
        image_data: Carrier image bytes
        pixel_key: Key for pixel/coefficient selection
        bits_per_channel: Bits per channel (LSB mode only)
        output_format: Force output format (LSB mode only)
        embed_mode: 'lsb' (default) or 'dct'
        dct_output_format: For DCT mode - 'png' (lossless) or 'jpeg' (smaller)
        dct_color_mode: For DCT mode - 'grayscale' (default) or 'color' (preserves colors)

    Returns:
        Tuple of (stego image bytes, stats, file extension)

    Raises:
        CapacityError: If data won't fit
        EmbeddingError: If embedding fails
        ImportError: If DCT mode requested but scipy unavailable
    """
    debug.print(f"embed_in_image: mode={embed_mode}, data={len(data)} bytes")
    debug.validate(
        embed_mode in VALID_EMBED_MODES, f"Invalid embed_mode: {embed_mode}. Use 'lsb' or 'dct'"
    )

    # Dispatch via backend registry
    from .backends import registry

    backend = registry.get(embed_mode)
    if not backend.is_available():
        raise ImportError(
            f"Dependencies for '{embed_mode}' mode are not installed. "
            f"Install with: pip install stegasoo[dct]"
        )

    if embed_mode == EMBED_MODE_DCT:
        # Validate DCT-specific options
        if dct_output_format not in (DCT_OUTPUT_PNG, DCT_OUTPUT_JPEG):
            debug.print(f"Invalid dct_output_format '{dct_output_format}', defaulting to PNG")
            dct_output_format = DCT_OUTPUT_PNG
        if dct_color_mode not in ("grayscale", "color"):
            debug.print(f"Invalid dct_color_mode '{dct_color_mode}', defaulting to color")
            dct_color_mode = "color"

        stego_bytes, dct_stats = backend.embed(
            data,
            image_data,
            pixel_key,
            progress_file=progress_file,
            dct_output_format=dct_output_format,
            dct_color_mode=dct_color_mode,
            quant_step=quant_step,
            jpeg_quality=jpeg_quality,
            max_dimension=max_dimension,
        )
        ext = "jpg" if dct_output_format == DCT_OUTPUT_JPEG else "png"
        debug.print(
            f"DCT embedding complete: {dct_output_format.upper()} output, "
            f"color_mode={dct_color_mode}, ext={ext}"
        )
        return stego_bytes, dct_stats, ext

    # LSB and other image backends
    stego_bytes, stats = backend.embed(
        data,
        image_data,
        pixel_key,
        progress_file=progress_file,
        bits_per_channel=bits_per_channel,
        output_format=output_format,
    )
    ext = getattr(stats, "output_extension", "png")
    return stego_bytes, stats, ext


def _embed_lsb(
    data: bytes,
    image_data: bytes,
    pixel_key: bytes,
    bits_per_channel: int = 1,
    output_format: str | None = None,
    progress_file: str | None = None,
) -> tuple[bytes, EmbedStats, str]:
    """
    Embed data using LSB steganography (internal implementation).
    """
    debug.print(f"LSB embedding {len(data)} bytes into image")
    debug.data(pixel_key, "Pixel key for embedding")
    debug.validate(
        bits_per_channel in (1, 2), f"bits_per_channel must be 1 or 2, got {bits_per_channel}"
    )
    debug.validate(len(pixel_key) == 32, f"Pixel key must be 32 bytes, got {len(pixel_key)}")

    img_file = None
    img = None
    stego_img = None

    try:
        img_file = Image.open(io.BytesIO(image_data))
        input_format = img_file.format

        debug.print(f"Carrier image: {img_file.size[0]}x{img_file.size[1]}, format: {input_format}")

        img = img_file.convert("RGB") if img_file.mode != "RGB" else img_file.copy()
        if img_file.mode != "RGB":
            debug.print(f"Converting image from {img_file.mode} to RGB")

        pixels = list(img.getdata())
        num_pixels = len(pixels)

        bits_per_pixel = 3 * bits_per_channel
        max_bytes = (num_pixels * bits_per_pixel) // 8

        debug.print(f"Image capacity: {max_bytes} bytes at {bits_per_channel} bit(s)/channel")

        data_with_len = struct.pack(">I", len(data)) + data

        if len(data_with_len) > max_bytes:
            debug.print(f"Capacity error: need {len(data_with_len)}, have {max_bytes}")
            raise CapacityError(len(data_with_len), max_bytes)

        debug.print(
            f"Total data to embed: {len(data_with_len)} bytes "
            f"({len(data_with_len)/max_bytes*100:.1f}% of capacity)"
        )

        binary_data = "".join(format(b, "08b") for b in data_with_len)
        pixels_needed = (len(binary_data) + bits_per_pixel - 1) // bits_per_pixel

        debug.print(f"Need {pixels_needed} pixels to embed {len(binary_data)} bits")

        selected_indices = generate_pixel_indices(pixel_key, num_pixels, pixels_needed)

        new_pixels = list(pixels)
        clear_mask = 0xFF ^ ((1 << bits_per_channel) - 1)

        bit_idx = 0
        modified_pixels = 0
        total_pixels_to_process = len(selected_indices)

        # Initial progress write - signals prep is done, embedding starting
        if progress_file:
            _write_progress(progress_file, 5, 100, "embedding")

        for progress_idx, pixel_idx in enumerate(selected_indices):
            if bit_idx >= len(binary_data):
                break

            r, g, b = new_pixels[pixel_idx]
            modified = False

            for channel_idx, channel_val in enumerate([r, g, b]):
                if bit_idx >= len(binary_data):
                    break
                bits = binary_data[bit_idx : bit_idx + bits_per_channel].ljust(
                    bits_per_channel, "0"
                )
                new_val = (channel_val & clear_mask) | int(bits, 2)

                if channel_val != new_val:
                    modified = True
                    if channel_idx == 0:
                        r = new_val
                    elif channel_idx == 1:
                        g = new_val
                    else:
                        b = new_val

                bit_idx += bits_per_channel

            if modified:
                new_pixels[pixel_idx] = (r, g, b)
                modified_pixels += 1

            # Report progress periodically
            if progress_file and progress_idx % PROGRESS_INTERVAL == 0:
                _write_progress(progress_file, progress_idx, total_pixels_to_process, "embedding")

        # Final progress before save
        if progress_file:
            _write_progress(
                progress_file, total_pixels_to_process, total_pixels_to_process, "saving"
            )

        debug.print(f"Modified {modified_pixels} pixels (out of {len(selected_indices)} selected)")

        stego_img = Image.new("RGB", img.size)
        stego_img.putdata(new_pixels)

        if output_format:
            out_fmt = output_format.upper()
            out_ext = FORMAT_TO_EXT.get(out_fmt, "png")
            debug.print(f"Using forced output format: {out_fmt}")
        else:
            out_fmt, out_ext = get_output_format(input_format)
            debug.print(f"Auto-selected output format: {out_fmt}")

        output = io.BytesIO()
        stego_img.save(output, out_fmt)
        output.seek(0)

        stats = EmbedStats(
            pixels_modified=modified_pixels,
            total_pixels=num_pixels,
            capacity_used=len(data_with_len) / max_bytes,
            bytes_embedded=len(data_with_len),
        )

        debug.print(f"LSB embedding complete: {out_fmt} image, {len(output.getvalue())} bytes")
        return output.getvalue(), stats, out_ext

    except CapacityError:
        raise
    except Exception as e:
        debug.exception(e, "embed_lsb")
        raise EmbeddingError(f"Failed to embed data: {e}") from e
    finally:
        # Properly close all PIL Images to prevent memory leaks
        if stego_img is not None:
            stego_img.close()
        if img is not None and img is not img_file:
            img.close()
        if img_file is not None:
            img_file.close()


# =============================================================================
# EXTRACTION FUNCTIONS
# =============================================================================


@debug.time
def extract_from_image(
    image_data: bytes,
    pixel_key: bytes,
    bits_per_channel: int = 1,
    embed_mode: str = EMBED_MODE_AUTO,
    progress_file: str | None = None,
    quant_step: int | None = None,
) -> bytes | None:
    """
    Extract hidden data from a stego image.

    Args:
        image_data: Stego image bytes
        pixel_key: Key for pixel/coefficient selection (must match encoding)
        bits_per_channel: Bits per channel (LSB mode only)
        embed_mode: 'auto' (try both), 'lsb', or 'dct'
        progress_file: Optional path to write progress JSON for UI polling

    Returns:
        Extracted data bytes, or None if extraction fails
    """
    debug.print(f"extract_from_image: mode={embed_mode}")

    from .backends import registry

    # AUTO MODE: Try LSB first (cheaper), then other backends
    if embed_mode == EMBED_MODE_AUTO:
        auto_order = [EMBED_MODE_LSB] + [
            m for m in registry.available_modes(carrier_type="image") if m != EMBED_MODE_LSB
        ]
        for mode in auto_order:
            backend = registry.get(mode)
            debug.print(f"Auto-detect: trying {mode}")
            result = backend.extract(
                image_data,
                pixel_key,
                progress_file=progress_file,
                bits_per_channel=bits_per_channel,
                quant_step=quant_step,
            )
            if result is not None:
                debug.print(f"Auto-detect: {mode} extraction succeeded")
                return result
        debug.print("Auto-detect: All modes failed")
        return None

    # EXPLICIT MODE
    backend = registry.get(embed_mode)
    if not backend.is_available():
        raise ImportError(f"Dependencies for '{embed_mode}' mode are not installed.")
    return backend.extract(
        image_data,
        pixel_key,
        progress_file=progress_file,
        bits_per_channel=bits_per_channel,
        quant_step=quant_step,
    )


def _extract_dct(
    image_data: bytes,
    pixel_key: bytes,
    progress_file: str | None = None,
) -> bytes | None:
    """Extract using DCT mode."""
    try:
        dct_mod = _get_dct_module()
        return dct_mod.extract_from_dct(image_data, pixel_key, progress_file)
    except Exception as e:
        debug.print(f"DCT extraction failed: {e}")
        return None


def _extract_lsb(image_data: bytes, pixel_key: bytes, bits_per_channel: int = 1) -> bytes | None:
    """
    Extract using LSB mode (internal implementation).
    """
    debug.print(f"LSB extracting from {len(image_data)} byte image")
    debug.data(pixel_key, "Pixel key for extraction")
    debug.validate(
        bits_per_channel in (1, 2), f"bits_per_channel must be 1 or 2, got {bits_per_channel}"
    )

    img_file = None
    img = None

    try:
        img_file = Image.open(io.BytesIO(image_data))
        debug.print(f"Image: {img_file.size[0]}x{img_file.size[1]}, format: {img_file.format}")

        img = img_file.convert("RGB") if img_file.mode != "RGB" else img_file.copy()
        if img_file.mode != "RGB":
            debug.print(f"Converting image from {img_file.mode} to RGB")

        pixels = list(img.getdata())
        num_pixels = len(pixels)
        bits_per_pixel = 3 * bits_per_channel

        debug.print(f"Image has {num_pixels} pixels, {bits_per_pixel} bits/pixel")

        initial_pixels = (32 + bits_per_pixel - 1) // bits_per_pixel + 10
        debug.print(f"Extracting initial {initial_pixels} pixels to find length")

        initial_indices = generate_pixel_indices(pixel_key, num_pixels, initial_pixels)

        binary_data = ""
        for pixel_idx in initial_indices:
            r, g, b = pixels[pixel_idx]
            for channel in [r, g, b]:
                for bit_pos in range(bits_per_channel - 1, -1, -1):
                    binary_data += str((channel >> bit_pos) & 1)

        try:
            length_bits = binary_data[:32]
            if len(length_bits) < 32:
                debug.print(f"Not enough bits for length: {len(length_bits)}/32")
                return None

            data_length = struct.unpack(">I", int(length_bits, 2).to_bytes(4, "big"))[0]
            debug.print(f"Extracted length: {data_length} bytes")
        except Exception as e:
            debug.print(f"Failed to parse length: {e}")
            return None

        max_possible = (num_pixels * bits_per_pixel) // 8 - 4
        if data_length > max_possible or data_length < 10:
            debug.print(f"Invalid data length: {data_length} (max possible: {max_possible})")
            return None

        total_bits = (4 + data_length) * 8
        pixels_needed = (total_bits + bits_per_pixel - 1) // bits_per_pixel

        debug.print(f"Need {pixels_needed} pixels to extract {data_length} bytes")

        selected_indices = generate_pixel_indices(pixel_key, num_pixels, pixels_needed)

        binary_data = ""
        for pixel_idx in selected_indices:
            r, g, b = pixels[pixel_idx]
            for channel in [r, g, b]:
                for bit_pos in range(bits_per_channel - 1, -1, -1):
                    binary_data += str((channel >> bit_pos) & 1)

        data_bits = binary_data[32 : 32 + (data_length * 8)]

        if len(data_bits) < data_length * 8:
            debug.print(f"Insufficient bits: {len(data_bits)} < {data_length * 8}")
            return None

        data_bytes = bytearray()
        for i in range(0, len(data_bits), 8):
            byte_bits = data_bits[i : i + 8]
            if len(byte_bits) == 8:
                data_bytes.append(int(byte_bits, 2))

        debug.print(f"LSB successfully extracted {len(data_bytes)} bytes")
        return bytes(data_bytes)

    except Exception as e:
        debug.exception(e, "extract_lsb")
        return None
    finally:
        # Properly close all PIL Images to prevent memory leaks
        if img is not None and img is not img_file:
            img.close()
        if img_file is not None:
            img_file.close()


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================


def get_image_dimensions(image_data: bytes) -> tuple[int, int]:
    """Get image dimensions without loading full image."""
    debug.validate(len(image_data) > 0, "Image data cannot be empty")
    img = Image.open(io.BytesIO(image_data))
    try:
        dimensions = img.size
        debug.print(f"Image dimensions: {dimensions[0]}x{dimensions[1]}")
        return dimensions
    finally:
        img.close()


def get_image_format(image_data: bytes) -> str | None:
    """Get image format (PIL format string like 'PNG', 'JPEG')."""
    try:
        img = Image.open(io.BytesIO(image_data))
        try:
            format_str = img.format
            debug.print(f"Image format: {format_str}")
            return format_str
        finally:
            img.close()
    except Exception as e:
        debug.print(f"Failed to get image format: {e}")
        return None


def is_lossless_format(image_data: bytes) -> bool:
    """Check if image is in a lossless format suitable for steganography."""
    fmt = get_image_format(image_data)
    is_lossless = fmt is not None and fmt.upper() in LOSSLESS_FORMATS
    debug.print(f"Image is lossless: {is_lossless} (format: {fmt})")
    return is_lossless


def peek_image(image_data: bytes) -> dict:
    """
    Check if an image contains Stegasoo hidden data without decrypting.

    Attempts to detect LSB and DCT headers by extracting the first few bytes
    and looking for Stegasoo magic markers.

    Args:
        image_data: Raw image bytes

    Returns:
        dict with:
            - has_stegasoo: bool - True if header detected
            - mode: str or None - 'lsb', 'dct', or None
            - confidence: str - 'high', 'low', or None

    Example:
        >>> result = peek_image(suspicious_image_bytes)
        >>> if result['has_stegasoo']:
        ...     print(f"Found {result['mode']} data!")
    """
    from .constants import EMBED_MODE_DCT, EMBED_MODE_LSB

    result = {"has_stegasoo": False, "mode": None, "confidence": None}

    # Try LSB extraction (look for header bytes)
    try:
        img = Image.open(io.BytesIO(image_data))
        pixels = list(img.getdata())
        img.close()

        # Extract first 32 bits (4 bytes) from LSB
        extracted = []
        for i in range(32):
            if i < len(pixels):
                pixel = pixels[i]
                if isinstance(pixel, tuple):
                    extracted.append(pixel[0] & 1)
                else:
                    extracted.append(pixel & 1)

        # Convert bits to bytes
        header_bytes = bytearray()
        for i in range(0, len(extracted), 8):
            byte = 0
            for j in range(8):
                if i + j < len(extracted):
                    byte = (byte << 1) | extracted[i + j]
            header_bytes.append(byte)

        # Check for LSB magic: \x89ST3
        if bytes(header_bytes[:4]) == b"\x89ST3":
            result["has_stegasoo"] = True
            result["mode"] = EMBED_MODE_LSB
            result["confidence"] = "high"
            return result
    except Exception:
        pass

    # Try DCT extraction (requires scipy/jpeglib)
    try:
        from .dct_steganography import HAS_JPEGLIB, HAS_SCIPY

        if HAS_SCIPY or HAS_JPEGLIB:
            from .dct_steganography import extract_from_dct

            # Extract first few bytes to check header
            extracted = extract_from_dct(image_data, seed=b"\x00" * 32, length=4)
            if extracted == b"\x89DCT":
                result["has_stegasoo"] = True
                result["mode"] = EMBED_MODE_DCT
                result["confidence"] = "high"
                return result
    except Exception:
        pass

    return result
