use anyhow::{anyhow, Result};
use image::ImageFormat;
use rayon::prelude::*;

use crate::mcap_reader::{FrameData, VideoFrame};

/// A decoded frame ready for encoding
#[derive(Debug, Clone)]
pub struct DecodedFrame {
    pub timestamp: f64,
    pub sequence: u64,
    pub width: u32,
    pub height: u32,
    pub rgb_data: Vec<u8>, // RGB24 format
}

/// Outcome of decoding a single frame.
///
/// `Ok` carries a fully decoded frame ready for the encoder.
/// `Bad` carries enough metadata to log the failure and count it toward the
/// bad-frame threshold without aborting the export.
#[derive(Debug)]
pub enum DecodeOutcome {
    Ok(DecodedFrame),
    Bad(BadFrame),
}

/// Metadata describing a frame that failed to decode.
#[derive(Debug, Clone)]
pub struct BadFrame {
    pub sequence: u64,
    pub timestamp: f64,
    pub size_bytes: usize,
    pub reason: String,
}

impl BadFrame {
    /// Truncate a decoder error message to a single ≤120-char line suitable for logging.
    fn sanitize_reason(raw: &str) -> String {
        let one_line = raw.replace(['\n', '\r'], " ");
        if one_line.chars().count() <= 120 {
            one_line
        } else {
            one_line.chars().take(120).collect()
        }
    }
}

/// Return the size in bytes of the underlying payload for a frame
/// (compressed bytes for compressed frames, raw row data for raw frames).
fn payload_size(frame: &VideoFrame) -> usize {
    match &frame.data {
        FrameData::Compressed { data, .. } => data.len(),
        FrameData::Raw { data, .. } => data.len(),
        FrameData::CompressedVideoPacket { data, .. } => data.len(),
    }
}

/// Decode a single frame to RGB24 format
pub fn decode_frame(frame: &VideoFrame) -> Result<DecodedFrame> {
    match &frame.data {
        FrameData::Compressed { format, data } => decode_compressed(frame, format, data),
        FrameData::CompressedVideoPacket { .. } => Err(anyhow!(
            "CompressedVideo packets cannot be decoded as still images; use the packet export path"
        )),
        FrameData::Raw {
            width,
            height,
            encoding,
            step,
            data,
        } => decode_raw(frame, *width, *height, *step, encoding, data),
    }
}

/// Decode compressed image (JPEG/PNG)
fn decode_compressed(frame: &VideoFrame, format: &str, data: &[u8]) -> Result<DecodedFrame> {
    // Determine image format from the format string
    let img_format = if format.contains("jpeg") || format.contains("jpg") {
        Some(ImageFormat::Jpeg)
    } else if format.contains("png") {
        Some(ImageFormat::Png)
    } else {
        None
    };

    let img = if let Some(fmt) = img_format {
        image::load_from_memory_with_format(data, fmt)?
    } else {
        // Try to guess format from data
        image::load_from_memory(data)?
    };

    let rgb = img.to_rgb8();
    let (width, height) = rgb.dimensions();

    Ok(DecodedFrame {
        timestamp: frame.timestamp,
        sequence: frame.sequence,
        width,
        height,
        rgb_data: rgb.into_raw(),
    })
}

/// Get bytes per pixel for an encoding
fn bytes_per_pixel(encoding: &str) -> Option<usize> {
    match encoding {
        "rgb8" | "RGB8" | "bgr8" | "BGR8" | "8UC3" => Some(3),
        "rgba8" | "RGBA8" | "bgra8" | "BGRA8" | "8UC4" => Some(4),
        "mono8" | "MONO8" | "8UC1" => Some(1),
        "16UC1" | "mono16" => Some(2),
        _ => None,
    }
}

/// Decode raw image data based on encoding
fn decode_raw(
    frame: &VideoFrame,
    width: u32,
    height: u32,
    step: u32,
    encoding: &str,
    data: &[u8],
) -> Result<DecodedFrame> {
    let bpp = bytes_per_pixel(encoding)
        .ok_or_else(|| anyhow!("Unsupported image encoding: {}", encoding))?;

    let row_size = (width as usize) * bpp;

    // Use step from message if valid, otherwise calculate from data
    let actual_stride = if step > 0 {
        step as usize
    } else if data.len() > 0 && height > 0 {
        // Calculate stride from data size
        data.len() / (height as usize)
    } else {
        row_size
    };

    // Validate we have enough data
    let min_data_size = actual_stride * (height as usize);
    if data.len() < min_data_size {
        return Err(anyhow!(
            "Insufficient image data: got {} bytes, need {} for {}x{} {} (stride: {})",
            data.len(),
            min_data_size,
            width,
            height,
            encoding,
            actual_stride
        ));
    }

    let output_size = (width as usize) * (height as usize) * 3;
    let mut rgb_data = Vec::with_capacity(output_size);

    // Process row by row to handle stride
    for row in 0..(height as usize) {
        let row_start = row * actual_stride;
        let row_data = &data[row_start..row_start + row_size];

        match encoding {
            "rgb8" | "RGB8" | "8UC3" => {
                // Already RGB, just copy
                rgb_data.extend_from_slice(row_data);
            }
            "bgr8" | "BGR8" => {
                // Convert BGR to RGB
                for chunk in row_data.chunks_exact(3) {
                    rgb_data.push(chunk[2]); // R
                    rgb_data.push(chunk[1]); // G
                    rgb_data.push(chunk[0]); // B
                }
            }
            "rgba8" | "RGBA8" | "8UC4" => {
                // Convert RGBA to RGB (drop alpha)
                for chunk in row_data.chunks_exact(4) {
                    rgb_data.push(chunk[0]); // R
                    rgb_data.push(chunk[1]); // G
                    rgb_data.push(chunk[2]); // B
                }
            }
            "bgra8" | "BGRA8" => {
                // Convert BGRA to RGB
                for chunk in row_data.chunks_exact(4) {
                    rgb_data.push(chunk[2]); // R
                    rgb_data.push(chunk[1]); // G
                    rgb_data.push(chunk[0]); // B
                }
            }
            "mono8" | "MONO8" | "8UC1" => {
                // Convert grayscale to RGB
                for &pixel in row_data {
                    rgb_data.push(pixel);
                    rgb_data.push(pixel);
                    rgb_data.push(pixel);
                }
            }
            "16UC1" | "mono16" => {
                // Convert 16-bit grayscale to 8-bit RGB
                for chunk in row_data.chunks_exact(2) {
                    // Take high byte for 8-bit representation
                    let pixel = chunk[1];
                    rgb_data.push(pixel);
                    rgb_data.push(pixel);
                    rgb_data.push(pixel);
                }
            }
            _ => unreachable!(),
        }
    }

    // Final validation
    if rgb_data.len() != output_size {
        return Err(anyhow!(
            "RGB conversion error: got {} bytes, expected {}",
            rgb_data.len(),
            output_size
        ));
    }

    Ok(DecodedFrame {
        timestamp: frame.timestamp,
        sequence: frame.sequence,
        width,
        height,
        rgb_data,
    })
}

/// Decode multiple frames in parallel using rayon.
///
/// Returns one `DecodeOutcome` per input frame, preserving sequence order.
/// Decode errors and panics inside the `image` crate are captured as
/// `DecodeOutcome::Bad` so a single corrupt frame cannot abort the whole export.
pub fn decode_frames_parallel(frames: &[VideoFrame]) -> Vec<DecodeOutcome> {
    let mut outcomes: Vec<DecodeOutcome> = frames
        .par_iter()
        .map(|f| {
            let seq = f.sequence;
            let ts = f.timestamp;
            let size = payload_size(f);

            let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| decode_frame(f)));

            match caught {
                Ok(Ok(decoded)) => DecodeOutcome::Ok(decoded),
                Ok(Err(err)) => DecodeOutcome::Bad(BadFrame {
                    sequence: seq,
                    timestamp: ts,
                    size_bytes: size,
                    reason: BadFrame::sanitize_reason(&err.to_string()),
                }),
                Err(panic) => {
                    let msg = if let Some(s) = panic.downcast_ref::<&str>() {
                        format!("panic: {}", s)
                    } else if let Some(s) = panic.downcast_ref::<String>() {
                        format!("panic: {}", s)
                    } else {
                        "panic: (non-string payload)".to_string()
                    };
                    DecodeOutcome::Bad(BadFrame {
                        sequence: seq,
                        timestamp: ts,
                        size_bytes: size,
                        reason: BadFrame::sanitize_reason(&msg),
                    })
                }
            }
        })
        .collect();

    outcomes.sort_by_key(|o| match o {
        DecodeOutcome::Ok(f) => f.sequence,
        DecodeOutcome::Bad(b) => b.sequence,
    });

    outcomes
}

#[cfg(test)]
mod sanitize_tests {
    use super::BadFrame;

    #[test]
    fn sanitize_reason_is_utf8_safe_on_multibyte_boundary() {
        // Construct a 200-char string of 2-byte chars so byte 120 falls mid-char.
        let input: String = std::iter::repeat('ß').take(200).collect();
        let out = BadFrame::sanitize_reason(&input);
        // Must not panic. Must be valid UTF-8 (guaranteed by String type).
        assert!(out.chars().count() <= 120);
        assert!(out.chars().all(|c| c == 'ß'));
    }

    #[test]
    fn sanitize_reason_replaces_newlines() {
        let input = "line1\nline2\rline3";
        let out = BadFrame::sanitize_reason(input);
        assert!(!out.contains('\n'));
        assert!(!out.contains('\r'));
        assert!(out.contains("line1 line2 line3"));
    }

    #[test]
    fn sanitize_reason_short_input_is_unchanged() {
        let input = "short error";
        let out = BadFrame::sanitize_reason(input);
        assert_eq!(out, "short error");
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcap_reader::{FrameData, VideoFrame};

    fn raw_frame(seq: u64, width: u32, height: u32, data: Vec<u8>) -> VideoFrame {
        VideoFrame {
            timestamp: seq as f64,
            sequence: seq,
            data: FrameData::Raw {
                width,
                height,
                encoding: "rgb8".to_string(),
                step: width * 3,
                data,
            },
        }
    }

    fn compressed_frame(seq: u64, data: Vec<u8>) -> VideoFrame {
        VideoFrame {
            timestamp: seq as f64,
            sequence: seq,
            data: FrameData::Compressed {
                format: "jpeg".to_string(),
                data,
            },
        }
    }

    #[test]
    fn decode_frames_parallel_returns_outcomes_for_mixed_input() {
        // Good: 2x2 rgb8 = 12 bytes
        let good = raw_frame(0, 2, 2, vec![0u8; 12]);
        // Bad: truncated rgb8 (need 12, got 3)
        let truncated = raw_frame(1, 2, 2, vec![0u8; 3]);
        // Bad: garbage JPEG
        let garbage = compressed_frame(2, vec![0xFF, 0xD8, 0xFF, 0xE0, 0xDE, 0xAD]);

        let outcomes = decode_frames_parallel(&[good, truncated, garbage]);

        assert_eq!(outcomes.len(), 3);
        match &outcomes[0] {
            DecodeOutcome::Ok(f) => {
                assert_eq!(f.sequence, 0);
                assert_eq!(f.width, 2);
                assert_eq!(f.height, 2);
            }
            DecodeOutcome::Bad(b) => panic!("expected Ok, got Bad: {}", b.reason),
        }
        match &outcomes[1] {
            DecodeOutcome::Bad(b) => {
                assert_eq!(b.sequence, 1);
                assert_eq!(b.size_bytes, 3);
                assert!(
                    b.reason.to_lowercase().contains("insufficient")
                        || b.reason.to_lowercase().contains("image data"),
                    "unexpected reason: {}",
                    b.reason
                );
            }
            DecodeOutcome::Ok(_) => panic!("expected Bad for truncated raw frame"),
        }
        match &outcomes[2] {
            DecodeOutcome::Bad(b) => {
                assert_eq!(b.sequence, 2);
                assert_eq!(b.size_bytes, 6);
                assert!(!b.reason.is_empty());
            }
            DecodeOutcome::Ok(_) => panic!("expected Bad for garbage JPEG"),
        }
    }

    #[test]
    fn decode_frames_parallel_sorts_outcomes_by_sequence() {
        let frames = vec![
            raw_frame(2, 1, 1, vec![0u8; 3]),
            raw_frame(0, 1, 1, vec![0u8; 3]),
            raw_frame(1, 1, 1, vec![0u8; 0]), // bad
        ];
        let outcomes = decode_frames_parallel(&frames);
        assert_eq!(outcomes.len(), 3);
        let seqs: Vec<u64> = outcomes
            .iter()
            .map(|o| match o {
                DecodeOutcome::Ok(f) => f.sequence,
                DecodeOutcome::Bad(b) => b.sequence,
            })
            .collect();
        assert_eq!(seqs, vec![0, 1, 2]);
    }
}
