//! Logo import pipeline: decode, measure the ink box, trim, downscale and derive the print and
//! knockout PNGs. Pure functions over bytes; storage lives in `commands::logo`.
use image::codecs::png::{CompressionType, FilterType, PngEncoder};
use image::imageops::{self, FilterType as ResizeFilter};
use image::{DynamicImage, ImageFormat, ImageReader, RgbImage, RgbaImage};
use serde::{Deserialize, Serialize};
use std::io::Cursor;

/// Largest accepted input file.
pub const MAX_INPUT_BYTES: usize = 20 * 1024 * 1024;
/// Largest accepted side of the decoded image, in pixels.
const MAX_SIDE: u32 = 8192;
/// The longer side of the derived print image never exceeds this.
const PRINT_MAX_SIDE: u32 = 1200;
/// Pixels with alpha at or below this are not ink (alpha logos).
const INK_ALPHA: u8 = 32;
/// An image counts as having transparency when any pixel is more transparent than this.
const ALPHA_OPAQUE: u8 = 250;
/// Alpha at or below this is treated as fully transparent when cleaning the print image.
const ALPHA_EMPTY: u8 = 5;
/// Colour distance from the background under which an opaque pixel is "empty".
const BG_DISTANCE: i32 = 24;
/// Rows and columns with fewer ink pixels than this are noise and do not widen the trim box.
const MIN_LINE_INK: usize = 2;

#[derive(Debug, thiserror::Error)]
pub enum LogoError {
    #[error("Use a PNG, JPEG or WebP image")]
    Format,
    #[error("The image is larger than 20 MB")]
    TooLarge,
    #[error("The image is larger than {MAX_SIDE} pixels on a side")]
    Dimensions,
    #[error("Could not read the image: {0}")]
    Decode(String),
    #[error("The image has no visible content")]
    Empty,
    #[error("Could not encode the logo: {0}")]
    Encode(String),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum LogoKind {
    Wordmark,
    Mark,
    Tall,
}

/// Measurements of the trimmed, downscaled print image. Stored as JSON in `app_settings.logo_meta`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct LogoMetrics {
    pub width: u32,
    pub height: u32,
    pub aspect: f64,
    pub ink_density: f64,
    pub mean_color: [u8; 3],
    pub kind: LogoKind,
    pub has_alpha: bool,
    pub opaque_background: bool,
    pub source_width: u32,
    pub source_height: u32,
}

#[derive(Debug)]
pub struct LogoDerived {
    pub print_png: Vec<u8>,
    /// White ink with the original alpha; only for logos that have transparency.
    pub knockout_png: Option<Vec<u8>>,
    pub metrics: LogoMetrics,
}

/// File extension for a sniffed original: png, jpg or webp.
pub fn sniff_extension(bytes: &[u8]) -> Result<&'static str, LogoError> {
    sniff(bytes).map(|(_, ext)| ext)
}

fn sniff(bytes: &[u8]) -> Result<(ImageFormat, &'static str), LogoError> {
    if bytes.starts_with(&[0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A]) {
        Ok((ImageFormat::Png, "png"))
    } else if bytes.starts_with(&[0xFF, 0xD8, 0xFF]) {
        Ok((ImageFormat::Jpeg, "jpg"))
    } else if bytes.len() >= 12 && &bytes[0..4] == b"RIFF" && &bytes[8..12] == b"WEBP" {
        Ok((ImageFormat::WebP, "webp"))
    } else {
        Err(LogoError::Format)
    }
}

fn decode(bytes: &[u8]) -> Result<RgbaImage, LogoError> {
    if bytes.len() > MAX_INPUT_BYTES {
        return Err(LogoError::TooLarge);
    }
    let (format, _) = sniff(bytes)?;
    let decode_err = |e: image::ImageError| LogoError::Decode(e.to_string());
    let decoder = ImageReader::with_format(Cursor::new(bytes), format)
        .into_decoder()
        .map_err(decode_err)?;
    // Check the declared size before allocating the pixels.
    let (w, h) = image::ImageDecoder::dimensions(&decoder);
    if w > MAX_SIDE || h > MAX_SIDE {
        return Err(LogoError::Dimensions);
    }
    DynamicImage::from_decoder(decoder)
        .map(|img| img.to_rgba8())
        .map_err(decode_err)
}

/// Median of each channel over the four 4x4 corner patches.
fn corner_background(img: &RgbaImage) -> [u8; 3] {
    let (w, h) = img.dimensions();
    let (pw, ph) = (4.min(w), 4.min(h));
    let mut channels: [Vec<u8>; 3] = Default::default();
    for (x0, y0) in [(0, 0), (w - pw, 0), (0, h - ph), (w - pw, h - ph)] {
        for y in y0..y0 + ph {
            for x in x0..x0 + pw {
                let p = img.get_pixel(x, y).0;
                for c in 0..3 {
                    channels[c].push(p[c]);
                }
            }
        }
    }
    let mut out = [0u8; 3];
    for c in 0..3 {
        channels[c].sort_unstable();
        out[c] = channels[c][channels[c].len() / 2];
    }
    out
}

/// What counts as ink: opaque enough for alpha logos, far enough from the background otherwise.
#[derive(Clone, Copy)]
enum InkTest {
    Alpha,
    Background([u8; 3]),
}

impl InkTest {
    fn is_ink(self, p: &[u8; 4]) -> bool {
        match self {
            InkTest::Alpha => p[3] > INK_ALPHA,
            InkTest::Background(bg) => {
                let d: i32 = (0..3).map(|c| (p[c] as i32 - bg[c] as i32).pow(2)).sum();
                d > BG_DISTANCE * BG_DISTANCE
            }
        }
    }
}

/// First and last index whose count reaches `MIN_LINE_INK`.
fn span(counts: &[usize]) -> Option<(usize, usize)> {
    let first = counts.iter().position(|&n| n >= MIN_LINE_INK)?;
    let last = counts.iter().rposition(|&n| n >= MIN_LINE_INK)?;
    Some((first, last))
}

fn premultiply(img: &mut RgbaImage) {
    for p in img.pixels_mut() {
        let a = p.0[3] as u32;
        for c in 0..3 {
            p.0[c] = ((p.0[c] as u32 * a + 127) / 255) as u8;
        }
    }
}

fn unpremultiply(img: &mut RgbaImage) {
    for p in img.pixels_mut() {
        let a = p.0[3] as u32;
        if a == 0 {
            continue;
        }
        for c in 0..3 {
            p.0[c] = ((p.0[c] as u32 * 255 + a / 2) / a).min(255) as u8;
        }
    }
}

/// Lanczos3 downscale to at most `PRINT_MAX_SIDE`. Colours are premultiplied while resampling so
/// transparent pixels do not bleed their (arbitrary) colour into the edges; the result is straight alpha.
fn downscale(img: RgbaImage, has_alpha: bool) -> RgbaImage {
    let (w, h) = img.dimensions();
    let longer = w.max(h);
    if longer <= PRINT_MAX_SIDE {
        return img;
    }
    let scale = PRINT_MAX_SIDE as f64 / longer as f64;
    let nw = ((w as f64 * scale).round() as u32).clamp(1, PRINT_MAX_SIDE);
    let nh = ((h as f64 * scale).round() as u32).clamp(1, PRINT_MAX_SIDE);
    let mut src = img;
    if has_alpha {
        premultiply(&mut src);
    }
    let mut out = imageops::resize(&src, nw, nh, ResizeFilter::Lanczos3);
    if has_alpha {
        unpremultiply(&mut out);
    }
    out
}

/// Exporters leave faint noise in "solid" and "empty" pixels (alpha 253 of 255, hidden colour under
/// alpha 0). It is invisible but triples the PNG size, so snap near-solid alpha to solid, near-empty
/// to empty, and drop the colour of empty pixels.
fn clean_alpha(img: &mut RgbaImage) {
    for p in img.pixels_mut() {
        match p.0[3] {
            a if a >= ALPHA_OPAQUE => p.0[3] = 255,
            a if a <= ALPHA_EMPTY => p.0 = [0, 0, 0, 0],
            _ => {}
        }
    }
}

/// Keep 64 levels per channel (steps of about 4/255, invisible in print). Anti-aliasing and encoder
/// noise of +-1 otherwise makes flat colours incompressible: a sample wordmark logo drops from 149 KB to 75 KB.
fn quantize_colour(img: &mut RgbaImage) {
    for p in img.pixels_mut() {
        for c in 0..3 {
            p.0[c] = ((((p.0[c] as u32 * 63 + 127) / 255) * 255 + 31) / 63) as u8;
        }
    }
}

fn encode_png(img: &DynamicImage) -> Result<Vec<u8>, LogoError> {
    let mut buf = Vec::new();
    let encoder =
        PngEncoder::new_with_quality(&mut buf, CompressionType::Best, FilterType::Adaptive);
    img.write_with_encoder(encoder)
        .map_err(|e| LogoError::Encode(e.to_string()))?;
    Ok(buf)
}

fn kind_for(aspect: f64) -> LogoKind {
    if aspect >= 2.5 {
        LogoKind::Wordmark
    } else if aspect <= 0.6 {
        LogoKind::Tall
    } else {
        LogoKind::Mark
    }
}

/// Decode a logo, trim it to its ink box, and derive the print PNG, the knockout PNG and the metrics.
pub fn process_logo(bytes: &[u8]) -> Result<LogoDerived, LogoError> {
    let source = decode(bytes)?;
    let (source_width, source_height) = source.dimensions();
    let has_alpha = source.pixels().any(|p| p.0[3] < ALPHA_OPAQUE);
    // The background estimate only drives measuring; no pixel is edited.
    let test = if has_alpha {
        InkTest::Alpha
    } else {
        InkTest::Background(corner_background(&source))
    };

    let mut rows = vec![0usize; source_height as usize];
    let mut cols = vec![0usize; source_width as usize];
    for (x, y, p) in source.enumerate_pixels() {
        if test.is_ink(&p.0) {
            rows[y as usize] += 1;
            cols[x as usize] += 1;
        }
    }
    let (Some((top, bottom)), Some((left, right))) = (span(&rows), span(&cols)) else {
        return Err(LogoError::Empty);
    };

    let cropped = imageops::crop_imm(
        &source,
        left as u32,
        top as u32,
        (right - left + 1) as u32,
        (bottom - top + 1) as u32,
    )
    .to_image();
    let mut print = downscale(cropped, has_alpha);
    if has_alpha {
        clean_alpha(&mut print);
    }
    quantize_colour(&mut print);
    let (width, height) = print.dimensions();

    let (mut ink, mut sum) = (0u64, [0u64; 3]);
    for p in print.pixels() {
        if test.is_ink(&p.0) {
            ink += 1;
            for (total, channel) in sum.iter_mut().zip(p.0) {
                *total += channel as u64;
            }
        }
    }
    let mean_color = if ink == 0 {
        [0; 3]
    } else {
        [0, 1, 2].map(|c| ((sum[c] + ink / 2) / ink) as u8)
    };
    let aspect = width as f64 / height as f64;

    let knockout_png = if has_alpha {
        let mut white = print.clone();
        for p in white.pixels_mut() {
            p.0[0] = 255;
            p.0[1] = 255;
            p.0[2] = 255;
        }
        Some(encode_png(&DynamicImage::ImageRgba8(white))?)
    } else {
        None
    };
    let print_image = if has_alpha {
        DynamicImage::ImageRgba8(print)
    } else {
        // No transparency to keep: RGB is smaller.
        DynamicImage::ImageRgb8(RgbImage::from_fn(width, height, |x, y| {
            let p = print.get_pixel(x, y).0;
            image::Rgb([p[0], p[1], p[2]])
        }))
    };

    Ok(LogoDerived {
        print_png: encode_png(&print_image)?,
        knockout_png,
        metrics: LogoMetrics {
            width,
            height,
            aspect,
            ink_density: ink as f64 / (width as f64 * height as f64),
            mean_color,
            kind: kind_for(aspect),
            has_alpha,
            opaque_background: !has_alpha,
            source_width,
            source_height,
        },
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use image::{codecs::jpeg::JpegEncoder, codecs::webp::WebPEncoder, ExtendedColorType, Rgba};

    fn png_bytes(img: &RgbaImage) -> Vec<u8> {
        let mut out = Vec::new();
        DynamicImage::ImageRgba8(img.clone())
            .write_to(&mut Cursor::new(&mut out), ImageFormat::Png)
            .unwrap();
        out
    }

    /// Transparent canvas with an opaque rectangle.
    fn alpha_logo(w: u32, h: u32, rect: (u32, u32, u32, u32), color: [u8; 3]) -> RgbaImage {
        let mut img = RgbaImage::from_pixel(w, h, Rgba([0, 0, 0, 0]));
        let (x0, y0, rw, rh) = rect;
        for y in y0..y0 + rh {
            for x in x0..x0 + rw {
                img.put_pixel(x, y, Rgba([color[0], color[1], color[2], 255]));
            }
        }
        img
    }

    fn decode_png(bytes: &[u8]) -> RgbaImage {
        image::load_from_memory_with_format(bytes, ImageFormat::Png)
            .unwrap()
            .to_rgba8()
    }

    #[test]
    fn alpha_logo_is_trimmed_to_the_rect_without_padding() {
        let img = alpha_logo(200, 100, (30, 20, 60, 25), [10, 20, 30]);
        let out = process_logo(&png_bytes(&img)).unwrap();
        let m = &out.metrics;
        assert_eq!((m.width, m.height), (60, 25));
        assert_eq!((m.source_width, m.source_height), (200, 100));
        assert!(m.has_alpha && !m.opaque_background);
        let printed = decode_png(&out.print_png);
        assert_eq!(printed.dimensions(), (60, 25));
        // Every edge pixel row/column is ink: nothing was left as padding.
        assert!(printed.pixels().all(|p| p.0[3] == 255));
        // Colours are quantised to 64 levels, so allow a few steps.
        for (got, want) in m.mean_color.iter().zip([10i32, 20, 30]) {
            assert!((*got as i32 - want).abs() <= 3, "{:?}", m.mean_color);
        }
    }

    #[test]
    fn noise_specks_do_not_widen_the_box() {
        let mut img = alpha_logo(200, 100, (50, 30, 40, 20), [0, 0, 0]);
        // Isolated single pixels far from the logo.
        img.put_pixel(2, 2, Rgba([0, 0, 0, 255]));
        img.put_pixel(197, 97, Rgba([0, 0, 0, 255]));
        img.put_pixel(100, 3, Rgba([0, 0, 0, 255]));
        let m = process_logo(&png_bytes(&img)).unwrap().metrics;
        assert_eq!((m.width, m.height), (40, 20));
    }

    #[test]
    fn faint_alpha_is_not_ink() {
        let mut img = alpha_logo(100, 100, (40, 40, 20, 20), [0, 0, 0]);
        // A soft shadow below the alpha threshold.
        for y in 70..90 {
            for x in 10..90 {
                img.put_pixel(x, y, Rgba([0, 0, 0, 30]));
            }
        }
        let m = process_logo(&png_bytes(&img)).unwrap().metrics;
        assert_eq!((m.width, m.height), (20, 20));
    }

    #[test]
    fn opaque_white_background_jpeg_is_trimmed_and_flagged() {
        let mut img = RgbaImage::from_pixel(160, 120, Rgba([255, 255, 255, 255]));
        for y in 40..80 {
            for x in 30..130 {
                img.put_pixel(x, y, Rgba([20, 40, 160, 255]));
            }
        }
        let rgb = DynamicImage::ImageRgba8(img).to_rgb8();
        let mut jpg = Vec::new();
        JpegEncoder::new_with_quality(&mut jpg, 95)
            .encode(rgb.as_raw(), 160, 120, ExtendedColorType::Rgb8)
            .unwrap();

        let out = process_logo(&jpg).unwrap();
        let m = &out.metrics;
        assert!(m.opaque_background && !m.has_alpha);
        // JPEG ringing may move an edge by a pixel or two.
        assert!((98..=102).contains(&m.width), "{}", m.width);
        assert!((38..=42).contains(&m.height), "{}", m.height);
        assert!(out.knockout_png.is_none());
        // The input pixels are not edited: the printed corners are still the logo colour, not transparent.
        let printed =
            image::load_from_memory_with_format(&out.print_png, ImageFormat::Png).unwrap();
        assert!(!printed.color().has_alpha());
    }

    #[test]
    fn lossless_webp_is_accepted() {
        let img = alpha_logo(80, 60, (10, 10, 30, 30), [200, 0, 0]);
        let mut webp = Vec::new();
        WebPEncoder::new_lossless(&mut webp)
            .encode(img.as_raw(), 80, 60, ExtendedColorType::Rgba8)
            .unwrap();
        assert_eq!(sniff_extension(&webp).unwrap(), "webp");
        let m = process_logo(&webp).unwrap().metrics;
        assert_eq!((m.width, m.height), (30, 30));
        assert!(m.has_alpha);
    }

    #[test]
    fn other_formats_and_garbage_are_rejected() {
        let msg = "Use a PNG, JPEG or WebP image";
        for bad in [
            &b"GIF89a\x01\x00\x01\x00"[..],
            b"<svg xmlns='http://www.w3.org/2000/svg'/>",
            b"not an image at all",
            b"",
            b"RIFF\x00\x00\x00\x00WAVEfmt ",
        ] {
            assert_eq!(process_logo(bad).unwrap_err().to_string(), msg);
        }
        // Right magic, broken body.
        let mut truncated = png_bytes(&alpha_logo(20, 20, (2, 2, 5, 5), [0; 3]));
        truncated.truncate(30);
        assert!(matches!(
            process_logo(&truncated),
            Err(LogoError::Decode(_))
        ));
    }

    #[test]
    fn oversized_input_and_dimensions_are_rejected() {
        let mut big = vec![0x89, b'P', b'N', b'G', 0x0D, 0x0A, 0x1A, 0x0A];
        big.resize(MAX_INPUT_BYTES + 1, 0);
        assert!(matches!(process_logo(&big), Err(LogoError::TooLarge)));

        // A blank image 8193 px wide compresses to a few KB, so it passes the byte cap.
        let wide = RgbaImage::new(8193, 2);
        assert!(matches!(
            process_logo(&png_bytes(&wide)),
            Err(LogoError::Dimensions)
        ));
    }

    #[test]
    fn blank_images_have_no_content() {
        assert!(matches!(
            process_logo(&png_bytes(&RgbaImage::new(40, 40))),
            Err(LogoError::Empty)
        ));
        let white = RgbaImage::from_pixel(40, 40, Rgba([255, 255, 255, 255]));
        assert!(matches!(
            process_logo(&png_bytes(&white)),
            Err(LogoError::Empty)
        ));
    }

    #[test]
    fn large_logos_are_downscaled_to_1200() {
        let img = alpha_logo(3000, 1000, (10, 10, 2980, 980), [0, 0, 0]);
        let out = process_logo(&png_bytes(&img)).unwrap();
        let m = &out.metrics;
        assert_eq!((m.width, m.height), (1200, 395));
        assert_eq!((m.source_width, m.source_height), (3000, 1000));
        assert_eq!(decode_png(&out.print_png).dimensions(), (1200, 395));
        assert!((m.aspect - 2980.0 / 980.0).abs() < 0.01);
    }

    #[test]
    fn density_aspect_and_kind() {
        // Half of a 100 x 40 box is ink: left half filled.
        let img = alpha_logo(100, 40, (0, 0, 50, 40), [0, 0, 0]);
        // Trimming shrinks the box to the ink, so add a second block to keep the density at 0.5.
        let mut img = img;
        for y in 0..40 {
            img.put_pixel(99, y, Rgba([0, 0, 0, 255]));
            img.put_pixel(98, y, Rgba([0, 0, 0, 255]));
        }
        let m = process_logo(&png_bytes(&img)).unwrap().metrics;
        assert_eq!((m.width, m.height), (100, 40));
        assert!((m.aspect - 2.5).abs() < 1e-9);
        assert_eq!(m.kind, LogoKind::Wordmark);
        assert!((m.ink_density - 0.52).abs() < 1e-9, "{}", m.ink_density);

        let square = alpha_logo(50, 50, (10, 10, 30, 30), [0; 3]);
        assert_eq!(
            process_logo(&png_bytes(&square)).unwrap().metrics.kind,
            LogoKind::Mark
        );
        let tall = alpha_logo(50, 100, (10, 10, 12, 40), [0; 3]);
        assert_eq!(
            process_logo(&png_bytes(&tall)).unwrap().metrics.kind,
            LogoKind::Tall
        );
        assert_eq!(kind_for(2.49), LogoKind::Mark);
        assert_eq!(kind_for(0.6), LogoKind::Tall);
    }

    #[test]
    fn knockout_is_white_with_the_original_alpha() {
        let mut img = alpha_logo(60, 40, (5, 5, 40, 20), [30, 60, 90]);
        // Anti-aliased edge column with partial alpha.
        for y in 5..25 {
            img.put_pixel(44, y, Rgba([30, 60, 90, 128]));
        }
        let out = process_logo(&png_bytes(&img)).unwrap();
        let print = decode_png(&out.print_png);
        let knock = decode_png(&out.knockout_png.expect("alpha logo has a knockout"));
        assert_eq!(print.dimensions(), knock.dimensions());
        for (a, b) in print.pixels().zip(knock.pixels()) {
            assert_eq!(&b.0[..3], &[255, 255, 255]);
            assert_eq!(a.0[3], b.0[3]);
        }
        assert!(knock.pixels().any(|p| p.0[3] == 128));
    }

    #[test]
    fn sample_logo_is_trimmed_and_small() {
        let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("../src/pdf/testing/fixtures/sample-wordmark-logo.png");
        let bytes = std::fs::read(path).unwrap();
        let out = process_logo(&bytes).unwrap();
        let m = &out.metrics;
        assert_eq!((m.source_width, m.source_height), (2172, 724));
        assert!((5.3..=5.6).contains(&m.aspect), "aspect {}", m.aspect);
        assert!(
            (0.50..=0.60).contains(&m.ink_density),
            "density {}",
            m.ink_density
        );
        assert_eq!(m.kind, LogoKind::Wordmark);
        assert!(
            out.print_png.len() <= 120 * 1024,
            "print is {} bytes",
            out.print_png.len()
        );
        assert!(out.knockout_png.is_some());
    }
}