exr 1.74.2

//! How to write either a single or a list of layers.

use crate::{
    block::BlockIndex,
    image::{
        recursive::{NoneMore, Recursive},
        write::channels::{ChannelsWriter, WritableChannels},
        Layer, Layers,
    },
    meta::{
        attribute::TileDescription,
        compute_chunk_count,
        header::{Header, ImageAttributes},
        Headers,
    },
    prelude::SmallVec,
};

/// Enables an image containing this list of layers to be written to a file.
pub trait WritableLayers<'slf> {
    /// Generate the file meta data for this list of layers
    fn infer_headers(&self, image_attributes: &ImageAttributes) -> Headers;

    /// The type of temporary writer
    type Writer: LayersWriter;

    /// Create a temporary writer for this list of layers
    fn create_writer(&'slf self, headers: &[Header]) -> Self::Writer;
}

/// A temporary writer for a list of channels
pub trait LayersWriter: Sync {
    /// Deliver a block of pixels from a single layer to be stored in the file
    fn extract_uncompressed_block(&self, headers: &[Header], block: BlockIndex) -> Vec<u8>;
}

/// A temporary writer for an arbitrary list of layers
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct AllLayersWriter<ChannelsWriter> {
    layers: SmallVec<[LayerWriter<ChannelsWriter>; 2]>,
}

/// A temporary writer for a single layer
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct LayerWriter<ChannelsWriter> {
    channels: ChannelsWriter, // impl ChannelsWriter
}

// impl for smallvec
impl<'slf, Channels: 'slf> WritableLayers<'slf> for Layers<Channels>
where
    Channels: WritableChannels<'slf>,
{
    type Writer = AllLayersWriter<Channels::Writer>;

    fn infer_headers(&self, image_attributes: &ImageAttributes) -> Headers {
        slice_infer_headers(self.as_slice(), image_attributes)
    }

    fn create_writer(&'slf self, headers: &[Header]) -> Self::Writer {
        slice_create_writer(self.as_slice(), headers)
    }
}

fn slice_infer_headers<'slf, Channels: 'slf + WritableChannels<'slf>>(
    slice: &[Layer<Channels>],
    image_attributes: &ImageAttributes,
) -> Headers {
    slice.iter().map(|layer| layer.infer_headers(image_attributes).remove(0)).collect()
    // TODO no array-vs-first
}

fn slice_create_writer<'slf, Channels: 'slf + WritableChannels<'slf>>(
    slice: &'slf [Layer<Channels>],
    headers: &[Header],
) -> AllLayersWriter<Channels::Writer> {
    AllLayersWriter {
        layers: slice
            .iter()
            .zip(headers.chunks_exact(1)) // TODO no array-vs-first
            .map(|(layer, header)| layer.create_writer(header))
            .collect(),
    }
}

impl<'slf, Channels: WritableChannels<'slf>> WritableLayers<'slf> for Layer<Channels> {
    type Writer = LayerWriter</* 'l, */ Channels::Writer>;

    fn infer_headers(&self, image_attributes: &ImageAttributes) -> Headers {
        let blocks = match self.encoding.blocks {
            crate::image::Blocks::ScanLines => crate::meta::BlockDescription::ScanLines,
            crate::image::Blocks::Tiles(tile_size) => {
                let (level_mode, rounding_mode) = self.channel_data.infer_level_modes();
                crate::meta::BlockDescription::Tiles(TileDescription {
                    level_mode,
                    rounding_mode,
                    tile_size,
                })
            }
        };

        let chunk_count = compute_chunk_count(self.encoding.compression, self.size, blocks);

        let header = Header {
            channels: self.channel_data.infer_channel_list(),
            compression: self.encoding.compression,

            blocks,
            chunk_count,

            line_order: self.encoding.line_order,
            layer_size: self.size,
            shared_attributes: image_attributes.clone(),
            own_attributes: self.attributes.clone(),

            deep: false, // TODO deep data
            deep_data_version: None,
            max_samples_per_pixel: None,
        };

        smallvec![header] // TODO no array-vs-first
    }

    fn create_writer(&'slf self, headers: &[Header]) -> Self::Writer {
        let channels =
            self.channel_data.create_writer(headers.first().expect("inferred header error")); // TODO no array-vs-first

        LayerWriter {
            channels,
        }
    }
}

impl<C> LayersWriter for AllLayersWriter<C>
where
    C: ChannelsWriter,
{
    fn extract_uncompressed_block(&self, headers: &[Header], block: BlockIndex) -> Vec<u8> {
        self.layers[block.layer]
            .extract_uncompressed_block(std::slice::from_ref(&headers[block.layer]), block)
        // TODO no array-vs-first
    }
}

impl<C> LayersWriter for LayerWriter<C>
where
    C: ChannelsWriter,
{
    fn extract_uncompressed_block(&self, headers: &[Header], block: BlockIndex) -> Vec<u8> {
        self.channels
            .extract_uncompressed_block(headers.first().expect("invalid inferred header"), block)
        // TODO no array-vs-first
    }
}

impl<'slf> WritableLayers<'slf> for NoneMore {
    type Writer = NoneMore;

    fn infer_headers(&self, _: &ImageAttributes) -> Headers {
        SmallVec::new()
    }

    fn create_writer(&'slf self, _: &[Header]) -> Self::Writer {
        NoneMore
    }
}

impl<'slf, InnerLayers, Channels> WritableLayers<'slf> for Recursive<InnerLayers, Layer<Channels>>
where
    InnerLayers: WritableLayers<'slf>,
    Channels: WritableChannels<'slf>,
{
    type Writer = RecursiveLayersWriter<InnerLayers::Writer, Channels::Writer>;

    fn infer_headers(&self, image_attributes: &ImageAttributes) -> Headers {
        let mut headers = self.inner.infer_headers(image_attributes);
        headers.push(self.value.infer_headers(image_attributes).remove(0)); // TODO no unwrap
        headers
    }

    fn create_writer(&'slf self, headers: &[Header]) -> Self::Writer {
        let (own_header, inner_headers) =
            headers.split_last().expect("header has not been inferred correctly");

        let layer_index = inner_headers.len();
        RecursiveLayersWriter {
            inner: self.inner.create_writer(inner_headers),
            value: (layer_index, self.value.create_writer(std::slice::from_ref(own_header))), /* TODO no slice */
        }
    }
}

type RecursiveLayersWriter<InnerLayersWriter, ChannelsWriter> =
    Recursive<InnerLayersWriter, (usize, LayerWriter<ChannelsWriter>)>;

impl LayersWriter for NoneMore {
    /// # Panics
    /// Panics if called, as this indicates a recursive length mismatch bug
    /// where a block is being extracted for a layer index that doesn't
    /// exist in the recursive layer structure.
    fn extract_uncompressed_block(&self, _: &[Header], _: BlockIndex) -> Vec<u8> {
        unreachable!(
            "recursive length mismatch bug: attempted to extract block for non-existent layer"
        )
    }
}

impl<InnerLayersWriter, Channels> LayersWriter
    for RecursiveLayersWriter<InnerLayersWriter, Channels>
where
    InnerLayersWriter: LayersWriter,
    Channels: ChannelsWriter,
{
    fn extract_uncompressed_block(&self, headers: &[Header], block: BlockIndex) -> Vec<u8> {
        let (layer_index, layer) = &self.value;
        if *layer_index == block.layer {
            let header = headers.get(*layer_index).expect("layer index bug");
            layer.extract_uncompressed_block(std::slice::from_ref(header), block)
        // TODO no slice?
        } else {
            self.inner.extract_uncompressed_block(headers, block)
        }
    }
}