//! Lazy tiled image.
//!
//! This supports OSM-style "tiled" images, but not all of the tiles have to be present. If a tile
//! is not present, a default pixel is returned. The tile is allocated with the first call to a
//! mutating operation.
use std::{
    fs::File,
    io::{BufWriter, Write},
    path::Path,
};

use color_eyre::eyre::Result;
use fnv::FnvHashMap;
use image::{
    codecs::png::{CompressionType, FilterType, PngEncoder},
    ColorType, ImageBuffer, ImageEncoder, Pixel, RgbaImage,
};
use num_traits::Zero;
use rayon::iter::{IntoParallelIterator, ParallelIterator};

/// Height of a single tile.
pub const TILE_HEIGHT: u64 = 256;
/// Width of a single tile.
pub const TILE_WIDTH: u64 = 256;

type TileIndex = (u64, u64);

/// Main "lazy image buffer" struct.
///
/// This lazily allocates a new tile (of size [`TILE_WIDTH`] × [`TILE_HEIGHT`]) for each mutable
/// pixel access. Each tile is pre-filled with the given default pixel.
#[derive(Debug, Clone)]
pub struct TileLayer<P: Pixel> {
    tiles: FnvHashMap<TileIndex, ImageBuffer<P, Vec<P::Subpixel>>>,
    default_pixel: P,
}

impl<P: Pixel> TileLayer<P> {
    /// Construct a new lazy buffer with the given default (background) pixel.
    ///
    /// Note that this does not yet allocate any image tiles.
    pub fn from_pixel(pixel: P) -> Self {
        TileLayer {
            tiles: Default::default(),
            default_pixel: pixel,
        }
    }

    /// Iterates over all tiles, together with their indices.
    pub fn enumerate_tiles(
        &self,
    ) -> impl Iterator<Item = (u64, u64, &ImageBuffer<P, Vec<P::Subpixel>>)> {
        self.tiles.iter().map(|((x, y), t)| (*x, *y, t))
    }

    /// Returns a mutable reference to the given tile.
    ///
    /// This allocates a new tile if the requested tile does not yet exist.
    pub fn tile_mut(&mut self, tile_x: u64, tile_y: u64) -> &mut ImageBuffer<P, Vec<P::Subpixel>> {
        self.tiles.entry((tile_x, tile_y)).or_insert_with(|| {
            ImageBuffer::from_pixel(TILE_WIDTH as u32, TILE_HEIGHT as u32, self.default_pixel)
        })
    }

    /// Enumerate all pixels that are allocated.
    ///
    /// This provides access to the pixel and its coordinates.
    pub fn enumerate_pixels(&self) -> impl Iterator<Item = (u64, u64, &P)> {
        self.tiles.iter().flat_map(|((tx, ty), tile)| {
            tile.enumerate_pixels().map(move |(x, y, p)| {
                (
                    u64::from(x) + tx * TILE_WIDTH,
                    u64::from(y) + ty * TILE_HEIGHT,
                    p,
                )
            })
        })
    }

    /// Iterate over all pixels that are allocated.
    pub fn pixels(&self) -> impl Iterator<Item = &P> {
        self.enumerate_pixels().map(|x| x.2)
    }

    /// Returns the number of allocated tiles.
    pub fn tile_count(&self) -> usize {
        self.tiles.len()
    }

    /// Copies the non-zero pixels from `source` to `self`.
    ///
    /// A zero-pixel is identified by comparing all its channels' values with `Zero::zero()`. If
    /// any channel is non-zero, the pixel is considered non-zero and is copied.
    ///
    /// The top-left pixel of `source` is copied to `(x, y)`.
    ///
    /// This method is more efficient than copying pixels one by one, as it groups them by tile and
    /// only does one tile lookup then.
    pub fn blit_nonzero(&mut self, x: u64, y: u64, source: &ImageBuffer<P, Vec<P::Subpixel>>) {
        let zero = zero_pixel::<P>();
        let source_width = u64::from(source.width());
        let source_height = u64::from(source.height());
        for tx in x / TILE_WIDTH..=(x + source_width) / TILE_WIDTH {
            for ty in y / TILE_HEIGHT..=(y + source_height) / TILE_HEIGHT {
                let tile = self.tile_mut(tx, ty);
                let offset_x = (tx * TILE_WIDTH).saturating_sub(x);
                let offset_y = (ty * TILE_HEIGHT).saturating_sub(y);
                let local_min_x = x.saturating_sub(tx * TILE_WIDTH);
                let local_min_y = y.saturating_sub(ty * TILE_HEIGHT);
                let local_max_x = TILE_WIDTH.min(x + source_width - tx * TILE_WIDTH);
                let local_max_y = TILE_HEIGHT.min(y + source_height - ty * TILE_HEIGHT);
                // Keep x in the inner loop for better cache locality!
                for (y, source_y) in (local_min_y..local_max_y).zip(offset_y..) {
                    for (x, source_x) in (local_min_x..local_max_x).zip(offset_x..) {
                        let pixel = source
                            .get_pixel(source_x.try_into().unwrap(), source_y.try_into().unwrap());
                        if pixel.channels() != zero.channels() {
                            *tile.get_pixel_mut(x.try_into().unwrap(), y.try_into().unwrap()) =
                                *pixel;
                        }
                    }
                }
            }
        }
    }
}

impl<P> TileLayer<P>
where
    P: Pixel + Send,
    P::Subpixel: Send,
{
    /// Turns this lazy tile layer into a parallelized iterator.
    pub fn into_parallel_tiles(
        self,
    ) -> impl ParallelIterator<Item = (u64, u64, ImageBuffer<P, Vec<P::Subpixel>>)> {
        IntoParallelIterator::into_par_iter(self.tiles).map(|((x, y), t)| (x, y, t))
    }
}

/// Saves the given image buffer to the given path.
pub fn compress_png<P: AsRef<Path>>(image: &RgbaImage, path: P) -> Result<()> {
    let outstream = BufWriter::new(File::create(path)?);
    compress_png_stream(image, outstream)
}

/// Saves the given image buffer to the given stream.
///
/// Note that this uses the best compression available.
pub fn compress_png_stream<W: Write>(image: &RgbaImage, outstream: W) -> Result<()> {
    let encoder =
        PngEncoder::new_with_quality(outstream, CompressionType::Best, FilterType::Adaptive);

    encoder.write_image(image, image.width(), image.height(), ColorType::Rgba8)?;

    Ok(())
}

/// Encodes the given image buffer and returns its data as a vector.
pub fn compress_png_as_bytes(image: &RgbaImage) -> Result<Vec<u8>> {
    let mut buffer = Vec::new();
    compress_png_stream(image, &mut buffer)?;
    Ok(buffer)
}

fn zero_pixel<P: Pixel>() -> P {
    let zeroes = vec![Zero::zero(); P::CHANNEL_COUNT as usize];
    *P::from_slice(&zeroes)
}