//! 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::{self, File},
    io::BufWriter,
    path::Path,
};

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

pub const TILE_HEIGHT: u64 = 256;
pub const TILE_WIDTH: u64 = 256;

/// Main "lazy image buffer" struct.
#[derive(Debug, Clone)]
pub struct TileLayer<P: Pixel> {
    tiles: FnvHashMap<(u64, u64), ImageBuffer<P, Vec<P::Subpixel>>>,
    default_pixel: P,
    width: u64,
    height: u64,
}

impl<P: Pixel> TileLayer<P> {
    pub fn from_pixel(width: u64, height: u64, pixel: P) -> Self {
        TileLayer {
            tiles: Default::default(),
            default_pixel: pixel,
            width,
            height,
        }
    }

    pub fn width(&self) -> u64 {
        self.width
    }

    pub fn height(&self) -> u64 {
        self.height
    }

    fn index(&self, x: u64, y: u64) -> ((u64, u64), (u32, u32)) {
        (
            (x / TILE_WIDTH, y / TILE_HEIGHT),
            (
                (x % TILE_WIDTH).try_into().unwrap(),
                (y % TILE_HEIGHT).try_into().unwrap(),
            ),
        )
    }

    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))
    }

    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)
        })
    }

    pub fn tile_for_mut(&mut self, x: u64, y: u64) -> &mut ImageBuffer<P, Vec<P::Subpixel>> {
        let ((tx, ty), _) = self.index(x, y);
        self.tile_mut(tx, ty)
    }

    pub fn get_pixel_checked(&self, x: u64, y: u64) -> Option<&P> {
        if x >= self.width || y >= self.height {
            return None;
        }

        let (outer_idx, (inner_x, inner_y)) = self.index(x, y);
        self.tiles
            .get(&outer_idx)
            .map(|tile| tile.get_pixel(inner_x, inner_y))
            .or(Some(&self.default_pixel))
    }

    pub fn get_pixel(&self, x: u64, y: u64) -> &P {
        // This is kinda cheating, but we care about the API for now, not the speed. We can
        // optimize this later.
        self.get_pixel_checked(x, y).unwrap()
    }

    pub fn get_pixel_mut_checked(&mut self, x: u64, y: u64) -> Option<&mut P> {
        if x >= self.width || y >= self.height {
            return None;
        }

        let ((outer_x, outer_y), (inner_x, inner_y)) = self.index(x, y);
        Some(
            self.tile_mut(outer_x, outer_y)
                .get_pixel_mut(inner_x, inner_y),
        )
    }

    pub fn get_pixel_mut(&mut self, x: u64, y: u64) -> &mut P {
        self.get_pixel_mut_checked(x, y).unwrap()
    }

    /// Enumerate all pixels that are explicitely set in this layer.
    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,
                )
            })
        })
    }

    /// Mutably enumerate all pixels that are explicitely set in this layer.
    pub fn enumerate_pixels_mut(&mut self) -> impl Iterator<Item = (u64, u64, &mut P)> {
        self.tiles.iter_mut().flat_map(|((tx, ty), tile)| {
            tile.enumerate_pixels_mut().map(move |(x, y, p)| {
                (
                    u64::from(x) + tx * TILE_WIDTH,
                    u64::from(y) + ty * TILE_HEIGHT,
                    p,
                )
            })
        })
    }

    pub fn pixels(&self) -> impl Iterator<Item = &P> {
        self.enumerate_pixels().map(|x| x.2)
    }

    pub fn pixels_mut(&mut self) -> impl Iterator<Item = &mut P> {
        self.enumerate_pixels_mut().map(|x| x.2)
    }

    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 repeatedly calling [`get_pixel_mut`], as it groups
    /// pixels by tile and only does one tile lookup.
    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 TileLayer<Rgba<u8>> {
    pub fn save_to_directory<S: AsRef<Path>>(&self, path: S) -> Result<()> {
        let path = path.as_ref();

        for ((x, y), tile) in self.tiles.iter() {
            let folder = path.join(&x.to_string());
            let metadata = folder.metadata();
            match metadata {
                Err(_) => fs::create_dir(&folder)?,
                Ok(m) if !m.is_dir() => bail!("Output path is not a directory"),
                _ => {}
            }
            let file = folder.join(&format!("{y}.png"));
            compress_png(tile, file)?;
        }

        Ok(())
    }
}

pub fn compress_png<P: AsRef<Path>>(image: &RgbaImage, path: P) -> Result<()> {
    let outstream = BufWriter::new(File::create(path)?);
    let encoder =
        PngEncoder::new_with_quality(outstream, CompressionType::Best, FilterType::Adaptive);

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

    Ok(())
}

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