use std::io::{ErrorKind, Read};

use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, image::Image, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quality_settings::{self, QualitySettings}, quantization, transform};

struct PixelTile {
    pixels: [i16; TILE_SZ2]
}

#[derive(Clone)]
struct CoefTile {
    coefs: [i16; TILE_SZ2]
}

#[derive(Clone, Copy)]
struct QuantTile {
    quality: QualitySettings,
    quants: [i16; TILE_SZ2]
}

pub enum ProgressiveImage {
    Present { image: Image, complete: bool },
    NotHereYet
}

pub fn decompress<R: Read>(
    reader: &mut ProtocolReader<R>
) -> ProtocolReaderResult<ProgressiveImage> {
    // read header
    let read_header_and_stuff: Result<_, ProtocolReaderError> = (|| {
        let header = reader.read_header()?;
        if header.bytes != MAGIC.bytes {
            return Err(ProtocolReaderError::WrongHeader);
        }
        if header.version != MAGIC.version {
            return Err(ProtocolReaderError::WrongVersion);
        }

        let width = reader.read_u32_wide()?;
        let height = reader.read_u32_wide()?;
        let n_layers = reader.read_u32_wide()?;

        // read quality settings for each layout
        let mut quality_settings = vec![];
        for _ in 0..n_layers {
            quality_settings.push(reader.read_quality_settings()?);
        }
        Ok((width, height, n_layers, quality_settings))
    })();

    let Ok((width, height, n_layers, quality_settings)) = read_header_and_stuff else {
        match read_header_and_stuff {
            Ok(_) => unreachable!(),
            Err(ProtocolReaderError::Io(io)) => {
                if io.kind() == ErrorKind::UnexpectedEof {
                    return ProtocolReaderResult::Ok(ProgressiveImage::NotHereYet)
                }
                return Err(ProtocolReaderError::Io(io))
            }
            Err(e) => return Err(e),
        }

    };

    let width_in_tiles = (width as usize + TILE_SZ - 1) / TILE_SZ;
    let height_in_tiles = (height as usize + TILE_SZ - 1) / TILE_SZ;

    let n_tiles_per_layer = width_in_tiles * height_in_tiles;

    let mut tiles = vec![QuantTile::new(); n_tiles_per_layer * n_layers as usize];
    let read_tiles_and_stuff: Result<(), ProtocolReaderError> = (|tiles: &mut [QuantTile]| {
        // read thumbnail block
        for l in 0..n_layers {
            for i in 0..n_tiles_per_layer {
                tiles[l as usize * n_tiles_per_layer + i].load_zero(
                    quality_settings[l as usize], reader
                )?;
            }
        }

        // read remaining tiles
        let mut prev: Option<QuantTile> = None;
        for i in 0..n_tiles_per_layer * n_layers as usize {
            tiles[i].load_rest(
                prev,
                reader
            )?;
            prev = Some(tiles[i]);
        }
        Ok(())
    })(&mut tiles);

    let complete;
    match read_tiles_and_stuff {
        Err(ProtocolReaderError::Io(i)) => {
            if i.kind() == ErrorKind::UnexpectedEof {
                complete = false
            } else {
                return Err(ProtocolReaderError::Io(i))
            }
        }
        Err(e) => return Err(e),
        Ok(()) => {
            complete = true;
        }
    };

    let mut tile_i = 0;
    let mut layers = vec![vec![0; width as usize * height as usize]; n_layers as usize];
    for layer in 0..n_layers {
        for y0 in (0..height).step_by(TILE_SZ) {
            for x0 in (0..width).step_by(TILE_SZ) {
                let pixel_tile = tiles[tile_i].to_coef_tile().to_pixel_tile();
                tile_i += 1;
                pixel_tile.to_layer(
                    x0 as usize, y0 as usize, &mut layers[layer as usize], 
                    width as usize, height as usize
                );
            }
        }
    }

    let mut image = Image::new(width as usize, height as usize);
    for (l, layer) in layers.into_iter().enumerate() {
        image.add(quality_settings[l], layer).expect("by construction, this layer should be OK");
    }

    return Ok(ProgressiveImage::Present { image: image, complete });
}

impl CoefTile {
    fn to_pixel_tile(&self) -> PixelTile {
        let mut pixels = self.coefs.clone();

        // columns
        for x in 0..TILE_SZ {
            transform::decode(&mut pixels, x, 8);
        }

        // rows
        for y in 0..TILE_SZ {
            transform::decode(&mut pixels, y * TILE_SZ, 1);
        }

        return PixelTile { pixels };
    }
}

impl PixelTile {
    fn to_layer(&self, x0: usize, y0: usize, layer: &mut [u8], width: usize, height: usize) {
        for y in 0..TILE_SZ {
            for x in 0..TILE_SZ {
                let dst_x = x0 + x;
                let dst_y = y0 + y;
                if dst_x < width && dst_y < height { 
                    layer[dst_y * width + dst_x] = self.clip(self.pixels[y * TILE_SZ + x]);
                }
            }
        }
    }
    
    fn clip(&self, x: i16) -> u8 {
        if x > u8::MAX as i16 {
            return u8::MAX
        }
        if x < u8::MIN as i16 {
            return u8::MIN;
        }
        return x as u8;
    }
}

impl QuantTile {
    fn new() -> QuantTile {
        QuantTile { 
            quality: QualitySettings { values: [0, 0, 0, 0, 0, 0, 0, 0] },
            quants: [0; TILE_SZ2] 
        }
    }
    
    fn to_coef_tile(&self) -> CoefTile {
        CoefTile {
            coefs: quantization::from_quantized(self.quality, self.quants)
        }
    }

    fn load_zero<R: Read>(
        &mut self, 
        quality: QualitySettings, 
        reader: &mut ProtocolReader<R>
    ) -> ProtocolReaderResult<()> {
        self.quality = quality;
        if quality.values[quality_settings::THUMBNAIL_IS_WIDE] != 0 {
            self.quants[0] = reader.read_i16_wide()?;
        } else {
            self.quants[0] = reader.read_i16_packed()?;
        }
        Ok(())
    }

    fn load_rest<R: Read>(&
        mut self, 
        prev: Option<QuantTile>,
        reader: &mut ProtocolReader<R>
    ) -> ProtocolReaderResult<()> {
        let mut i: usize = 0;
        loop {
            let cmd = reader.read_u8_wide()? as usize;
            if let Some(p) = prev {
                if cmd == 0xff {
                    self.quants[1..].copy_from_slice(&p.quants[1..]);
                    return Ok(());
                }
            }

            i += cmd;
            if i >= TILE_SZ2 { break; }

            self.quants[i] = reader.read_i16_packed()?;
            i += 1;
        }
 
        Ok(())
    }
    
}