use std::{
    fs::File,
    io::{BufReader, Read},
    path::Path,
};

use libcolors::surface::{Surface, VecSurface};

use crate::error::Error;

pub struct ImageData {
    pub data: Vec<u8>,
    pub pixel_format: PixelFormat,
    pub bit_depth: BitDepth,
    pub stride: usize,
    pub width: u32,
    pub height: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PixelFormat {
    Rgb,
    Rgba,
    #[allow(unused)]
    Bgr,
    Bgra,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BitDepth {
    Eight,
}

struct ImageLoader {
    #[allow(unused)]
    name: &'static str,
    probe: fn(&Path) -> Result<bool, Error>,
    load: fn(&Path) -> Result<ImageData, Error>,
}

static IMAGE_LOADERS: &[ImageLoader] = &[
    ImageLoader {
        name: "png",
        probe: png_probe,
        load: png_load,
    },
    ImageLoader {
        name: "jpeg",
        probe: jpeg_probe,
        load: jpeg_load,
    },
];

fn png_probe(path: &Path) -> Result<bool, Error> {
    let file = File::open(path).map(BufReader::new)?;
    let decoder = png::Decoder::new(file);
    Ok(decoder.read_info().is_ok())
}
fn jpeg_probe(path: &Path) -> Result<bool, Error> {
    let mut file = File::open(path).map(BufReader::new)?;
    let mut magic = [0; 4];
    let len = file.read(&mut magic)?;
    if len != 4 || magic != [0xFF, 0xD8, 0xFF, 0xE0] {
        Ok(false)
    } else {
        Ok(true)
    }
}

fn png_load(path: &Path) -> Result<ImageData, Error> {
    let file = File::open(path).map(BufReader::new)?;
    let decoder = png::Decoder::new(file);
    let mut info = decoder.read_info()?;
    let output_buffer_size = info.output_buffer_size().ok_or(Error::CorruptImage)?;
    let mut buffer = vec![0; output_buffer_size];
    let frame = info.next_frame(&mut buffer)?;
    buffer.drain(frame.buffer_size()..);
    let pixel_format = match frame.color_type {
        png::ColorType::Rgb => PixelFormat::Rgb,
        png::ColorType::Rgba => PixelFormat::Rgba,
        t => todo!("Unsupported color type: {t:?}"),
    };
    let bit_depth = match frame.bit_depth {
        png::BitDepth::Eight => BitDepth::Eight,
        t => todo!("Unsupported bit depth: {t:?}"),
    };
    Ok(ImageData {
        data: buffer,
        pixel_format,
        bit_depth,
        stride: frame.line_size,
        width: frame.width,
        height: frame.height,
    })
}
fn jpeg_load(path: &Path) -> Result<ImageData, Error> {
    let file = File::open(path).map(BufReader::new)?;
    let mut decoder = Box::new(pixie::JpegDecoder::new(file));
    let rgb = decoder.decode_rgb()?;

    Ok(ImageData {
        data: rgb.data,
        pixel_format: PixelFormat::Rgb,
        bit_depth: BitDepth::Eight,
        stride: rgb.stride,
        width: rgb.width as u32,
        height: rgb.height as u32,
    })
}

impl ImageData {
    pub fn load_image<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
        let path = path.as_ref();
        let mut loader = None;
        for l in IMAGE_LOADERS {
            if (l.probe)(path)? {
                loader = Some(l);
                break;
            }
        }
        let loader = loader.ok_or(Error::UnknownFormat)?;
        (loader.load)(path)
    }

    #[inline]
    fn convert_pixel(
        source: &[u8],
        source_format: PixelFormat,
        source_depth: BitDepth,
        output: &mut [u8],
        output_format: PixelFormat,
        output_depth: BitDepth,
    ) {
        let (r, g, b, a) = match (source_format, source_depth) {
            (PixelFormat::Rgb, BitDepth::Eight) => (source[0], source[1], source[2], 255),
            (PixelFormat::Rgba, BitDepth::Eight) => (source[0], source[1], source[2], source[3]),
            (PixelFormat::Bgr, BitDepth::Eight) => (source[2], source[1], source[0], 255),
            (PixelFormat::Bgra, BitDepth::Eight) => (source[2], source[1], source[0], source[3]),
        };

        match (output_format, output_depth) {
            (PixelFormat::Rgb, BitDepth::Eight) => {
                output[0] = r;
                output[1] = g;
                output[2] = b;
            }
            (PixelFormat::Rgba, BitDepth::Eight) => {
                output[0] = r;
                output[1] = g;
                output[2] = b;
                output[3] = a;
            }
            (PixelFormat::Bgr, BitDepth::Eight) => {
                output[0] = b;
                output[1] = g;
                output[2] = r;
            }
            (PixelFormat::Bgra, BitDepth::Eight) => {
                output[0] = b;
                output[1] = g;
                output[2] = r;
                output[3] = a;
            }
        }
    }

    // TODO optimized format conversion
    pub fn convert_color(
        &self,
        target_pixel_format: PixelFormat,
        target_bit_depth: BitDepth,
    ) -> ImageData {
        // TODO less than 8bpp
        let source_bytes_per_pixel = match (self.pixel_format, self.bit_depth) {
            (PixelFormat::Rgb | PixelFormat::Bgr, BitDepth::Eight) => 3,
            (PixelFormat::Rgba | PixelFormat::Bgra, BitDepth::Eight) => 4,
        };
        let output_bytes_per_pixel = match (target_pixel_format, target_bit_depth) {
            (PixelFormat::Rgb | PixelFormat::Bgr, BitDepth::Eight) => 3,
            (PixelFormat::Rgba | PixelFormat::Bgra, BitDepth::Eight) => 4,
        };
        let output_stride = (self.width as usize * output_bytes_per_pixel + 3) & !3;
        let output_buffer_size = output_stride * self.height as usize;
        let mut output_buffer = vec![0u8; output_buffer_size];

        for y in 0..self.height {
            let source_row_i = y as usize * self.stride;
            let source_row = &self.data[source_row_i..source_row_i + self.stride];
            let output_row_i = y as usize * output_stride;
            let output_row = &mut output_buffer[output_row_i..output_row_i + output_stride];

            // TODO SIMD impl
            for x in 0..self.width {
                let source_pixel_i = x as usize * source_bytes_per_pixel;
                let source_pixel =
                    &source_row[source_pixel_i..source_pixel_i + source_bytes_per_pixel];
                let output_pixel_i = x as usize * output_bytes_per_pixel;
                let output_pixel =
                    &mut output_row[output_pixel_i..output_pixel_i + output_bytes_per_pixel];

                Self::convert_pixel(
                    source_pixel,
                    self.pixel_format,
                    self.bit_depth,
                    output_pixel,
                    target_pixel_format,
                    target_bit_depth,
                );
            }
        }

        ImageData {
            data: output_buffer,
            pixel_format: target_pixel_format,
            bit_depth: target_bit_depth,
            stride: output_stride,
            width: self.width,
            height: self.height,
        }
    }

    pub fn create_surface(&self) -> VecSurface {
        if self.pixel_format == PixelFormat::Bgra && self.bit_depth == BitDepth::Eight {
            let mut surface = VecSurface::new(self.width, self.height);
            assert_eq!(surface.stride() * 4, self.stride);
            // TODO misaligned
            surface.copy_from_slice(bytemuck::cast_slice(&self.data[..]));
            surface
        } else {
            self.convert_color(PixelFormat::Bgra, BitDepth::Eight)
                .create_surface()
        }
    }
}

// TODO optimized scaling
pub fn scale_surface(source: &VecSurface, destination: &mut VecSurface) {
    let scale_x = destination.width() as f32 / source.width() as f32;
    let scale_y = destination.height() as f32 / source.height() as f32;

    let (source_scale_x, destination_scale_x, w) = if destination.width() > source.width() {
        (1.0 / scale_x, 1.0, destination.width())
    } else {
        (1.0, scale_x, source.width())
    };
    let (source_scale_y, destination_scale_y, h) = if destination.height() > source.height() {
        (1.0 / scale_y, 1.0, destination.height())
    } else {
        (1.0, scale_y, source.height())
    };

    for y in 0..h {
        for x in 0..w {
            let sx = (source_scale_x * x as f32) as u32;
            let sy = (source_scale_y * y as f32) as u32;
            let dx = (destination_scale_x * x as f32) as u32;
            let dy = (destination_scale_y * y as f32) as u32;

            destination.set_pixel(dx, dy, source.pixel(sx, sy));
        }
    }
}
