"""
Stegasoo Audio Steganography — LSB Embedding/Extraction (v4.3.0)

LSB (Least Significant Bit) embedding for PCM audio samples.

Hides data in the least significant bit(s) of audio samples, analogous to
how steganography.py hides data in pixel LSBs. The carrier audio must be
lossless (WAV or FLAC) — lossy codecs (MP3, OGG, AAC) destroy LSBs.

Uses ChaCha20 as a CSPRNG for pseudo-random sample index selection,
ensuring that without the key an attacker cannot determine which samples
were modified.

Supports:
- 16-bit PCM (int16 samples)
- 24-bit PCM (int32 samples from soundfile)
- Float audio (converted to int16 before embedding)
- 1 or 2 bits per sample embedding depth
- Mono and multi-channel audio (flattened for embedding)
"""

import io
import struct

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

from .constants import (
    AUDIO_MAGIC_LSB,
    EMBED_MODE_AUDIO_LSB,
)
from .debug import debug
from .exceptions import AudioCapacityError, AudioError
from .models import AudioEmbedStats
from .steganography import ENCRYPTION_OVERHEAD

# Progress reporting interval — write every N samples
PROGRESS_INTERVAL = 5000


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


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


# =============================================================================
# CAPACITY
# =============================================================================


def calculate_audio_lsb_capacity(
    audio_data: bytes,
    bits_per_sample: int = 1,
) -> int:
    """
    Calculate the maximum bytes that can be embedded in a WAV/FLAC file via LSB.

    Reads the carrier audio with soundfile, counts the total number of individual
    sample values (num_frames * channels), and computes how many payload bytes
    can be hidden at the given bit depth, minus the fixed encryption overhead.

    Args:
        audio_data: Raw bytes of a WAV or FLAC file.
        bits_per_sample: Number of LSBs to use per sample (1 or 2).

    Returns:
        Maximum embeddable payload size in bytes (after subtracting overhead).

    Raises:
        AudioError: If the audio cannot be read or is in an unsupported format.
    """
    debug.validate(
        bits_per_sample in (1, 2), f"bits_per_sample must be 1 or 2, got {bits_per_sample}"
    )

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

    num_samples = info.frames * info.channels
    total_bits = num_samples * bits_per_sample
    max_bytes = total_bits // 8

    capacity = max(0, max_bytes - ENCRYPTION_OVERHEAD)
    debug.print(
        f"Audio LSB capacity: {capacity} bytes "
        f"({num_samples} samples, {bits_per_sample} bit(s)/sample, "
        f"{info.samplerate} Hz, {info.channels} ch)"
    )
    return capacity


# =============================================================================
# SAMPLE INDEX GENERATION (ChaCha20 CSPRNG)
# =============================================================================
#
# Identical strategy to generate_pixel_indices in steganography.py:
# - >= 50% capacity utilisation: full Fisher-Yates shuffle, take first N
# - <  50%: direct random sampling with collision handling
#
# The key MUST be 32 bytes (same derivation path as the pixel key).


@debug.time
def generate_sample_indices(key: bytes, num_samples: int, num_needed: int) -> list[int]:
    """
    Generate pseudo-random sample indices using ChaCha20 as a CSPRNG.

    Produces a deterministic sequence of unique sample indices so that
    the same key always yields the same embedding locations.

    Args:
        key: 32-byte key for the ChaCha20 cipher.
        num_samples: Total number of samples in the carrier audio.
        num_needed: How many unique sample indices are required.

    Returns:
        List of ``num_needed`` unique indices in [0, num_samples).

    Raises:
        AssertionError (via debug.validate): On invalid arguments.
    """
    debug.validate(len(key) == 32, f"Sample key must be 32 bytes, got {len(key)}")
    debug.validate(num_samples > 0, f"Number of samples must be positive, got {num_samples}")
    debug.validate(num_needed > 0, f"Number needed must be positive, got {num_needed}")
    debug.validate(
        num_needed <= num_samples,
        f"Cannot select {num_needed} samples from {num_samples} available",
    )

    debug.print(f"Generating {num_needed} sample indices from {num_samples} total samples")

    # Strategy 1: Full Fisher-Yates shuffle when we need many indices
    if num_needed >= num_samples // 2:
        debug.print(f"Using full shuffle (needed {num_needed}/{num_samples} samples)")
        nonce = b"\x00" * 16
        cipher = Cipher(algorithms.ChaCha20(key, nonce), mode=None, backend=default_backend())
        encryptor = cipher.encryptor()

        indices = list(range(num_samples))
        random_bytes = encryptor.update(b"\x00" * (num_samples * 4))

        for i in range(num_samples - 1, 0, -1):
            j_bytes = random_bytes[(num_samples - 1 - i) * 4 : (num_samples - 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 for lower utilisation
    debug.print(f"Using optimized selection (needed {num_needed}/{num_samples} samples)")
    selected: list[int] = []
    used: set[int] = set()

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

    # Pre-generate 2x bytes to handle expected collisions
    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_samples
        byte_offset += 4

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

    # 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_samples
            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
# =============================================================================


@debug.time
def embed_in_audio_lsb(
    data: bytes,
    carrier_audio: bytes,
    sample_key: bytes,
    bits_per_sample: int = 1,
    progress_file: str | None = None,
) -> tuple[bytes, AudioEmbedStats]:
    """
    Embed data into PCM audio samples using LSB steganography.

    The payload is prepended with a 4-byte magic header (``AUDIO_MAGIC_LSB``)
    and a 4-byte big-endian length prefix, then converted to a binary string.
    Pseudo-random sample indices are generated from ``sample_key`` and the
    corresponding sample LSBs are overwritten.

    The modified audio is written back as a 16-bit PCM WAV file.

    Args:
        data: Encrypted payload bytes to embed.
        carrier_audio: Raw bytes of the carrier WAV/FLAC file.
        sample_key: 32-byte key for sample index generation.
        bits_per_sample: LSBs to use per sample (1 or 2).
        progress_file: Optional path for progress JSON (frontend 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.
    """
    debug.print(f"Audio LSB embedding {len(data)} bytes")
    debug.data(sample_key, "Sample key for embedding")
    debug.validate(
        bits_per_sample in (1, 2), f"bits_per_sample must be 1 or 2, got {bits_per_sample}"
    )
    debug.validate(len(sample_key) == 32, f"Sample key must be 32 bytes, got {len(sample_key)}")

    try:
        # 1. Read carrier audio
        samples, samplerate = sf.read(io.BytesIO(carrier_audio), dtype="int16", always_2d=True)
        # samples shape: (num_frames, channels)
        original_shape = samples.shape
        channels = original_shape[1]
        duration = original_shape[0] / samplerate

        debug.print(
            f"Carrier audio: {samplerate} Hz, {channels} ch, "
            f"{original_shape[0]} frames, {duration:.2f}s"
        )

        # Flatten to 1D for embedding
        flat_samples = samples.flatten().copy()
        num_samples = len(flat_samples)

        # 2. Prepend magic + length prefix
        header = AUDIO_MAGIC_LSB + struct.pack(">I", len(data))
        payload = header + data
        debug.print(
            f"Payload with header: {len(payload)} bytes (magic 4 + len 4 + data {len(data)})"
        )

        # 3. Check capacity
        max_bytes = (num_samples * bits_per_sample) // 8
        if len(payload) > max_bytes:
            debug.print(f"Capacity error: need {len(payload)}, have {max_bytes}")
            raise AudioCapacityError(len(payload), max_bytes)

        debug.print(
            f"Capacity usage: {len(payload)}/{max_bytes} bytes "
            f"({len(payload) / max_bytes * 100:.1f}%)"
        )

        # 4. Convert payload to binary string
        binary_data = "".join(format(b, "08b") for b in payload)
        samples_needed = (len(binary_data) + bits_per_sample - 1) // bits_per_sample

        debug.print(f"Need {samples_needed} samples to embed {len(binary_data)} bits")

        # 5. Generate pseudo-random sample indices
        selected_indices = generate_sample_indices(sample_key, num_samples, samples_needed)

        # 6. Modify LSBs of selected samples
        lsb_mask = (1 << bits_per_sample) - 1
        bit_idx = 0
        modified_count = 0
        total_to_process = len(selected_indices)

        # Initial progress
        if progress_file:
            _write_progress(progress_file, 5, 100, "embedding")

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

            bits = binary_data[bit_idx : bit_idx + bits_per_sample].ljust(bits_per_sample, "0")
            bit_val = int(bits, 2)

            sample_val = flat_samples[sample_idx]
            # Work in unsigned 16-bit space to avoid overflow
            unsigned_val = int(sample_val) & 0xFFFF
            new_unsigned = (unsigned_val & ~lsb_mask) | bit_val
            # Convert back to signed int16
            new_val = np.int16(new_unsigned if new_unsigned < 32768 else new_unsigned - 65536)

            if sample_val != new_val:
                flat_samples[sample_idx] = new_val
                modified_count += 1

            bit_idx += bits_per_sample

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

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

        debug.print(f"Modified {modified_count} samples (out of {samples_needed} selected)")

        # 7. Reshape and write back as WAV
        stego_samples = flat_samples.reshape(original_shape)

        output_buf = io.BytesIO()
        sf.write(output_buf, stego_samples, samplerate, format="WAV", subtype="PCM_16")
        output_buf.seek(0)
        stego_bytes = output_buf.getvalue()

        stats = AudioEmbedStats(
            samples_modified=modified_count,
            total_samples=num_samples,
            capacity_used=len(payload) / max_bytes,
            bytes_embedded=len(payload),
            sample_rate=samplerate,
            channels=channels,
            duration_seconds=duration,
            embed_mode=EMBED_MODE_AUDIO_LSB,
        )

        debug.print(f"Audio LSB embedding complete: {len(stego_bytes)} byte WAV")
        return stego_bytes, stats

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


# =============================================================================
# EXTRACTION
# =============================================================================


@debug.time
def extract_from_audio_lsb(
    audio_data: bytes,
    sample_key: bytes,
    bits_per_sample: int = 1,
    progress_file: str | None = None,
) -> bytes | None:
    """
    Extract hidden data from audio using LSB steganography.

    Reads the stego audio, generates the same pseudo-random sample indices
    from ``sample_key``, extracts the LSBs, and reconstructs the payload.
    Verifies the ``AUDIO_MAGIC_LSB`` header before returning.

    Args:
        audio_data: Raw bytes of the stego WAV file.
        sample_key: 32-byte key (must match the one used for embedding).
        bits_per_sample: LSBs per sample (must match embedding).
        progress_file: Optional path for progress JSON.

    Returns:
        Extracted payload bytes (without magic/length prefix), or ``None``
        if extraction fails (wrong key, no data, corrupted).
    """
    debug.print(f"Audio LSB extracting from {len(audio_data)} byte audio")
    debug.data(sample_key, "Sample key for extraction")
    debug.validate(
        bits_per_sample in (1, 2), f"bits_per_sample must be 1 or 2, got {bits_per_sample}"
    )

    try:
        # 1. Read audio
        samples, samplerate = sf.read(io.BytesIO(audio_data), dtype="int16", always_2d=True)
        flat_samples = samples.flatten()
        num_samples = len(flat_samples)

        debug.print(f"Audio: {samplerate} Hz, {samples.shape[1]} ch, {num_samples} total samples")

        # 2. Extract initial samples to find magic bytes + length (8 bytes = 64 bits)
        header_bits_needed = 64  # 4 bytes magic + 4 bytes length
        header_samples_needed = (header_bits_needed + bits_per_sample - 1) // bits_per_sample + 10

        if header_samples_needed > num_samples:
            debug.print("Audio too small to contain header")
            return None

        initial_indices = generate_sample_indices(sample_key, num_samples, header_samples_needed)

        binary_data = ""
        for sample_idx in initial_indices:
            val = int(flat_samples[sample_idx]) & 0xFFFF
            for bit_pos in range(bits_per_sample - 1, -1, -1):
                binary_data += str((val >> bit_pos) & 1)

        # 3. Verify magic bytes
        if len(binary_data) < 64:
            debug.print(f"Not enough bits for header: {len(binary_data)}/64")
            return None

        magic_bits = binary_data[:32]
        magic_bytes = int(magic_bits, 2).to_bytes(4, "big")

        if magic_bytes != AUDIO_MAGIC_LSB:
            debug.print(f"Magic mismatch: got {magic_bytes!r}, expected {AUDIO_MAGIC_LSB!r}")
            return None

        debug.print("Magic bytes verified: AUDL")

        # 4. Parse length
        length_bits = binary_data[32:64]
        data_length = struct.unpack(">I", int(length_bits, 2).to_bytes(4, "big"))[0]
        debug.print(f"Extracted length: {data_length} bytes")

        # Sanity check length
        max_possible = (num_samples * bits_per_sample) // 8 - 8  # minus header
        if data_length > max_possible or data_length < 1:
            debug.print(f"Invalid data length: {data_length} (max possible: {max_possible})")
            return None

        # 5. Extract full payload
        total_bits = (8 + data_length) * 8  # header (8 bytes) + payload
        total_samples_needed = (total_bits + bits_per_sample - 1) // bits_per_sample

        if total_samples_needed > num_samples:
            debug.print(f"Need {total_samples_needed} samples but only {num_samples} available")
            return None

        debug.print(f"Need {total_samples_needed} samples to extract {data_length} bytes")

        selected_indices = generate_sample_indices(sample_key, num_samples, total_samples_needed)

        # Initial progress
        if progress_file:
            _write_progress(progress_file, 5, 100, "extracting")

        binary_data = ""
        for progress_idx, sample_idx in enumerate(selected_indices):
            val = int(flat_samples[sample_idx]) & 0xFFFF
            for bit_pos in range(bits_per_sample - 1, -1, -1):
                binary_data += str((val >> bit_pos) & 1)

            if progress_file and progress_idx % PROGRESS_INTERVAL == 0:
                _write_progress(progress_file, progress_idx, total_samples_needed, "extracting")

        if progress_file:
            _write_progress(progress_file, total_samples_needed, total_samples_needed, "extracting")

        # Skip the 8-byte header (magic + length) = 64 bits
        data_bits = binary_data[64 : 64 + (data_length * 8)]

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

        # Convert bits back to bytes
        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"Audio LSB successfully extracted {len(data_bytes)} bytes")
        return bytes(data_bytes)

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