"""
Stegasoo Video Steganography — LSB Embedding/Extraction (v4.4.0)

Frame-based LSB embedding for video files.

Hides data in the least significant bits of video frame pixels. Uses the
existing image steganography engine for per-frame embedding, providing
high capacity across multiple I-frames.

Strategy:
1. Extract I-frames (keyframes) from video using ffmpeg
2. Embed payload across I-frames using existing LSB engine
3. Re-encode video with modified frames using FFV1 lossless codec
4. Output: MKV container with embedded data

Uses ChaCha20 as a CSPRNG for pseudo-random frame selection and pixel
selection within frames, ensuring that without the key an attacker cannot
determine which frames/pixels were modified.
"""

import struct
import tempfile
from pathlib import Path

from .constants import (
    EMBED_MODE_VIDEO_LSB,
    VIDEO_MAGIC_LSB,
    VIDEO_OUTPUT_CODEC,
)
from .debug import debug
from .exceptions import VideoCapacityError, VideoError
from .models import VideoEmbedStats
from .steganography import ENCRYPTION_OVERHEAD, _embed_lsb, _extract_lsb
from .video_utils import extract_frames, get_video_info, reassemble_video

# Progress reporting interval — write every N frames
PROGRESS_INTERVAL = 5


# =============================================================================
# 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_video_lsb_capacity(video_data: bytes) -> int:
    """
    Calculate the maximum bytes that can be embedded in a video via LSB.

    Calculates capacity based on I-frames (keyframes) only. Each I-frame
    provides capacity proportional to its pixel count.

    Args:
        video_data: Raw bytes of a video file.

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

    Raises:
        VideoError: If the video cannot be read or is in an unsupported format.
    """
    from .video_utils import calculate_video_capacity

    capacity_info = calculate_video_capacity(video_data, EMBED_MODE_VIDEO_LSB)

    debug.print(
        f"Video LSB capacity: {capacity_info.usable_capacity_bytes} bytes "
        f"({capacity_info.i_frames} I-frames, {capacity_info.resolution[0]}x{capacity_info.resolution[1]})"
    )

    return capacity_info.usable_capacity_bytes


# =============================================================================
# FRAME INDEX GENERATION (ChaCha20 CSPRNG)
# =============================================================================


def generate_frame_indices(key: bytes, num_frames: int, num_needed: int) -> list[int]:
    """
    Generate pseudo-random frame indices using ChaCha20 as a CSPRNG.

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

    Args:
        key: 32-byte key for the ChaCha20 cipher.
        num_frames: Total number of frames available.
        num_needed: How many unique frame indices are required.

    Returns:
        List of ``num_needed`` unique indices in [0, num_frames).
    """
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms

    debug.validate(len(key) == 32, f"Frame key must be 32 bytes, got {len(key)}")
    debug.validate(num_frames > 0, f"Number of frames must be positive, got {num_frames}")
    debug.validate(num_needed > 0, f"Number needed must be positive, got {num_needed}")
    debug.validate(
        num_needed <= num_frames,
        f"Cannot select {num_needed} frames from {num_frames} available",
    )

    debug.print(f"Generating {num_needed} frame indices from {num_frames} total frames")

    # Use a different nonce offset for frame selection (vs pixel selection)
    nonce = b"\x01" + b"\x00" * 15  # Different from pixel selection nonce

    if num_needed >= num_frames // 2:
        # Full Fisher-Yates shuffle
        cipher = Cipher(algorithms.ChaCha20(key, nonce), mode=None, backend=default_backend())
        encryptor = cipher.encryptor()

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

        for i in range(num_frames - 1, 0, -1):
            j_bytes = random_bytes[(num_frames - 1 - i) * 4 : (num_frames - i) * 4]
            j = int.from_bytes(j_bytes, "big") % (i + 1)
            indices[i], indices[j] = indices[j], indices[i]

        return indices[:num_needed]

    # Direct sampling
    selected: list[int] = []
    used: set[int] = set()

    cipher = Cipher(algorithms.ChaCha20(key, nonce), mode=None, backend=default_backend())
    encryptor = cipher.encryptor()

    bytes_needed = (num_needed * 2) * 4
    random_bytes = encryptor.update(b"\x00" * bytes_needed)

    byte_offset = 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_frames
        byte_offset += 4

        if idx not in used:
            used.add(idx)
            selected.append(idx)

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


# =============================================================================
# EMBEDDING
# =============================================================================


@debug.time
def embed_in_video_lsb(
    data: bytes,
    carrier_video: bytes,
    pixel_key: bytes,
    progress_file: str | None = None,
) -> tuple[bytes, VideoEmbedStats]:
    """
    Embed data into video frames using LSB steganography.

    The payload is prepended with a 4-byte magic header and a 4-byte
    big-endian length prefix. Data is distributed across I-frames using
    pseudo-random selection based on the pixel_key.

    The output video uses FFV1 lossless codec in MKV container to
    preserve the embedded data perfectly.

    Args:
        data: Encrypted payload bytes to embed.
        carrier_video: Raw bytes of the carrier video file.
        pixel_key: 32-byte key for frame and pixel selection.
        progress_file: Optional path for progress JSON (frontend polling).

    Returns:
        Tuple of (stego video bytes, VideoEmbedStats).

    Raises:
        VideoCapacityError: If the payload is too large for the carrier.
        VideoError: On any other embedding failure.
    """
    debug.print(f"Video LSB embedding {len(data)} bytes")
    debug.data(pixel_key, "Pixel key for embedding")
    debug.validate(len(pixel_key) == 32, f"Pixel key must be 32 bytes, got {len(pixel_key)}")

    try:
        # Get video info
        video_info = get_video_info(carrier_video)
        debug.print(
            f"Carrier video: {video_info.width}x{video_info.height}, "
            f"{video_info.fps:.2f} fps, {video_info.duration_seconds:.1f}s, "
            f"{video_info.i_frame_count} I-frames"
        )

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

        # Calculate capacity and check fit
        capacity = calculate_video_lsb_capacity(carrier_video)
        if len(payload) > capacity + ENCRYPTION_OVERHEAD:
            raise VideoCapacityError(len(payload), capacity)

        # Extract I-frames to temp directory
        with tempfile.TemporaryDirectory(prefix="stegasoo_video_") as temp_dir_str:
            temp_dir = Path(temp_dir_str)

            _write_progress(progress_file, 5, 100, "extracting_frames")

            frames, _ = extract_frames(carrier_video, temp_dir, keyframes_only=True)
            num_frames = len(frames)

            debug.print(f"Extracted {num_frames} I-frames for embedding")

            if num_frames == 0:
                raise VideoError("No I-frames found in video")

            # Calculate bytes per frame (minus 4 byte length prefix used by _embed_lsb)
            pixels_per_frame = video_info.width * video_info.height
            bytes_per_frame = (pixels_per_frame * 3) // 8 - 4  # 3 bits per pixel, minus len prefix

            # For simplicity, embed entire payload in first frame if it fits
            # This makes extraction straightforward
            if len(payload) <= bytes_per_frame:
                debug.print(f"Payload fits in single frame ({len(payload)} <= {bytes_per_frame})")
                frame_path = frames[0]

                with open(frame_path, "rb") as f:
                    frame_data = f.read()

                try:
                    stego_frame, stats, ext = _embed_lsb(
                        payload,
                        frame_data,
                        pixel_key,
                        bits_per_channel=1,
                        output_format="PNG",
                    )

                    with open(frame_path, "wb") as f:
                        f.write(stego_frame)

                    modified_frames = 1

                except Exception as e:
                    debug.print(f"Failed to embed in frame: {e}")
                    raise VideoError(f"Failed to embed in frame: {e}")
            else:
                # For larger payloads, we need to split across frames
                # Each frame stores: 4-byte chunk length + chunk data
                debug.print("Splitting payload across multiple frames")

                frames_needed = (len(payload) + bytes_per_frame - 1) // bytes_per_frame
                frames_needed = min(frames_needed, num_frames)

                debug.print(f"Using {frames_needed} frames to embed {len(payload)} bytes")

                # For now, use sequential frames for simplicity
                modified_frames = 0
                bytes_remaining = len(payload)
                payload_offset = 0

                for frame_idx in range(frames_needed):
                    if bytes_remaining <= 0:
                        break

                    frame_path = frames[frame_idx]

                    with open(frame_path, "rb") as f:
                        frame_data = f.read()

                    chunk_size = min(bytes_remaining, bytes_per_frame)
                    chunk = payload[payload_offset : payload_offset + chunk_size]

                    try:
                        stego_frame, stats, ext = _embed_lsb(
                            chunk,
                            frame_data,
                            pixel_key,
                            bits_per_channel=1,
                            output_format="PNG",
                        )

                        with open(frame_path, "wb") as f:
                            f.write(stego_frame)

                        modified_frames += 1
                        payload_offset += chunk_size
                        bytes_remaining -= chunk_size

                    except Exception as e:
                        debug.print(f"Failed to embed in frame {frame_idx}: {e}")
                        raise VideoError(f"Failed to embed in frame {frame_idx}: {e}")

                    if progress_file and frame_idx % PROGRESS_INTERVAL == 0:
                        pct = 10 + int((frame_idx / frames_needed) * 70)
                        _write_progress(progress_file, pct, 100, "embedding")

            _write_progress(progress_file, 80, 100, "reassembling")

            # Reassemble video with modified frames
            stego_video = reassemble_video(
                frames,
                carrier_video,
                fps=1.0,  # I-frame only videos use 1 fps
            )

            _write_progress(progress_file, 100, 100, "complete")

            video_stats = VideoEmbedStats(
                frames_modified=modified_frames,
                total_frames=video_info.total_frames,
                capacity_used=len(payload) / (capacity + ENCRYPTION_OVERHEAD),
                bytes_embedded=len(payload),
                width=video_info.width,
                height=video_info.height,
                fps=video_info.fps,
                duration_seconds=video_info.duration_seconds,
                embed_mode=EMBED_MODE_VIDEO_LSB,
                codec=VIDEO_OUTPUT_CODEC,
            )

            debug.print(
                f"Video LSB embedding complete: {len(stego_video)} bytes, "
                f"{modified_frames} frames modified"
            )

            return stego_video, video_stats

    except VideoCapacityError:
        raise
    except VideoError:
        raise
    except Exception as e:
        debug.exception(e, "embed_in_video_lsb")
        raise VideoError(f"Failed to embed data in video: {e}") from e


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


@debug.time
def extract_from_video_lsb(
    video_data: bytes,
    pixel_key: bytes,
    progress_file: str | None = None,
) -> bytes | None:
    """
    Extract hidden data from video using LSB steganography.

    Extracts I-frames, reads LSBs from the same pseudo-random locations
    used during embedding, and reconstructs the payload.

    Args:
        video_data: Raw bytes of the stego video file.
        pixel_key: 32-byte key (must match the one used for 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"Video LSB extracting from {len(video_data)} byte video")
    debug.data(pixel_key, "Pixel key for extraction")

    try:
        # Get video info
        video_info = get_video_info(video_data)
        debug.print(
            f"Video: {video_info.width}x{video_info.height}, "
            f"{video_info.i_frame_count} I-frames"
        )

        # Extract I-frames
        with tempfile.TemporaryDirectory(prefix="stegasoo_video_extract_") as temp_dir_str:
            temp_dir = Path(temp_dir_str)

            _write_progress(progress_file, 5, 100, "extracting_frames")

            frames, _ = extract_frames(video_data, temp_dir, keyframes_only=True)
            num_frames = len(frames)

            if num_frames == 0:
                debug.print("No I-frames found in video")
                return None

            debug.print(f"Extracted {num_frames} I-frames for extraction")

            _write_progress(progress_file, 20, 100, "extracting_data")

            # First, try to extract from frame 0 to get magic and total length
            frame_path = frames[0]
            with open(frame_path, "rb") as f:
                frame_data = f.read()

            first_chunk = _extract_lsb(frame_data, pixel_key, bits_per_channel=1)
            if first_chunk is None or len(first_chunk) < 8:
                debug.print("Failed to extract initial data from first frame")
                return None

            # Check magic bytes
            magic = first_chunk[:4]
            if magic != VIDEO_MAGIC_LSB:
                debug.print(f"Magic mismatch: got {magic!r}, expected {VIDEO_MAGIC_LSB!r}")
                return None

            # Get total payload length
            total_length = struct.unpack(">I", first_chunk[4:8])[0]
            debug.print(f"Total payload length: {total_length} bytes")

            # Sanity check
            pixels_per_frame = video_info.width * video_info.height
            bytes_per_frame = (pixels_per_frame * 3) // 8 - 4  # minus length prefix
            max_possible = bytes_per_frame * num_frames

            if total_length > max_possible or total_length < 1:
                debug.print(f"Invalid payload length: {total_length}")
                return None

            # If the entire payload fits in the first frame, return it directly
            # This matches the simplified single-frame embedding approach
            if len(first_chunk) >= 8 + total_length:
                debug.print("Payload fits in single frame, extracting directly")
                payload = first_chunk[8 : 8 + total_length]
            else:
                # Multi-frame extraction
                debug.print("Multi-frame extraction needed")
                frames_needed = (total_length + 8 + bytes_per_frame - 1) // bytes_per_frame
                frames_needed = min(frames_needed, num_frames)

                # Extract sequentially (matching the embedding approach)
                extracted_chunks = [first_chunk]
                for frame_idx in range(1, frames_needed):
                    frame_path = frames[frame_idx]
                    with open(frame_path, "rb") as f:
                        frame_data = f.read()

                    chunk = _extract_lsb(frame_data, pixel_key, bits_per_channel=1)
                    if chunk:
                        extracted_chunks.append(chunk)

                    if progress_file and frame_idx % PROGRESS_INTERVAL == 0:
                        pct = 20 + int((frame_idx / frames_needed) * 70)
                        _write_progress(progress_file, pct, 100, "extracting_data")

                # Combine chunks
                combined = b"".join(extracted_chunks)

                if len(combined) < 8 + total_length:
                    debug.print(
                        f"Insufficient data: have {len(combined) - 8}, need {total_length}"
                    )
                    return None

                payload = combined[8 : 8 + total_length]

            _write_progress(progress_file, 100, 100, "complete")

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

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