//! S3 storage client for image uploads.
//! Delegates S3 operations to the shared `s3_storage` crate.

use uuid::Uuid;

use crate::config::S3Config;

/// S3 client wrapper for image storage.
#[derive(Clone)]
pub struct S3Storage {
    inner: s3_storage::S3Client,
}

/// Maximum image size: 5 MB.
pub const MAX_IMAGE_SIZE: usize = 5 * 1024 * 1024;

/// Maximum decoded image dimensions (pixels). The 5 MB byte cap does NOT bound
/// the decoded pixel count, a lossless PNG/WebP of a solid color compresses a
/// gigapixel canvas into a few KB (a decompression / "pixel bomb"). We never
/// decode server-side, but a viewer's browser would, so we reject absurd
/// declared dimensions up front. 50 MP comfortably clears any real photo
/// (an 8000x6000 shot is 48 MP).
pub const MAX_IMAGE_PIXELS: u64 = 50_000_000;

/// Allowed image content types.
const ALLOWED_CONTENT_TYPES: &[&str] = &["image/png", "image/jpeg", "image/gif", "image/webp"];

/// Allowed file extensions (lowercase).
#[cfg(test)]
const ALLOWED_EXTENSIONS: &[&str] = &["png", "jpg", "jpeg", "gif", "webp"];

impl S3Storage {
    /// Create a new S3 client from configuration.
    pub async fn new(config: &S3Config) -> Result<Self, String> {
        let s3_config = s3_storage::S3Config {
            endpoint: config.endpoint.clone(),
            bucket: config.bucket.clone(),
            access_key: config.access_key.clone(),
            secret_key: config.secret_key.clone(),
            region: config.region.clone(),
        };

        let inner = s3_storage::S3Client::new(&s3_config).await?;
        Ok(Self { inner })
    }

    /// Upload bytes to S3.
    #[tracing::instrument(skip_all)]
    pub async fn upload(
        &self,
        s3_key: &str,
        content_type: &str,
        data: Vec<u8>,
    ) -> Result<(), String> {
        self.inner.upload(s3_key, content_type, data, None).await
    }

    /// Download bytes from S3.
    #[tracing::instrument(skip_all)]
    pub async fn download(&self, s3_key: &str) -> Result<(Vec<u8>, String), String> {
        self.inner.download(s3_key).await
    }

    /// Open a streaming download from S3 as an HTTP response body, without
    /// buffering the object in memory.
    ///
    /// The image-serve handler streams the result to the client chunk-by-chunk
    /// instead of materialising a full `Vec<u8>` per request (ultra-fuzz S1). The
    /// aws `ByteStream` is adapted to a body via its tokio `AsyncBufRead` reader,
    /// keeping the aws/tokio-util plumbing contained here.
    #[tracing::instrument(skip_all)]
    pub async fn download_stream(&self, s3_key: &str) -> Result<axum::body::Body, String> {
        let stream = self.inner.download_stream(s3_key).await?;
        let reader = stream.into_async_read();
        Ok(axum::body::Body::from_stream(
            tokio_util::io::ReaderStream::new(reader),
        ))
    }

    /// Delete an object from S3.
    #[tracing::instrument(skip_all)]
    pub async fn delete(&self, s3_key: &str) -> Result<(), String> {
        self.inner.delete(s3_key).await
    }

    /// Batch-delete objects in a single request (S3 allows up to 1000 keys).
    /// Returns the keys that failed to delete, paired with the error message;
    /// an empty vec means every key was deleted (or was already absent).
    #[tracing::instrument(skip_all)]
    pub async fn delete_objects(
        &self,
        s3_keys: &[String],
    ) -> Result<Vec<(String, String)>, String> {
        self.inner.delete_objects(s3_keys).await
    }
}

/// Generate an S3 key for a forum image.
/// Format: `mt/{community_slug}/{uuid}.{ext}`
pub fn generate_image_key(community_slug: &str, ext: &str) -> String {
    let id = Uuid::new_v4();
    format!("mt/{community_slug}/{id}.{ext}")
}

/// Sniff the real image format from leading magic bytes, independent of the
/// client-declared filename or Content-Type. Returns the canonical extension
/// family (`png`/`jpg`/`gif`/`webp`) or `None` if the bytes are not one of the
/// four allowed image formats. This is the only attacker-uncontrollable signal
/// in the upload, the filename extension and the multipart Content-Type are
/// both client-supplied, so a content check that only cross-references those
/// two is trivially satisfied by a spoofed file.
pub fn sniff_image_format(data: &[u8]) -> Option<&'static str> {
    // PNG: 89 50 4E 47 0D 0A 1A 0A
    if data.starts_with(&[0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]) {
        return Some("png");
    }
    // JPEG: starts with FF D8 FF
    if data.starts_with(&[0xFF, 0xD8, 0xFF]) {
        return Some("jpg");
    }
    // GIF: "GIF87a" or "GIF89a"
    if data.starts_with(b"GIF87a") || data.starts_with(b"GIF89a") {
        return Some("gif");
    }
    // WebP: "RIFF" .... "WEBP" (bytes 0-3 = RIFF, 8-11 = WEBP)
    if data.len() >= 12 && &data[0..4] == b"RIFF" && &data[8..12] == b"WEBP" {
        return Some("webp");
    }
    None
}

/// Read the declared pixel dimensions `(width, height)` straight from the format
/// header, without decoding the image. Returns `None` if the header is too short
/// or the dimensions can't be located, callers treat that as "can't tell",
/// which is safe because the format was already authenticated by
/// [`sniff_image_format`] and a real pixel bomb carries parseable dimensions.
fn image_dimensions(format: &str, data: &[u8]) -> Option<(u32, u32)> {
    let be16 = |i: usize| -> Option<u32> {
        Some(u32::from(u16::from_be_bytes([
            *data.get(i)?,
            *data.get(i + 1)?,
        ])))
    };
    let le16 = |i: usize| -> Option<u32> {
        Some(u32::from(u16::from_le_bytes([
            *data.get(i)?,
            *data.get(i + 1)?,
        ])))
    };
    match format {
        // PNG: IHDR is the first chunk; width/height are big-endian u32 at 16/20.
        "png" => {
            let w = u32::from_be_bytes([
                *data.get(16)?,
                *data.get(17)?,
                *data.get(18)?,
                *data.get(19)?,
            ]);
            let h = u32::from_be_bytes([
                *data.get(20)?,
                *data.get(21)?,
                *data.get(22)?,
                *data.get(23)?,
            ]);
            Some((w, h))
        }
        // GIF: logical-screen width/height are little-endian u16 at 6/8.
        "gif" => Some((le16(6)?, le16(8)?)),
        // JPEG: scan segments for a Start-Of-Frame marker (C0-CF except the
        // non-SOF C4/C8/CC); height/width are big-endian u16 right after the
        // 3-byte (length + precision) preamble.
        "jpg" => {
            let mut i = 2; // skip SOI (FF D8)
            while i + 9 < data.len() {
                if data[i] != 0xFF {
                    i += 1;
                    continue;
                }
                let marker = data[i + 1];
                if (0xC0..=0xCF).contains(&marker)
                    && marker != 0xC4
                    && marker != 0xC8
                    && marker != 0xCC
                {
                    return Some((be16(i + 7)?, be16(i + 5)?));
                }
                // Standalone markers (RSTn, SOI, EOI, TEM) carry no length.
                if marker == 0xD8
                    || marker == 0xD9
                    || (0xD0..=0xD7).contains(&marker)
                    || marker == 0x01
                {
                    i += 2;
                    continue;
                }
                let seg_len = be16(i + 2)? as usize;
                if seg_len < 2 {
                    return None;
                }
                i += 2 + seg_len;
            }
            None
        }
        // WebP: three sub-chunk layouts after the "WEBP" tag at byte 12.
        "webp" => match data.get(12..16)? {
            b"VP8X" => {
                // Canvas width/height minus one, 24-bit little-endian at 24/27.
                let w = 1
                    + (u32::from(*data.get(24)?)
                        | u32::from(*data.get(25)?) << 8
                        | u32::from(*data.get(26)?) << 16);
                let h = 1
                    + (u32::from(*data.get(27)?)
                        | u32::from(*data.get(28)?) << 8
                        | u32::from(*data.get(29)?) << 16);
                Some((w, h))
            }
            b"VP8L" => {
                // 14-bit width/height minus one, packed after the 0x2F signature.
                let b = data.get(21..25)?;
                let bits = u32::from(b[0])
                    | u32::from(b[1]) << 8
                    | u32::from(b[2]) << 16
                    | u32::from(b[3]) << 24;
                Some((1 + (bits & 0x3FFF), 1 + ((bits >> 14) & 0x3FFF)))
            }
            b"VP8 " => {
                // Lossy: 14-bit width/height little-endian at 26/28 (after the
                // 3-byte start code 9D 01 2A).
                Some((le16(26)? & 0x3FFF, le16(28)? & 0x3FFF))
            }
            _ => None,
        },
        _ => None,
    }
}

/// Validate an uploaded file. Returns the sanitized extension and content type.
///
/// Defence in depth: the filename extension and the multipart Content-Type must
/// agree AND the actual file bytes must sniff to the same image format. A
/// polyglot or content-spoofed file (HTML/SVG bytes named `evil.png` with a
/// faked `image/png` Content-Type) fails the byte sniff even though it satisfies
/// the metadata cross-check.
pub fn validate_image(
    filename: &str,
    content_type: &str,
    data: &[u8],
) -> Result<(&'static str, &'static str), &'static str> {
    let size = data.len();
    if size > MAX_IMAGE_SIZE {
        return Err("Image exceeds 5 MB limit.");
    }
    if size == 0 {
        return Err("Empty file.");
    }

    let ext = filename
        .rsplit('.')
        .next()
        .map(str::to_lowercase)
        .unwrap_or_default();

    let ext_str: &'static str = match ext.as_str() {
        "png" => "png",
        "jpg" | "jpeg" => "jpg",
        "gif" => "gif",
        "webp" => "webp",
        _ => return Err("Allowed types: png, jpg, gif, webp."),
    };

    let ct: &'static str = ALLOWED_CONTENT_TYPES
        .iter()
        .find(|&&ct| ct == content_type)
        .copied()
        .ok_or("Invalid image content type.")?;

    // Cross-validate: extension should match content type
    let ext_matches = match ext_str {
        "png" => ct == "image/png",
        "jpg" => ct == "image/jpeg",
        "gif" => ct == "image/gif",
        "webp" => ct == "image/webp",
        _ => false,
    };
    if !ext_matches {
        return Err("File extension does not match content type.");
    }

    // Authoritative check: the bytes themselves must be the declared format.
    match sniff_image_format(data) {
        Some(sniffed) if sniffed == ext_str => {
            // Reject decompression / pixel bombs: a tiny file can still declare a
            // gigapixel canvas that detonates in a viewer's browser. If we can
            // read the dimensions and they exceed the cap, refuse.
            if let Some((w, h)) = image_dimensions(ext_str, data)
                && u64::from(w) * u64::from(h) > MAX_IMAGE_PIXELS
            {
                return Err("Image dimensions are too large.");
            }
            Ok((ext_str, ct))
        }
        Some(_) => Err("File contents do not match the declared image type."),
        None => Err("File is not a valid PNG, JPEG, GIF, or WebP image."),
    }
}

/// Strip EXIF metadata from JPEG data.
/// Works by copying all JPEG segments except APP1 (EXIF) and APP13 (IPTC).
/// Returns the cleaned data, or the original data if parsing fails.
pub fn strip_exif_jpeg(data: &[u8]) -> Vec<u8> {
    if data.len() < 4 || data[0] != 0xFF || data[1] != 0xD8 {
        return data.to_vec(); // Not a valid JPEG
    }

    let mut out = Vec::with_capacity(data.len());
    out.extend_from_slice(&[0xFF, 0xD8]); // SOI marker

    let mut i = 2;
    while i + 1 < data.len() {
        if data[i] != 0xFF {
            // Not a marker, copy rest as-is (image data)
            out.extend_from_slice(&data[i..]);
            break;
        }

        let marker = data[i + 1];

        // SOS (Start of Scan), copy the rest verbatim (compressed data follows)
        if marker == 0xDA {
            out.extend_from_slice(&data[i..]);
            break;
        }

        // Markers without length (RST0-RST7, SOI, EOI, TEM)
        if marker == 0x00 || marker == 0x01 || (0xD0..=0xD9).contains(&marker) {
            out.extend_from_slice(&data[i..i + 2]);
            i += 2;
            continue;
        }

        if i + 3 >= data.len() {
            out.extend_from_slice(&data[i..]);
            break;
        }
        let seg_len = ((data[i + 2] as usize) << 8) | (data[i + 3] as usize);
        let total = 2 + seg_len; // marker (2) + length includes itself

        if i + total > data.len() {
            out.extend_from_slice(&data[i..]);
            break;
        }

        // Skip APP1 (0xE1 = EXIF) and APP13 (0xED = IPTC/Photoshop)
        if marker == 0xE1 || marker == 0xED {
            i += total;
            continue;
        }

        out.extend_from_slice(&data[i..i + total]);
        i += total;
    }

    out
}

#[cfg(test)]
mod tests {
    use super::*;

    // Minimal byte fixtures carrying each format's real magic bytes.
    fn png_bytes() -> Vec<u8> {
        let mut v = vec![0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A];
        v.extend_from_slice(&[0u8; 64]);
        v
    }
    fn jpeg_bytes() -> Vec<u8> {
        let mut v = vec![0xFF, 0xD8, 0xFF, 0xE0];
        v.extend_from_slice(&[0u8; 64]);
        v
    }
    fn gif_bytes() -> Vec<u8> {
        let mut v = b"GIF89a".to_vec();
        v.extend_from_slice(&[0u8; 64]);
        v
    }
    fn webp_bytes() -> Vec<u8> {
        let mut v = b"RIFF\x00\x00\x00\x00WEBP".to_vec();
        v.extend_from_slice(&[0u8; 64]);
        v
    }

    #[test]
    fn validate_valid_png() {
        let (ext, ct) = validate_image("photo.png", "image/png", &png_bytes()).unwrap();
        assert_eq!(ext, "png");
        assert_eq!(ct, "image/png");
    }

    #[test]
    fn validate_valid_jpeg() {
        let (ext, _) = validate_image("photo.jpg", "image/jpeg", &jpeg_bytes()).unwrap();
        assert_eq!(ext, "jpg");
    }

    #[test]
    fn validate_rejects_oversized() {
        let mut big = png_bytes();
        big.resize(6 * 1024 * 1024, 0);
        let err = validate_image("big.png", "image/png", &big).unwrap_err();
        assert!(err.contains("5 MB"));
    }

    #[test]
    fn validate_rejects_bad_extension() {
        let err = validate_image("file.exe", "application/octet-stream", &png_bytes()).unwrap_err();
        assert!(err.contains("Allowed types"));
    }

    #[test]
    fn validate_rejects_mismatched_type() {
        let err = validate_image("photo.png", "image/jpeg", &png_bytes()).unwrap_err();
        assert!(err.contains("does not match"));
    }

    #[test]
    fn validate_rejects_png_pixel_bomb() {
        // A tiny, valid-looking PNG that declares a 30000x30000 (900 MP) canvas.
        let mut bomb = png_bytes();
        let w = 30_000u32.to_be_bytes();
        let h = 30_000u32.to_be_bytes();
        bomb[16..20].copy_from_slice(&w);
        bomb[20..24].copy_from_slice(&h);
        let err = validate_image("bomb.png", "image/png", &bomb).unwrap_err();
        assert!(err.contains("dimensions are too large"));
    }

    #[test]
    fn validate_accepts_reasonable_png_dimensions() {
        let mut img = png_bytes();
        img[16..20].copy_from_slice(&1920u32.to_be_bytes());
        img[20..24].copy_from_slice(&1080u32.to_be_bytes());
        assert!(validate_image("ok.png", "image/png", &img).is_ok());
    }

    #[test]
    fn validate_rejects_empty() {
        let err = validate_image("photo.png", "image/png", &[]).unwrap_err();
        assert!(err.contains("Empty"));
    }

    #[test]
    fn validate_rejects_content_spoof() {
        // Metadata says PNG and they agree, but the bytes are HTML, the classic
        // polyglot/spoof the old metadata-only check waved through.
        let html = b"<!DOCTYPE html><script>alert(1)</script>";
        let err = validate_image("evil.png", "image/png", html).unwrap_err();
        assert!(err.contains("not a valid"), "got: {err}");
    }

    #[test]
    fn validate_rejects_format_mismatch_bytes() {
        // Declared PNG (ext + content-type agree) but the bytes are a real JPEG.
        let err = validate_image("photo.png", "image/png", &jpeg_bytes()).unwrap_err();
        assert!(err.contains("do not match the declared"), "got: {err}");
    }

    #[test]
    fn sniff_detects_each_format() {
        assert_eq!(sniff_image_format(&png_bytes()), Some("png"));
        assert_eq!(sniff_image_format(&jpeg_bytes()), Some("jpg"));
        assert_eq!(sniff_image_format(&gif_bytes()), Some("gif"));
        assert_eq!(sniff_image_format(&webp_bytes()), Some("webp"));
    }

    #[test]
    fn sniff_rejects_non_image() {
        assert_eq!(sniff_image_format(b"<svg xmlns=...>"), None);
        assert_eq!(sniff_image_format(b""), None);
        assert_eq!(sniff_image_format(b"RIFF\x00\x00\x00\x00AVI "), None); // RIFF but not WEBP
    }

    #[test]
    fn validate_accepts_gif_and_webp() {
        assert!(validate_image("a.gif", "image/gif", &gif_bytes()).is_ok());
        assert!(validate_image("a.webp", "image/webp", &webp_bytes()).is_ok());
    }

    #[test]
    fn strip_exif_preserves_non_jpeg() {
        let data = b"not a jpeg";
        let result = strip_exif_jpeg(data);
        assert_eq!(result, data);
    }

    #[test]
    fn strip_exif_minimal_jpeg() {
        // Minimal JPEG: SOI + APP0 (JFIF) + SOS + EOI
        let mut jpeg = vec![0xFF, 0xD8]; // SOI
        // APP0 segment (marker + length + data)
        jpeg.extend_from_slice(&[0xFF, 0xE0, 0x00, 0x04, 0x00, 0x00]); // APP0, len=4
        // APP1 segment (EXIF, should be stripped)
        jpeg.extend_from_slice(&[0xFF, 0xE1, 0x00, 0x04, 0xAA, 0xBB]); // APP1, len=4
        // SOS + fake data
        jpeg.extend_from_slice(&[0xFF, 0xDA, 0x00, 0x02]);
        jpeg.extend_from_slice(&[0xFF, 0xD9]); // EOI

        let result = strip_exif_jpeg(&jpeg);
        // Should contain SOI, APP0, SOS+data, but NOT APP1
        assert!(result.len() < jpeg.len(), "EXIF should be stripped");
        let has_app1 = result.windows(2).any(|w| w == [0xFF, 0xE1]);
        assert!(!has_app1, "APP1 should be removed");
        let has_app0 = result.windows(2).any(|w| w == [0xFF, 0xE0]);
        assert!(has_app0, "APP0 should be preserved");
    }

    #[test]
    fn generate_key_format() {
        let key = generate_image_key("test-comm", "png");
        assert!(key.starts_with("mt/test-comm/"));
        assert!(
            std::path::Path::new(&key)
                .extension()
                .is_some_and(|ext| ext.eq_ignore_ascii_case("png"))
        );
    }

    #[test]
    fn allowed_extensions_match() {
        // Verify our ALLOWED_EXTENSIONS list matches validate_image behavior
        for ext in ALLOWED_EXTENSIONS {
            let filename = format!("test.{ext}");
            let (ct, bytes) = match *ext {
                "png" => ("image/png", png_bytes()),
                "jpg" | "jpeg" => ("image/jpeg", jpeg_bytes()),
                "gif" => ("image/gif", gif_bytes()),
                "webp" => ("image/webp", webp_bytes()),
                _ => continue,
            };
            assert!(
                validate_image(&filename, ct, &bytes).is_ok(),
                "Extension {ext} should be valid"
            );
        }
    }
}