// Copyright 2020-2022 Masaki Hara
// Licensed under MIT or Apache-2.0.
// See LICENSE and LICENSE-Apache-2.0 in the repository.

use std::io::{self, Read, Write};
use std::ops::Deref;

use image::error::{DecodingError, EncodingError, ImageFormatHint};
use image::{
    ColorType, DynamicImage, ImageBuffer, ImageDecoder, ImageError, ImageResult, Rgb, RgbImage,
    Rgba, RgbaImage,
};
use libwebp::boxed::WebpBox;

#[derive(Debug)]
pub struct WebpReader<R: Read> {
    reader: Reader<R>,
    index: usize,
}

impl<R: Read> WebpReader<R> {
    fn new(reader: Reader<R>) -> Self {
        Self { reader, index: 0 }
    }
}

impl<R: Read> Read for WebpReader<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let (_, _, _, image_buf) = self.reader.data().unwrap();
        let new_index = (self.index + buf.len()).min(image_buf.len());
        let num_written = new_index - self.index;
        buf[..num_written].copy_from_slice(&image_buf[self.index..new_index]);
        self.index = new_index;
        Ok(num_written)
    }
}

#[derive(Debug)]
pub struct WebpDecoder<R: Read> {
    reader: Reader<R>,
}

#[allow(clippy::upper_case_acronyms)]
#[derive(Debug)]
enum WebpColor {
    RGB,
    RGBA,
}

impl<R: Read> WebpDecoder<R> {
    pub fn new(reader: R) -> ImageResult<Self> {
        Self::new_inner(reader, WebpColor::RGBA)
    }

    pub fn new_rgba(reader: R) -> ImageResult<Self> {
        Self::new_inner(reader, WebpColor::RGBA)
    }

    pub fn new_rgb(reader: R) -> ImageResult<Self> {
        Self::new_inner(reader, WebpColor::RGB)
    }

    fn new_inner(reader: R, colortype: WebpColor) -> ImageResult<Self> {
        let mut reader = Reader::new(reader, colortype);
        reader.read_info()?;
        Ok(Self { reader })
    }
}

impl<'a, R: Read + 'a> ImageDecoder<'a> for WebpDecoder<R> {
    type Reader = WebpReader<R>;

    fn dimensions(&self) -> (u32, u32) {
        self.reader.info().unwrap()
    }
    fn color_type(&self) -> ColorType {
        match self.reader.colortype {
            WebpColor::RGB => ColorType::Rgb8,
            WebpColor::RGBA => ColorType::Rgba8,
        }
    }
    fn into_reader(mut self) -> ImageResult<Self::Reader> {
        self.reader.read_data()?;
        Ok(WebpReader::new(self.reader))
    }
}

const READER_READ_UNIT: usize = 1024;

#[derive(Debug)]
struct Reader<R: Read> {
    reader: R,
    colortype: WebpColor,
    buf: Vec<u8>,
    info: Option<(u32, u32)>,
    data: Option<(u32, u32, u32, WebpBox<[u8]>)>,
}

impl<R: Read> Reader<R> {
    fn new(reader: R, colortype: WebpColor) -> Self {
        Self {
            reader,
            colortype,
            buf: Vec::new(),
            info: None,
            data: None,
        }
    }

    fn info(&self) -> Option<(u32, u32)> {
        self.info
    }

    fn read_info(&mut self) -> io::Result<()> {
        if self.info.is_some() {
            return Ok(());
        }
        loop {
            let read_len = self.read_into_buf(READER_READ_UNIT)?;
            if let Ok(info) = libwebp::WebPGetInfo(&self.buf) {
                self.info = Some(info);
                return Ok(());
            }
            if read_len == 0 {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    "Invalid webp header",
                ));
            }
        }
    }

    fn data(&self) -> Option<(u32, u32, u32, &[u8])> {
        let (w, h, s, ref buf) = *self.data.as_ref()?;
        Some((w, h, s, buf))
    }

    fn read_data(&mut self) -> io::Result<()> {
        self.read_info()?;
        if self.data.is_some() {
            return Ok(());
        }

        self.reader.read_to_end(&mut self.buf)?;
        let data = match self.colortype {
            WebpColor::RGB => libwebp::WebPDecodeRGB(&self.buf).map(|(w, h, b)| (w, h, w * 3, b)),
            WebpColor::RGBA => libwebp::WebPDecodeRGBA(&self.buf).map(|(w, h, b)| (w, h, w * 4, b)),
        }
        .map_err(|_| io::Error::new(io::ErrorKind::InvalidInput, "Invalid webp data"))?;
        self.data = Some(data);
        Ok(())
    }

    fn read_into_buf(&mut self, by: usize) -> io::Result<usize> {
        let old_len = self.buf.len();
        self.buf.resize(old_len + by, 0);
        let result = self.reader.read(&mut self.buf[old_len..]);
        self.buf.resize(old_len + result.as_ref().unwrap_or(&0), 0);
        result
    }
}

pub fn webp_load<R: Read>(r: R) -> ImageResult<DynamicImage> {
    Ok(DynamicImage::ImageRgba8(webp_load_rgba(r)?))
}

pub fn webp_load_rgba<R: Read>(mut r: R) -> ImageResult<RgbaImage> {
    let mut buf = Vec::new();
    r.read_to_end(&mut buf)?;
    webp_load_rgba_from_memory(&buf)
}

pub fn webp_load_rgb<R: Read>(mut r: R) -> ImageResult<RgbImage> {
    let mut buf = Vec::new();
    r.read_to_end(&mut buf)?;
    webp_load_rgb_from_memory(&buf)
}

pub fn webp_load_from_memory(buf: &[u8]) -> ImageResult<DynamicImage> {
    Ok(DynamicImage::ImageRgba8(webp_load_rgba_from_memory(buf)?))
}

pub fn webp_load_rgba_from_memory(buf: &[u8]) -> ImageResult<RgbaImage> {
    let (width, height, buf) = libwebp::WebPDecodeRGBA(buf)
        .map_err(|_| DecodingError::new(ImageFormatHint::Unknown, "Webp Format Error".to_string()))
        .map_err(ImageError::Decoding)?;
    Ok(ImageBuffer::from_raw(width, height, buf.to_vec()).unwrap())
}

pub fn webp_load_rgb_from_memory(buf: &[u8]) -> ImageResult<RgbImage> {
    let (width, height, buf) = libwebp::WebPDecodeRGB(buf)
        .map_err(|_| DecodingError::new(ImageFormatHint::Unknown, "Webp Format Error".to_string()))
        .map_err(ImageError::Decoding)?;
    Ok(ImageBuffer::from_raw(width, height, buf.to_vec()).unwrap())
}

pub fn webp_write<W: Write>(img: &DynamicImage, w: &mut W) -> ImageResult<()> {
    match img {
        DynamicImage::ImageRgb8(img) => webp_write_rgb(img, w),
        DynamicImage::ImageRgba8(img) => webp_write_rgba(img, w),
        DynamicImage::ImageLuma8(_) => webp_write_rgb(&img.to_rgb8(), w),
        DynamicImage::ImageLumaA8(_) => webp_write_rgba(&img.to_rgba8(), w),
        DynamicImage::ImageRgb16(_) => webp_write_rgb(&img.to_rgb8(), w),
        DynamicImage::ImageRgba16(_) => webp_write_rgba(&img.to_rgba8(), w),
        DynamicImage::ImageLuma16(_) => webp_write_rgb(&img.to_rgb8(), w),
        DynamicImage::ImageLumaA16(_) => webp_write_rgba(&img.to_rgba8(), w),
        DynamicImage::ImageRgb32F(_) => webp_write_rgb(&img.to_rgb8(), w),
        DynamicImage::ImageRgba32F(_) => webp_write_rgba(&img.to_rgba8(), w),
        _ => webp_write_rgba(&img.to_rgba8(), w),
    }
}

pub fn webp_write_rgba<W: Write, C>(img: &ImageBuffer<Rgba<u8>, C>, w: &mut W) -> ImageResult<()>
where
    C: Deref<Target = [u8]>,
{
    let buf = libwebp::WebPEncodeRGBA(&img, img.width(), img.height(), img.width() * 4, 75.0)
        .map_err(|_| EncodingError::new(ImageFormatHint::Unknown, "Webp Format Error".to_string()))
        .map_err(ImageError::Encoding)?;
    w.write_all(&buf)?;
    Ok(())
}

pub fn webp_write_rgb<W: Write, C>(img: &ImageBuffer<Rgb<u8>, C>, w: &mut W) -> ImageResult<()>
where
    C: Deref<Target = [u8]>,
{
    let buf = libwebp::WebPEncodeRGB(&img, img.width(), img.height(), img.width() * 3, 75.0)
        .map_err(|_| EncodingError::new(ImageFormatHint::Unknown, "Webp Format Error".to_string()))
        .map_err(ImageError::Encoding)?;
    w.write_all(&buf)?;
    Ok(())
}