use std::io::Read;

use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{ProtocolReader, ProtocolReaderError, ProtocolReaderResult}, quantization, transform};

struct PixelTile {
    // i32: representation that supports Walsh-Hadamard
    pixels: [i16; TILE_SZ2]
}

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

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

pub fn decompress<R: Read>(
    reader: &mut ProtocolReader<R>
) -> ProtocolReaderResult<(u32, u32, Vec<Vec<u8>>)> {
    // read header
    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()?;

    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 = width_in_tiles * height_in_tiles * (n_layers as usize);
    let mut tiles = vec![QuantTile::new(); n_tiles];
    for i in 0..n_tiles {
        tiles[i].load_zero(reader)?;
    }
    for i in 0..n_tiles {
        tiles[i].load_rest(reader)?;
    }

    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 x0 in (0..width).step_by(TILE_SZ) {
            for y0 in (0..height).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
                );
            }
        }
    }
    return Ok((width, height, layers))
}

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 x in 0..TILE_SZ {
            for y 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 { quants: [0; TILE_SZ2] }
    }
    
    fn to_coef_tile(&self) -> CoefTile {
        CoefTile {
            coefs: quantization::from_quantized(self.quants)
        }
    }

    fn load_zero<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> {
        self.quants[0] = reader.read_i16_wide()?;
        Ok(())
    }

    fn load_rest<R: Read>(&mut self, reader: &mut ProtocolReader<R>) -> ProtocolReaderResult<()> {
        for i in 1..self.quants.len() {
            self.quants[i] = reader.read_i16_packed()?;
        }
        Ok(())
    }
    
}