"""
Spread Spectrum Audio Steganography Module (v4.4.0 — Per-Channel Hybrid)

Hides data in audio by adding keyed pseudo-random noise (spread spectrum)
below the threshold of audibility. Designed to survive lossy compression
(MP3, AAC, Opus) better than LSB embedding, which requires lossless carriers.

v4.4.0 changes:
  - Per-channel independent embedding (preserves spatial stereo/surround mix)
  - Adaptive chip length tiers (256/512/1024) for capacity vs robustness
  - LFE channel skipping for 5.1+ layouts
  - Round-robin bit distribution across embeddable channels
  - v2 header format with backward-compatible v0 decode fallback

How it works:
  Each payload bit is "spread" over chip_length audio samples using a unique
  ChaCha20-derived chip sequence. A '1' bit adds the chip pattern; a '0' bit
  subtracts it. On extraction, correlating the stego audio against the same
  chip sequence recovers each bit.

v2 data layout in the carrier:
  Header (20 bytes, channel 0 only, chip=1024 for robustness):
    [4B magic AUDS] [1B version=0x02] [1B chip_tier] [1B num_channels]
    [1B flags] [4B length x3 copies]
  Payload (round-robin across embeddable channels, chip=tier-selected):
    RS-encoded payload bits distributed: bit 0→ch0, bit 1→ch1, ...

  Legacy v0 layout (backward compat):
    [4B magic AUDS] [4B length x3 copies] [RS-encoded payload] — all mono

Error correction:
  The raw payload is protected with Reed-Solomon coding (AUDIO_SS_RS_NSYM
  parity symbols per 255-byte block) so that bit errors introduced by
  compression or DAC/ADC round-trips can be corrected transparently.

Requires: soundfile, numpy, cryptography, reedsolo (optional but recommended)
"""

from __future__ import annotations

import io
import struct

import numpy as np
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms

from .constants import (
    AUDIO_LFE_CHANNEL_INDEX,
    AUDIO_LFE_MIN_CHANNELS,
    AUDIO_MAGIC_SPREAD,
    AUDIO_SS_AMPLITUDE,
    AUDIO_SS_CHIP_LENGTH,
    AUDIO_SS_CHIP_LENGTHS,
    AUDIO_SS_DEFAULT_CHIP_TIER,
    AUDIO_SS_HEADER_VERSION,
    AUDIO_SS_RS_NSYM,
    EMBED_MODE_AUDIO_SPREAD,
)
from .debug import debug
from .exceptions import AudioCapacityError, AudioError
from .models import AudioCapacityInfo, AudioEmbedStats

# Lazy import for soundfile
try:
    import soundfile as sf

    HAS_SOUNDFILE = True
except ImportError:
    HAS_SOUNDFILE = False
    sf = None  # type: ignore[assignment]

# Lazy import for reedsolo
try:
    from reedsolo import ReedSolomonError, RSCodec

    HAS_REEDSOLO = True
except ImportError:
    HAS_REEDSOLO = False
    RSCodec = None  # type: ignore[assignment,misc]
    ReedSolomonError = None  # type: ignore[assignment,misc]


# Header sizes
_V0_HEADER_SIZE = 16  # Legacy: 4B magic + 3x4B length
_V2_HEADER_SIZE = 20  # v2: 4B magic + 1B ver + 1B tier + 1B nch + 1B flags + 3x4B length
_MAGIC_SIZE = 4
_LENGTH_COPIES = 3

# Header chip length — always 1024 for maximum robustness regardless of payload tier
_HEADER_CHIP_LENGTH = 1024

# v2 header flags
_FLAG_LFE_SKIPPED = 0x01

# Progress reporting interval (every N bits)
_PROGRESS_INTERVAL = 500


# =============================================================================
# PROGRESS REPORTING
# =============================================================================


def _write_progress(
    progress_file: str | None, current: int, total: int, phase: str = "embedding"
) -> None:
    """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


# =============================================================================
# REED-SOLOMON
# =============================================================================


def _rs_encode(data: bytes) -> bytes:
    """
    Wrap data in Reed-Solomon error correction.

    Adds AUDIO_SS_RS_NSYM parity symbols per 255-byte block, allowing
    recovery of up to RS_NSYM/2 byte errors per block.
    """
    if not HAS_REEDSOLO:
        return data
    rs = RSCodec(AUDIO_SS_RS_NSYM)
    return bytes(rs.encode(data))


def _rs_decode(data: bytes) -> bytes | None:
    """
    Decode Reed-Solomon protected data.

    Returns the corrected payload bytes, or None if the data is
    too corrupted for error correction to recover.
    """
    if not HAS_REEDSOLO:
        return data
    rs = RSCodec(AUDIO_SS_RS_NSYM)
    try:
        decoded, _, errata_pos = rs.decode(data)
        if errata_pos:
            debug.print(f"RS corrected {len(errata_pos)} byte errors")
        return bytes(decoded)
    except ReedSolomonError as e:
        debug.print(f"RS decode failed (too many errors): {e}")
        return None


def _rs_encoded_size(data_length: int) -> int:
    """Calculate the RS-encoded size for a given data length."""
    if not HAS_REEDSOLO or AUDIO_SS_RS_NSYM <= 0:
        return data_length
    data_block_size = 255 - AUDIO_SS_RS_NSYM
    num_blocks = (data_length + data_block_size - 1) // data_block_size
    return data_length + num_blocks * AUDIO_SS_RS_NSYM


# =============================================================================
# CHANNEL MANAGEMENT
# =============================================================================


def _embeddable_channels(num_channels: int) -> list[int]:
    """
    Return the list of channel indices suitable for embedding.

    Skips the LFE channel (index 3) for 5.1+ layouts because LFE is
    bandlimited to ~120Hz — terrible as a spread spectrum carrier.
    """
    if num_channels < AUDIO_LFE_MIN_CHANNELS:
        return list(range(num_channels))
    return [i for i in range(num_channels) if i != AUDIO_LFE_CHANNEL_INDEX]


# =============================================================================
# CHIP SEQUENCE GENERATION (ChaCha20 CSPRNG)
# =============================================================================


def _generate_chip_sequence(
    seed: bytes, chip_index: int, length: int, channel_index: int = 0
) -> np.ndarray:
    """
    Generate a pseudo-random chip sequence for spread spectrum embedding.

    Uses ChaCha20 as a CSPRNG keyed by ``seed``, with ``chip_index`` and
    ``channel_index`` encoded into the nonce so that each (channel, bit)
    pair gets a unique, deterministic spreading code.

    Args:
        seed: 32-byte key for ChaCha20. Padded/hashed to 32B if shorter.
        chip_index: Index of the bit being embedded (used as nonce material).
        length: Number of samples in the chip.
        channel_index: Audio channel index (0 for legacy/header, per-channel for v2).

    Returns:
        Float64 numpy array of ``length`` elements in bipolar {-1, +1}.
    """
    # Ensure seed is exactly 32 bytes
    if len(seed) < 32:
        import hashlib

        seed = hashlib.sha256(seed).digest()
    elif len(seed) > 32:
        seed = seed[:32]

    # Build a 16-byte nonce encoding both channel and chip index
    # v2: channel_index in high 4 bytes, chip_index in low 12 bytes
    # v0 compat: channel_index=0 → nonce is just chip_index in full 16 bytes,
    #            which matches the old layout since high bytes are 0
    nonce = channel_index.to_bytes(4, byteorder="big") + chip_index.to_bytes(12, byteorder="big")

    cipher = Cipher(algorithms.ChaCha20(seed, nonce), mode=None, backend=default_backend())
    encryptor = cipher.encryptor()
    random_bytes = encryptor.update(b"\x00" * length)

    # Map bytes to bipolar ±1 spreading code (DSSS standard)
    raw = np.frombuffer(random_bytes, dtype=np.uint8)
    chip = np.where(raw < 128, np.float64(-1.0), np.float64(1.0))

    return chip


def _generate_chip_sequence_v0(seed: bytes, chip_index: int, length: int) -> np.ndarray:
    """Legacy v0 chip sequence generation — chip_index fills entire 16-byte nonce."""
    if len(seed) < 32:
        import hashlib

        seed = hashlib.sha256(seed).digest()
    elif len(seed) > 32:
        seed = seed[:32]

    nonce = chip_index.to_bytes(16, byteorder="big")
    cipher = Cipher(algorithms.ChaCha20(seed, nonce), mode=None, backend=default_backend())
    encryptor = cipher.encryptor()
    random_bytes = encryptor.update(b"\x00" * length)

    raw = np.frombuffer(random_bytes, dtype=np.uint8)
    return np.where(raw < 128, np.float64(-1.0), np.float64(1.0))


# =============================================================================
# BIT CONVERSION UTILITIES
# =============================================================================


def _bytes_to_bits(data: bytes) -> list[int]:
    """Convert a byte string to a list of 0/1 ints (MSB first per byte)."""
    bits: list[int] = []
    for byte in data:
        for shift in range(7, -1, -1):
            bits.append((byte >> shift) & 1)
    return bits


def _bits_to_bytes(bits: list[int]) -> bytes:
    """Convert a list of 0/1 ints back to bytes (MSB first per byte)."""
    result = bytearray()
    for i in range(0, len(bits) - 7, 8):
        byte_val = 0
        for j in range(8):
            byte_val = (byte_val << 1) | bits[i + j]
        result.append(byte_val)
    return bytes(result)


# =============================================================================
# MAJORITY VOTING
# =============================================================================


def _majority_vote_length(length_bytes: bytes) -> int | None:
    """
    Extract the payload length from three 4-byte copies via majority voting.

    Each copy is a big-endian uint32. The value that appears at least twice
    wins. Returns None if all three disagree.
    """
    if len(length_bytes) < 12:
        return None

    copies = [
        struct.unpack(">I", length_bytes[0:4])[0],
        struct.unpack(">I", length_bytes[4:8])[0],
        struct.unpack(">I", length_bytes[8:12])[0],
    ]

    debug.print(f"Length copies for majority vote: {copies}")

    if copies[0] == copies[1] or copies[0] == copies[2]:
        return copies[0]
    if copies[1] == copies[2]:
        return copies[1]

    debug.print("Majority vote failed: all three length copies disagree")
    return None


# =============================================================================
# HEADER CONSTRUCTION & PARSING
# =============================================================================


def _build_header_v0(data_length: int) -> bytes:
    """Build legacy v0 header: 4B magic + 3x4B length = 16 bytes."""
    length_packed = struct.pack(">I", data_length)
    return AUDIO_MAGIC_SPREAD + length_packed * _LENGTH_COPIES


def _build_header_v2(
    data_length: int, chip_tier: int, num_embeddable_channels: int, lfe_skipped: bool
) -> bytes:
    """
    Build v2 header (20 bytes).

    Layout:
      [4B magic AUDS] [1B version=0x02] [1B chip_tier] [1B num_channels]
      [1B flags] [4B length x3 copies]
    """
    flags = _FLAG_LFE_SKIPPED if lfe_skipped else 0
    length_packed = struct.pack(">I", data_length)
    return (
        AUDIO_MAGIC_SPREAD
        + struct.pack("BBBB", AUDIO_SS_HEADER_VERSION, chip_tier, num_embeddable_channels, flags)
        + length_packed * _LENGTH_COPIES
    )


def _parse_header(
    header_bytes: bytes,
) -> tuple[bool, int, int | None, int | None, int | None, bool]:
    """
    Parse and validate the spread spectrum header (version-dispatching).

    Reads byte 4 after magic to determine version:
      - 0x02 → v2 header (20 bytes)
      - anything else → v0 header (16 bytes, legacy mono)

    Returns:
        (magic_valid, version, payload_length, chip_tier, num_channels, lfe_skipped)
        version is 0 for legacy, 2 for v2.
        payload_length is None if majority vote fails.
        chip_tier/num_channels are None for v0.
    """
    if len(header_bytes) < _V0_HEADER_SIZE:
        return False, 0, None, None, None, False

    magic = header_bytes[:_MAGIC_SIZE]
    if magic != AUDIO_MAGIC_SPREAD:
        debug.print(f"Magic mismatch: got {magic!r}, expected {AUDIO_MAGIC_SPREAD!r}")
        return False, 0, None, None, None, False

    # Check version byte
    version_byte = header_bytes[4]

    if version_byte == AUDIO_SS_HEADER_VERSION:
        # v2 header — 20 bytes
        if len(header_bytes) < _V2_HEADER_SIZE:
            return False, 2, None, None, None, False

        chip_tier = header_bytes[5]
        num_channels = header_bytes[6]
        flags = header_bytes[7]
        lfe_skipped = bool(flags & _FLAG_LFE_SKIPPED)

        length = _majority_vote_length(header_bytes[8:20])
        debug.print(
            f"v2 header: tier={chip_tier}, channels={num_channels}, "
            f"lfe_skip={lfe_skipped}, length={length}"
        )
        return True, 2, length, chip_tier, num_channels, lfe_skipped
    else:
        # v0 header — 16 bytes, byte 4 is part of the first length copy
        length = _majority_vote_length(header_bytes[_MAGIC_SIZE:_V0_HEADER_SIZE])
        debug.print(f"v0 header: length={length}")
        return True, 0, length, None, None, False


# =============================================================================
# SPREAD SPECTRUM CORE — PER-CHANNEL (v2)
# =============================================================================


def _embed_channel(
    channel_samples: np.ndarray,
    bits: list[int],
    seed: bytes,
    amplitude: float,
    chip_length: int,
    channel_index: int,
    offset: int = 0,
) -> int:
    """
    Embed bits into a single channel's samples. Returns count of bits embedded.
    """
    count = 0
    for i, bit in enumerate(bits):
        start = offset + i * chip_length
        end = start + chip_length

        if end > len(channel_samples):
            break

        chip = _generate_chip_sequence(seed, i, chip_length, channel_index=channel_index)

        if bit == 1:
            channel_samples[start:end] += amplitude * chip
        else:
            channel_samples[start:end] -= amplitude * chip
        count += 1

    return count


def _extract_channel(
    channel_samples: np.ndarray,
    num_bits: int,
    seed: bytes,
    chip_length: int,
    channel_index: int,
    offset: int = 0,
) -> list[int]:
    """Extract bits from a single channel's samples."""
    bits: list[int] = []
    for i in range(num_bits):
        start = offset + i * chip_length
        end = start + chip_length

        if end > len(channel_samples):
            break

        chip = _generate_chip_sequence(seed, i, chip_length, channel_index=channel_index)
        correlation = np.dot(channel_samples[start:end], chip)
        bits.append(1 if correlation > 0 else 0)

    return bits


def _distribute_bits_round_robin(
    bits: list[int], num_channels: int
) -> list[list[int]]:
    """
    Distribute bits round-robin across channels.

    bit 0 → ch0, bit 1 → ch1, ..., bit N → ch(N % num_channels), ...

    Returns a list of per-channel bit lists.
    """
    per_channel: list[list[int]] = [[] for _ in range(num_channels)]
    for i, bit in enumerate(bits):
        per_channel[i % num_channels].append(bit)
    return per_channel


def _collect_bits_round_robin(per_channel_bits: list[list[int]]) -> list[int]:
    """
    Reassemble bits from per-channel lists back into original order.

    Inverse of _distribute_bits_round_robin.
    """
    num_channels = len(per_channel_bits)
    if num_channels == 0:
        return []

    max_len = max(len(ch) for ch in per_channel_bits)
    result: list[int] = []
    for i in range(max_len):
        for ch in range(num_channels):
            if i < len(per_channel_bits[ch]):
                result.append(per_channel_bits[ch][i])
    return result


# =============================================================================
# LEGACY MONO CORE (v0 compat)
# =============================================================================


def _embed_spread_spectrum_v0(
    samples: np.ndarray,
    bits: list[int],
    seed: bytes,
    amplitude: float,
    offset: int = 0,
    progress_file: str | None = None,
) -> np.ndarray:
    """Legacy v0 mono embedding with original nonce layout."""
    total_bits = len(bits)
    for i, bit in enumerate(bits):
        start = offset + i * AUDIO_SS_CHIP_LENGTH
        end = start + AUDIO_SS_CHIP_LENGTH

        if end > len(samples):
            debug.print(f"Warning: ran out of samples at bit {i}/{total_bits}")
            break

        chip = _generate_chip_sequence_v0(seed, i, AUDIO_SS_CHIP_LENGTH)

        if bit == 1:
            samples[start:end] += amplitude * chip
        else:
            samples[start:end] -= amplitude * chip

        if progress_file and i % _PROGRESS_INTERVAL == 0:
            _write_progress(progress_file, i, total_bits, "embedding")

    return samples


def _extract_spread_spectrum_v0(
    samples: np.ndarray,
    num_bits: int,
    seed: bytes,
    offset: int = 0,
    progress_file: str | None = None,
) -> list[int]:
    """Legacy v0 mono extraction with original nonce layout."""
    bits: list[int] = []
    for i in range(num_bits):
        start = offset + i * AUDIO_SS_CHIP_LENGTH
        end = start + AUDIO_SS_CHIP_LENGTH

        if end > len(samples):
            debug.print(f"Warning: ran out of samples at bit {i}/{num_bits}")
            break

        chip = _generate_chip_sequence_v0(seed, i, AUDIO_SS_CHIP_LENGTH)
        correlation = np.dot(samples[start:end], chip)
        bits.append(1 if correlation > 0 else 0)

        if progress_file and i % _PROGRESS_INTERVAL == 0:
            _write_progress(progress_file, i, num_bits, "extracting")

    return bits


# =============================================================================
# PUBLIC API
# =============================================================================


def calculate_audio_spread_capacity(
    audio_data: bytes, chip_tier: int = AUDIO_SS_DEFAULT_CHIP_TIER
) -> AudioCapacityInfo:
    """
    Calculate embedding capacity for spread spectrum audio steganography.

    Args:
        audio_data: Raw bytes of a WAV file.
        chip_tier: Chip tier (0=lossless/256, 1=high_lossy/512, 2=low_lossy/1024).

    Returns:
        AudioCapacityInfo with capacity details.

    Raises:
        AudioError: If the audio cannot be read.
    """
    if not HAS_SOUNDFILE:
        raise AudioError("soundfile is required for audio spread spectrum steganography")

    try:
        info = sf.info(io.BytesIO(audio_data))
    except Exception as e:
        raise AudioError(f"Failed to read audio file: {e}") from e

    chip_length = AUDIO_SS_CHIP_LENGTHS.get(chip_tier, AUDIO_SS_CHIP_LENGTH)
    embed_channels = _embeddable_channels(info.channels)
    num_embed_ch = len(embed_channels)
    num_frames = info.frames
    duration = num_frames / info.samplerate

    # Header is always in channel 0 with chip=1024
    header_bits = _V2_HEADER_SIZE * 8  # 160 bits
    header_samples = header_bits * _HEADER_CHIP_LENGTH  # always 1024

    # Payload samples available per channel (after header eats into ch0)
    samples_per_channel_for_payload = num_frames - header_samples
    if samples_per_channel_for_payload <= 0:
        return AudioCapacityInfo(
            total_samples=num_frames * info.channels,
            usable_capacity_bytes=0,
            embed_mode=EMBED_MODE_AUDIO_SPREAD,
            sample_rate=info.samplerate,
            duration_seconds=duration,
            chip_tier=chip_tier,
            chip_length=chip_length,
            embeddable_channels=num_embed_ch,
            total_channels=info.channels,
        )

    # Bits per channel for payload
    bits_per_channel = samples_per_channel_for_payload // chip_length

    # Total payload bits across all embeddable channels
    total_payload_bits = bits_per_channel * num_embed_ch
    total_payload_bytes = total_payload_bits // 8

    # Account for RS overhead
    if HAS_REEDSOLO and AUDIO_SS_RS_NSYM > 0:
        usable_bytes = int(total_payload_bytes * (255 - AUDIO_SS_RS_NSYM) / 255)
    else:
        usable_bytes = total_payload_bytes

    debug.print(
        f"Spread capacity (tier {chip_tier}, chip={chip_length}): {usable_bytes} bytes "
        f"({num_frames} frames, {num_embed_ch}/{info.channels} ch, "
        f"{info.samplerate} Hz, {duration:.2f}s)"
    )

    return AudioCapacityInfo(
        total_samples=num_frames * info.channels,
        usable_capacity_bytes=usable_bytes,
        embed_mode=EMBED_MODE_AUDIO_SPREAD,
        sample_rate=info.samplerate,
        duration_seconds=duration,
        chip_tier=chip_tier,
        chip_length=chip_length,
        embeddable_channels=num_embed_ch,
        total_channels=info.channels,
    )


def embed_in_audio_spread(
    data: bytes,
    carrier_audio: bytes,
    seed: bytes,
    chip_tier: int = AUDIO_SS_DEFAULT_CHIP_TIER,
    progress_file: str | None = None,
) -> tuple[bytes, AudioEmbedStats]:
    """
    Embed data into audio using per-channel spread spectrum steganography.

    v4.4.0: Embeds independently per channel with round-robin bit distribution,
    preserving the spatial stereo/surround mix. Header is always in channel 0
    with chip=1024 for robustness. Payload uses the selected chip tier.

    Args:
        data: Raw payload bytes to embed (already encrypted by caller).
        carrier_audio: Raw bytes of the carrier WAV file.
        seed: Key material for chip sequence generation.
        chip_tier: Chip tier (0=256, 1=512, 2=1024). Default 2 (most robust).
        progress_file: Optional path for frontend progress polling.

    Returns:
        Tuple of (stego WAV bytes, AudioEmbedStats).

    Raises:
        AudioCapacityError: If the payload is too large for the carrier.
        AudioError: On any other embedding failure.
    """
    if not HAS_SOUNDFILE:
        raise AudioError("soundfile is required for audio spread spectrum steganography")

    debug.print(f"Spread spectrum v2 embedding {len(data)} bytes, tier={chip_tier}")

    try:
        # 1. Read carrier audio as float64 2D (frames x channels)
        buf = io.BytesIO(carrier_audio)
        samples, sample_rate = sf.read(buf, dtype="float64", always_2d=True)
        num_frames, channels = samples.shape
        duration = num_frames / sample_rate

        buf.seek(0)
        carrier_info = sf.info(buf)
        output_subtype = carrier_info.subtype if carrier_info.subtype else "PCM_16"

        # 2. Build embeddable channels list
        embed_ch = _embeddable_channels(channels)
        num_embed_ch = len(embed_ch)
        lfe_skipped = len(embed_ch) < channels
        chip_length = AUDIO_SS_CHIP_LENGTHS.get(chip_tier, AUDIO_SS_CHIP_LENGTH)

        debug.print(
            f"Carrier: {sample_rate} Hz, {channels} ch ({num_embed_ch} embeddable), "
            f"{num_frames} frames, {duration:.2f}s, chip={chip_length}"
        )

        # 3. RS-encode the payload
        rs_data = _rs_encode(data)
        debug.print(f"RS-encoded payload: {len(data)} -> {len(rs_data)} bytes")

        # 4. Build v2 header
        header = _build_header_v2(len(data), chip_tier, num_embed_ch, lfe_skipped)
        header_bits = _bytes_to_bits(header)

        # 5. Embed header in channel 0 only, chip=1024, starting at sample 0
        header_samples_needed = len(header_bits) * _HEADER_CHIP_LENGTH
        if header_samples_needed > num_frames:
            raise AudioCapacityError(len(header), num_frames // _HEADER_CHIP_LENGTH // 8)

        _embed_channel(
            samples[:, embed_ch[0]],
            header_bits,
            seed,
            AUDIO_SS_AMPLITUDE,
            _HEADER_CHIP_LENGTH,
            channel_index=0,
            offset=0,
        )

        # 6. Calculate payload offset (in samples, same for all channels)
        payload_offset = header_samples_needed

        # 7. Convert payload to bits and distribute round-robin
        payload_bits = _bytes_to_bits(rs_data)
        total_payload_bits = len(payload_bits)

        # Check capacity across all channels
        payload_samples_available = num_frames - payload_offset
        bits_per_channel = payload_samples_available // chip_length
        total_capacity_bits = bits_per_channel * num_embed_ch

        if total_payload_bits > total_capacity_bits:
            max_bytes = (total_capacity_bits // 8)
            if HAS_REEDSOLO and AUDIO_SS_RS_NSYM > 0:
                max_bytes = int(max_bytes * (255 - AUDIO_SS_RS_NSYM) / 255)
            raise AudioCapacityError(len(data), max_bytes)

        per_channel_bits = _distribute_bits_round_robin(payload_bits, num_embed_ch)

        debug.print(
            f"Distributing {total_payload_bits} payload bits across {num_embed_ch} channels "
            f"(~{total_payload_bits // num_embed_ch} bits/ch)"
        )

        _write_progress(progress_file, 0, total_payload_bits, "embedding")

        # 8. Embed payload bits per channel
        total_embedded = 0
        for ch_idx, ch in enumerate(embed_ch):
            bits_for_ch = per_channel_bits[ch_idx]
            if not bits_for_ch:
                continue
            count = _embed_channel(
                samples[:, ch],
                bits_for_ch,
                seed,
                AUDIO_SS_AMPLITUDE,
                chip_length,
                channel_index=ch,
                offset=payload_offset,
            )
            total_embedded += count

            if progress_file:
                _write_progress(progress_file, total_embedded, total_payload_bits, "embedding")

        # 9. Clip to [-1.0, 1.0]
        np.clip(samples, -1.0, 1.0, out=samples)

        _write_progress(progress_file, total_payload_bits, total_payload_bits, "saving")

        # 10. Write back as WAV
        output_buf = io.BytesIO()
        sf.write(output_buf, samples, sample_rate, format="WAV", subtype=output_subtype)
        output_buf.seek(0)
        stego_bytes = output_buf.getvalue()

        # Calculate capacity used
        header_capacity = header_samples_needed  # in ch0
        payload_capacity = total_embedded * chip_length  # across all channels
        total_sample_slots = num_frames * num_embed_ch
        capacity_used = (header_capacity + payload_capacity) / total_sample_slots

        stats = AudioEmbedStats(
            samples_modified=header_capacity + payload_capacity,
            total_samples=num_frames * channels,
            capacity_used=capacity_used,
            bytes_embedded=len(header) + len(rs_data),
            sample_rate=sample_rate,
            channels=channels,
            duration_seconds=duration,
            embed_mode=EMBED_MODE_AUDIO_SPREAD,
            chip_tier=chip_tier,
            chip_length=chip_length,
            embeddable_channels=num_embed_ch,
        )

        debug.print(
            f"Spread spectrum v2 embedding complete: {len(stego_bytes)} byte WAV, "
            f"capacity used {capacity_used * 100:.1f}%"
        )
        return stego_bytes, stats

    except AudioCapacityError:
        raise
    except Exception as e:
        debug.exception(e, "embed_in_audio_spread")
        raise AudioError(f"Failed to embed data in audio via spread spectrum: {e}") from e


def extract_from_audio_spread(
    audio_data: bytes,
    seed: bytes,
    progress_file: str | None = None,
) -> bytes | None:
    """
    Extract hidden data from audio using spread spectrum correlation.

    Auto-detects header version:
      - v2: per-channel extraction with round-robin reassembly
      - v0: legacy mono extraction (backward compat)

    Args:
        audio_data: Raw bytes of the stego WAV file.
        seed: Key material (must match the seed used for embedding).
        progress_file: Optional path for frontend progress polling.

    Returns:
        Extracted payload bytes, or None if extraction fails.
    """
    if not HAS_SOUNDFILE:
        debug.print("soundfile not available for spread spectrum extraction")
        return None

    debug.print(f"Spread spectrum extracting from {len(audio_data)} byte audio")

    try:
        # 1. Read stego audio as float64 2D
        samples, sample_rate = sf.read(io.BytesIO(audio_data), dtype="float64", always_2d=True)
        num_frames, channels = samples.shape

        debug.print(f"Stego audio: {sample_rate} Hz, {channels} ch, {num_frames} frames")

        # 2. Extract header from channel 0 with chip=1024
        #    We need at least v2 header size (20 bytes = 160 bits)
        header_bits_needed = _V2_HEADER_SIZE * 8
        header_samples_needed = header_bits_needed * _HEADER_CHIP_LENGTH

        # Get channel 0 (first embeddable channel — always index 0 for header)
        embed_ch = _embeddable_channels(channels)
        ch0_samples = samples[:, embed_ch[0]]

        if header_samples_needed > len(ch0_samples):
            debug.print("Audio too short to contain spread spectrum header")
            return None

        _write_progress(progress_file, 0, header_bits_needed, "extracting header")

        # Extract header bits from channel 0 with channel_index=0
        header_bits = _extract_channel(
            ch0_samples,
            header_bits_needed,
            seed,
            _HEADER_CHIP_LENGTH,
            channel_index=0,
            offset=0,
        )

        if len(header_bits) < header_bits_needed:
            debug.print(f"Short header extraction: {len(header_bits)}/{header_bits_needed}")
            return None

        header_bytes = _bits_to_bytes(header_bits)

        # 3. Parse header (version-dispatching)
        magic_valid, version, data_length, chip_tier, num_ch, lfe_skipped = _parse_header(
            header_bytes
        )

        if not magic_valid:
            # Try legacy v0 with mono extraction
            debug.print("v2 header not found, trying legacy v0 mono extraction")
            return _extract_v0_fallback(samples, seed, progress_file)

        if data_length is None:
            debug.print("Could not determine payload length (majority vote failed)")
            return None

        if version == 0:
            # v0 header detected — use legacy mono path
            debug.print(f"v0 header detected, payload_length={data_length}")
            return _extract_v0_with_known_length(samples, seed, data_length, progress_file)

        # v2 extraction
        debug.print(
            f"v2 header: payload_length={data_length}, tier={chip_tier}, "
            f"channels={num_ch}, lfe_skip={lfe_skipped}"
        )

        chip_length = AUDIO_SS_CHIP_LENGTHS.get(chip_tier, AUDIO_SS_CHIP_LENGTH)
        payload_offset = _V2_HEADER_SIZE * 8 * _HEADER_CHIP_LENGTH

        # 4. Calculate RS-encoded size
        rs_encoded_size = _rs_encoded_size(data_length)
        total_payload_bits = rs_encoded_size * 8

        # Sanity check
        payload_samples_available = num_frames - payload_offset
        bits_per_channel = payload_samples_available // chip_length
        max_capacity_bits = bits_per_channel * num_ch

        if total_payload_bits > max_capacity_bits:
            debug.print(
                f"Payload too large for carrier: need {total_payload_bits} bits, "
                f"have {max_capacity_bits}"
            )
            return None

        # 5. Calculate how many bits per channel
        bits_per_ch_needed = [0] * num_ch
        for i in range(total_payload_bits):
            bits_per_ch_needed[i % num_ch] += 1

        debug.print(
            f"Extracting {total_payload_bits} payload bits from {num_ch} channels "
            f"(chip={chip_length})"
        )

        _write_progress(progress_file, 0, total_payload_bits, "extracting")

        # 6. Extract per-channel bits
        per_channel_bits: list[list[int]] = []
        total_extracted = 0

        for ch_idx in range(num_ch):
            # Map ch_idx back to actual channel index
            ch = embed_ch[ch_idx] if ch_idx < len(embed_ch) else ch_idx

            ch_bits = _extract_channel(
                samples[:, ch],
                bits_per_ch_needed[ch_idx],
                seed,
                chip_length,
                channel_index=ch,
                offset=payload_offset,
            )
            per_channel_bits.append(ch_bits)
            total_extracted += len(ch_bits)

            if progress_file:
                _write_progress(progress_file, total_extracted, total_payload_bits, "extracting")

        # 7. Reassemble bits in round-robin order
        all_payload_bits = _collect_bits_round_robin(per_channel_bits)

        if len(all_payload_bits) < total_payload_bits:
            debug.print(
                f"Short extraction: {len(all_payload_bits)}/{total_payload_bits} bits"
            )
            return None

        _write_progress(progress_file, total_payload_bits, total_payload_bits, "decoding")

        # 8. Convert to bytes and RS-decode
        rs_payload = _bits_to_bytes(all_payload_bits[:total_payload_bits])

        if len(rs_payload) < rs_encoded_size:
            debug.print(f"RS payload too short: {len(rs_payload)}/{rs_encoded_size} bytes")
            return None

        rs_payload = rs_payload[:rs_encoded_size]
        decoded = _rs_decode(rs_payload)
        if decoded is None:
            debug.print("Reed-Solomon decoding failed — data too corrupted")
            return None

        if len(decoded) < data_length:
            debug.print(f"Decoded data shorter than expected: {len(decoded)}/{data_length}")
            return None

        payload = decoded[:data_length]
        debug.print(f"Spread spectrum v2 extraction successful: {len(payload)} bytes")
        return payload

    except Exception as e:
        debug.exception(e, "extract_from_audio_spread")
        return None


# =============================================================================
# LEGACY V0 EXTRACTION (backward compat)
# =============================================================================


def _extract_v0_fallback(
    samples: np.ndarray, seed: bytes, progress_file: str | None = None
) -> bytes | None:
    """Full v0 extraction — read header + payload from mono mix."""
    channels = samples.shape[1]

    if channels > 1:
        mono_samples = np.mean(samples, axis=1)
    else:
        mono_samples = samples[:, 0].copy()

    total_samples = len(mono_samples)

    # Extract v0 header (16 bytes = 128 bits)
    header_bits_needed = _V0_HEADER_SIZE * 8
    header_samples_needed = header_bits_needed * AUDIO_SS_CHIP_LENGTH

    if header_samples_needed > total_samples:
        debug.print("Audio too short for v0 header")
        return None

    header_bits = _extract_spread_spectrum_v0(mono_samples, header_bits_needed, seed, offset=0)
    if len(header_bits) < header_bits_needed:
        return None

    header_bytes = _bits_to_bytes(header_bits)

    magic = header_bytes[:_MAGIC_SIZE]
    if magic != AUDIO_MAGIC_SPREAD:
        debug.print("v0 fallback: magic mismatch")
        return None

    data_length = _majority_vote_length(header_bytes[_MAGIC_SIZE:_V0_HEADER_SIZE])
    if data_length is None:
        return None

    return _extract_v0_with_known_length(samples, seed, data_length, progress_file)


def _extract_v0_with_known_length(
    samples: np.ndarray, seed: bytes, data_length: int, progress_file: str | None = None
) -> bytes | None:
    """v0 extraction with known payload length — mono mix path."""
    channels = samples.shape[1]

    if channels > 1:
        mono_samples = np.mean(samples, axis=1)
    else:
        mono_samples = samples[:, 0].copy()

    total_samples = len(mono_samples)

    # Sanity check
    max_payload = (total_samples // AUDIO_SS_CHIP_LENGTH) // 8 - _V0_HEADER_SIZE
    if data_length < 1 or data_length > max_payload:
        debug.print(f"v0: invalid payload length {data_length} (max {max_payload})")
        return None

    rs_encoded_size = _rs_encoded_size(data_length)
    total_payload_bytes = _V0_HEADER_SIZE + rs_encoded_size
    total_bits_needed = total_payload_bytes * 8
    total_samples_needed = total_bits_needed * AUDIO_SS_CHIP_LENGTH

    if total_samples_needed > total_samples:
        debug.print(f"v0: need {total_samples_needed} samples, have {total_samples}")
        return None

    debug.print(f"v0 extraction: {total_bits_needed} bits ({data_length}B payload)")

    _write_progress(progress_file, 0, total_bits_needed, "extracting")

    all_bits = _extract_spread_spectrum_v0(
        mono_samples, total_bits_needed, seed, offset=0, progress_file=progress_file
    )

    if len(all_bits) < total_bits_needed:
        return None

    _write_progress(progress_file, total_bits_needed, total_bits_needed, "decoding")

    all_bytes = _bits_to_bytes(all_bits)
    rs_payload = all_bytes[_V0_HEADER_SIZE : _V0_HEADER_SIZE + rs_encoded_size]

    if len(rs_payload) < rs_encoded_size:
        return None

    decoded = _rs_decode(rs_payload)
    if decoded is None:
        debug.print("v0: Reed-Solomon decoding failed")
        return None

    if len(decoded) < data_length:
        return None

    payload = decoded[:data_length]
    debug.print(f"v0 extraction successful: {len(payload)} bytes")
    return payload
