//! TRACES: FR-EXP-1 | FR-EXP-2 | FR-EXP-8
//! The encoders.
//!
//! All four write RGB, not RGBA. The pipeline produces an opaque frame — no
//! operation makes a pixel transparent, and the shader writes 1.0 into alpha
//! unconditionally — so a fourth channel would be a third more bytes carrying
//! the same value in every pixel, and a PNG that some tools then treat as
//! having meaningful transparency.
//!
//! # The colour profile
//!
//! All four embed one, in the place their container puts it: a JPEG APP2
//! segment, a PNG `iCCP` chunk, TIFF tag 34675. Encoding correctly and
//! labelling correctly are separate jobs and both are required — pixels in
//! Display P3 with no profile are read as sRGB and come out desaturated,
//! which is a worse outcome than not offering the space at all.
//!
//! sRGB gets one too, rather than relying on it being everyone's default.
//! Untagged is not the same as tagged sRGB: it means "guess", and the guess
//! differs between a browser, a phone gallery and a print shop.
//!
//! # Metadata
//!
//! Nothing is written. `strip_location` defaults to on (FR-EXP-8) and this
//! satisfies it in the strongest possible way: there is no EXIF block, so
//! there is no GPS tag, no serial number, and no lens history in the file
//! that leaves the machine.
//!
//! The other half of FR-EXP-8 — *retaining* camera and copyright metadata
//! when the user asks for it — is not implemented, and cannot be faked by
//! omission. It needs the source's EXIF carried through `dr-decode` and
//! re-serialised here, which is a piece of work in its own right and belongs
//! with the batch-export path that would make it worth having.

use dr_types::{ExportFormat, ExportSettings};

use crate::{icc, ExportError};

/// Encode a resized, sharpened RGBA buffer to the requested format.
///
/// The buffer is already encoded into `settings.colour_space` — that happened
/// in the shader, at the only point where the unclipped colour still existed.
/// All that is left here is to say so.
pub fn encode(
    rgba: &[u8],
    width: u32,
    height: u32,
    settings: &ExportSettings,
) -> Result<Vec<u8>, ExportError> {
    let profile = icc::profile(settings.colour_space);
    match settings.format {
        ExportFormat::Jpeg => jpeg(rgba, width, height, settings.quality, &profile),
        ExportFormat::Png => png(rgba, width, height, &profile),
        ExportFormat::Tiff8 => tiff8(rgba, width, height, &profile),
        ExportFormat::Tiff16 => tiff16(rgba, width, height, &profile),
        other => Err(ExportError::FormatUnsupported(other)),
    }
}

/// Drop alpha, which the pipeline never varies.
fn rgb(rgba: &[u8]) -> Vec<u8> {
    let mut out = Vec::with_capacity(rgba.len() / 4 * 3);
    for px in rgba.chunks_exact(4) {
        out.extend_from_slice(&px[..3]);
    }
    out
}

fn jpeg(
    rgba: &[u8],
    width: u32,
    height: u32,
    quality: u8,
    profile: &[u8],
) -> Result<Vec<u8>, ExportError> {
    let mut bytes = Vec::new();
    let mut encoder = jpeg_encoder::Encoder::new(&mut bytes, quality);
    // Splits across APP2 segments itself if it has to. The profiles this crate
    // generates fit in one, but the branch is the encoder's rather than ours.
    encoder
        .add_icc_profile(profile)
        .map_err(|e| ExportError::Encode(e.to_string()))?;
    encoder
        .encode(
            &rgb(rgba),
            width as u16,
            height as u16,
            jpeg_encoder::ColorType::Rgb,
        )
        .map_err(|e| ExportError::Encode(e.to_string()))?;
    Ok(bytes)
}

fn png(rgba: &[u8], width: u32, height: u32, profile: &[u8]) -> Result<Vec<u8>, ExportError> {
    let mut bytes = Vec::new();
    {
        // Built through `Info` rather than the setters, because the profile is
        // the one field `png::Encoder` has no setter for. The `sRGB` chunk is
        // deliberately left unset: the crate writes `iCCP` only in its
        // absence, and an sRGB chunk beside a P3 profile is a contradiction a
        // reader has to pick a side of.
        let mut info = png::Info::with_size(width, height);
        info.color_type = png::ColorType::Rgb;
        info.bit_depth = png::BitDepth::Eight;
        info.icc_profile = Some(std::borrow::Cow::Borrowed(profile));

        let encoder = png::Encoder::with_info(&mut bytes, info)
            .map_err(|e| ExportError::Encode(e.to_string()))?;
        let mut writer = encoder
            .write_header()
            .map_err(|e| ExportError::Encode(e.to_string()))?;
        writer
            .write_image_data(&rgb(rgba))
            .map_err(|e| ExportError::Encode(e.to_string()))?;
        writer
            .finish()
            .map_err(|e| ExportError::Encode(e.to_string()))?;
    }
    Ok(bytes)
}

/// The ICC profile as a TIFF tag value.
///
/// A newtype only because the tag's field type is `UNDEFINED` (7) and the
/// `tiff` crate maps a plain `&[u8]` to `BYTE` (1). Both are single bytes and
/// most readers do not look, but libtiff declares `TIFFTAG_ICCPROFILE` as
/// undefined and a strict reader is entitled to agree with it.
struct IccTag<'a>(&'a [u8]);

impl tiff::encoder::TiffValue for IccTag<'_> {
    const BYTE_LEN: u8 = 1;
    const FIELD_TYPE: tiff::tags::Type = tiff::tags::Type::UNDEFINED;

    fn count(&self) -> usize {
        self.0.len()
    }

    fn data(&self) -> std::borrow::Cow<'_, [u8]> {
        std::borrow::Cow::Borrowed(self.0)
    }
}

/// Tag 34675, `InterColourProfile`. Not in the `tiff` crate's `Tag` enum.
const TAG_ICC_PROFILE: u16 = 34675;

fn tiff8(rgba: &[u8], width: u32, height: u32, profile: &[u8]) -> Result<Vec<u8>, ExportError> {
    use tiff::encoder::{colortype, TiffEncoder};

    let mut bytes = std::io::Cursor::new(Vec::new());
    let mut encoder =
        TiffEncoder::new(&mut bytes).map_err(|e| ExportError::Encode(e.to_string()))?;
    let mut image = encoder
        .new_image::<colortype::RGB8>(width, height)
        .map_err(|e| ExportError::Encode(e.to_string()))?;
    tag_profile(image.encoder(), profile)?;
    image
        .write_data(&rgb(rgba))
        .map_err(|e| ExportError::Encode(e.to_string()))?;
    Ok(bytes.into_inner())
}

/// Add the profile tag to a directory being built.
///
/// Shared by both TIFF widths, and separate from them because `write_image`
/// cannot be used once there is a tag to add — the directory has to be opened,
/// written into, and closed by hand.
fn tag_profile<W, K>(
    dir: &mut tiff::encoder::DirectoryEncoder<'_, W, K>,
    profile: &[u8],
) -> Result<(), ExportError>
where
    W: std::io::Write + std::io::Seek,
    K: tiff::encoder::TiffKind,
{
    dir.write_tag(tiff::tags::Tag::Unknown(TAG_ICC_PROFILE), IccTag(profile))
        .map_err(|e| ExportError::Encode(e.to_string()))
}

/// 16-bit TIFF, for work continuing in another editor.
///
/// **Honest about what it carries.** The adjust pass renders to an 8-bit
/// target (`AdjustPass::FORMAT` is `Rgba8Unorm`, and the generated shader
/// declares `texture_storage_2d<rgba8unorm, write>`), so the samples widened
/// here hold eight bits of information in a sixteen-bit container. The file
/// is a correct 16-bit TIFF and will round-trip through any editor without
/// further loss — but it does not resurrect precision the pipeline already
/// quantised away.
///
/// Making it mean what it says is a pipeline change rather than an encoder
/// one: the composer has to be told what format to write, and export has to
/// ask for the wide one (FR-EXP-9). Until then this is a container promotion,
/// which is still the right thing to hand an editor that works in 16-bit.
fn tiff16(rgba: &[u8], width: u32, height: u32, profile: &[u8]) -> Result<Vec<u8>, ExportError> {
    use tiff::encoder::{colortype, TiffEncoder};

    // `x * 257` rather than `x << 8`: it maps 255 to 65535 exactly, where the
    // shift maps it to 65280 and makes white slightly grey.
    let wide: Vec<u16> = rgb(rgba).iter().map(|&v| u16::from(v) * 257).collect();

    let mut bytes = std::io::Cursor::new(Vec::new());
    let mut encoder =
        TiffEncoder::new(&mut bytes).map_err(|e| ExportError::Encode(e.to_string()))?;
    let mut image = encoder
        .new_image::<colortype::RGB16>(width, height)
        .map_err(|e| ExportError::Encode(e.to_string()))?;
    tag_profile(image.encoder(), profile)?;
    image
        .write_data(&wide)
        .map_err(|e| ExportError::Encode(e.to_string()))?;
    Ok(bytes.into_inner())
}

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

    #[test]
    fn alpha_is_dropped_before_encoding() {
        let rgba = vec![1, 2, 3, 255, 4, 5, 6, 255];
        assert_eq!(rgb(&rgba), vec![1, 2, 3, 4, 5, 6]);
    }

    #[test]
    fn white_widens_to_full_scale_not_almost() {
        // The bug a left-shift introduces: 255 << 8 is 65280, so pure white
        // comes out a quarter of a percent grey in every 16-bit export.
        assert_eq!(u16::from(255u8) * 257, u16::MAX);
        assert_eq!(u16::from(0u8) * 257, 0);
        // And the midpoint stays the midpoint.
        assert_eq!(u16::from(128u8) * 257, 32896);
    }

    #[test]
    fn a_png_round_trips_its_pixels_exactly() {
        // PNG is lossless, so this is a real end-to-end check that the buffer
        // reaching the encoder is the one we think it is — channel order
        // included, which a size assertion would not catch.
        let rgba: Vec<u8> = vec![
            255, 0, 0, 255, // red
            0, 255, 0, 255, // green
            0, 0, 255, 255, // blue
            10, 20, 30, 255,
        ];
        let bytes = png(&rgba, 2, 2, &icc::profile(ColourSpace::Srgb)).unwrap();

        let decoder = png::Decoder::new(std::io::Cursor::new(&bytes));
        let mut reader = decoder.read_info().unwrap();
        let mut out = vec![0; reader.output_buffer_size().unwrap()];
        let info = reader.next_frame(&mut out).unwrap();

        assert_eq!((info.width, info.height), (2, 2));
        assert_eq!(info.color_type, png::ColorType::Rgb);
        assert_eq!(&out[..12], &[255, 0, 0, 0, 255, 0, 0, 0, 255, 10, 20, 30]);
    }

    /// A flat frame, for the tests that care only about what surrounds the
    /// pixels.
    fn flat(w: u32, h: u32) -> Vec<u8> {
        vec![128; (w * h * 4) as usize]
    }

    #[test]
    fn a_png_carries_a_profile_a_decoder_gets_back_intact() {
        // Read out through the PNG decoder, so this exercises the iCCP
        // chunk's deflate round-trip rather than asserting the encoder was
        // called. A truncated or mis-deflated profile is one a reader
        // discards silently, leaving the file to be guessed at as sRGB.
        for space in ColourSpace::ALL {
            let want = icc::profile(space);
            let bytes = png(&flat(4, 4), 4, 4, &want).unwrap();

            let decoder = png::Decoder::new(std::io::Cursor::new(&bytes));
            let reader = decoder.read_info().unwrap();
            let got = reader
                .info()
                .icc_profile
                .as_ref()
                .unwrap_or_else(|| panic!("{space:?} PNG carries no profile"));
            assert_eq!(got.as_ref(), want.as_slice(), "{space:?}");
        }
    }

    #[test]
    fn a_jpeg_carries_its_profile_in_a_well_formed_app2_segment() {
        // The APP2 form is exacting: the marker, a length, the string
        // "ICC_PROFILE\0", then a chunk index and count before the payload.
        // A reader that does not find that header ignores the segment, so
        // walking it here is the only way to know the file is really tagged.
        for space in ColourSpace::ALL {
            let want = icc::profile(space);
            let bytes = jpeg(&flat(4, 4), 4, 4, 90, &want).unwrap();
            let got = jpeg_icc(&bytes)
                .unwrap_or_else(|| panic!("{space:?} JPEG has no ICC_PROFILE segment"));
            assert_eq!(got, want, "{space:?}");
        }
    }

    /// Walk a JPEG's marker segments and reassemble the ICC profile.
    ///
    /// Hand-rolled because `zune-jpeg` decodes pixels and this is about what
    /// travels beside them.
    fn jpeg_icc(bytes: &[u8]) -> Option<Vec<u8>> {
        const TAG: &[u8] = b"ICC_PROFILE\0";
        let mut out = Vec::new();
        let mut i = 2; // past the SOI
        while i + 4 <= bytes.len() {
            if bytes[i] != 0xFF {
                return None;
            }
            let marker = bytes[i + 1];
            // Start of scan: entropy-coded data follows and there are no more
            // parseable segments.
            if marker == 0xDA {
                break;
            }
            let len = usize::from(u16::from_be_bytes([bytes[i + 2], bytes[i + 3]]));
            let payload = &bytes[i + 4..i + 2 + len];
            if marker == 0xE2 && payload.starts_with(TAG) {
                // Two bytes of chunk index and count follow the tag.
                out.extend_from_slice(&payload[TAG.len() + 2..]);
            }
            i += 2 + len;
        }
        (!out.is_empty()).then_some(out)
    }

    #[test]
    fn both_tiff_widths_carry_the_profile_in_tag_34675() {
        // Read back through the `tiff` decoder's own tag lookup. Writing the
        // tag with the wrong field type or a stale offset produces a file that
        // still opens — with no profile, and therefore the wrong colours.
        use tiff::decoder::Decoder;
        use tiff::tags::Tag;

        for space in ColourSpace::ALL {
            let want = icc::profile(space);
            for (label, bytes) in [
                ("8-bit", tiff8(&flat(4, 4), 4, 4, &want).unwrap()),
                ("16-bit", tiff16(&flat(4, 4), 4, 4, &want).unwrap()),
            ] {
                let mut d = Decoder::new(std::io::Cursor::new(&bytes)).expect("decode");
                let got = d
                    .get_tag_u8_vec(Tag::Unknown(TAG_ICC_PROFILE))
                    .unwrap_or_else(|e| panic!("{space:?} {label} TIFF: {e}"));
                assert_eq!(got, want, "{space:?} {label}");
            }
        }
    }

    #[test]
    fn a_tiff_still_decodes_to_its_pixels_with_the_extra_tag_present() {
        // Adding a tag means opening the directory by hand instead of using
        // `write_image`, which is the sort of change that produces a valid
        // header over unreadable strips.
        use tiff::decoder::{Decoder, DecodingResult};

        let rgba: Vec<u8> = vec![
            255, 0, 0, 255, // red
            0, 255, 0, 255, // green
            0, 0, 255, 255, // blue
            10, 20, 30, 255,
        ];
        let bytes = tiff8(&rgba, 2, 2, &icc::profile(ColourSpace::Srgb)).unwrap();
        let mut d = Decoder::new(std::io::Cursor::new(&bytes)).expect("decode");
        assert_eq!(d.dimensions().expect("dimensions"), (2, 2));
        let DecodingResult::U8(pixels) = d.read_image().expect("read") else {
            panic!("expected 8-bit samples");
        };
        assert_eq!(
            &pixels[..12],
            &[255, 0, 0, 0, 255, 0, 0, 0, 255, 10, 20, 30]
        );
    }
}