use std::io::Write;

use crate::{constants::{MAGIC, TILE_SZ, TILE_SZ2}, protocol::{self, ProtocolWriter, ProtocolWriterResult}, quality_settings::{self, QualitySettings}, quantization, transform};


struct PixelTile {
    quality: QualitySettings,
    pixels: [i16; TILE_SZ2]
}
struct CoefTile {
    quality: QualitySettings,
    coefs: [i16; TILE_SZ2]
}

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


pub fn compress<W: Write>(
    width: u32,
    height: u32,
    // n_components: u16,
    layers: &[(QualitySettings, &[u8])],
    writer: &mut protocol::ProtocolWriter<W>
) -> ProtocolWriterResult<()> {

    // validation
    for l in 0..layers.len() {
        assert!(layers[l].1.len() == width as usize * height as usize);
    }

    // write header
    writer.write_header(MAGIC)?;
    writer.write_u32_wide(width)?;
    writer.write_u32_wide(height)?;
    writer.write_u32_wide(layers.len() as u32)?;

    // write quality settings for each layout
    for (quality, _) in layers.iter() {
        writer.write_quality_settings(*quality)?;
    }

    // build list of tiles
    let mut tiles = vec![];
    for (quality, layer) in layers.iter() {
        for x0 in (0..width).step_by(TILE_SZ) {
            for y0 in (0..height).step_by(TILE_SZ) {
                let pixel_tile = PixelTile::from_layer(
                    x0 as usize, y0 as usize, *quality, layer, 
                    width as usize, height as usize
                );
                let coef_tile = CoefTile::from_pixel_tile(&pixel_tile);
                let quant_tile = QuantTile::from_coef_tile(&coef_tile);
                tiles.push(quant_tile);
            }
        }
    }

    // write the thumbnail block
    for t in tiles.iter() {
        t.write_zero(writer)?;
    }

    // write the tiles
    let mut prev: Option<QuantTile> = None;
    for t in tiles {
        t.write_rest(prev, writer)?;
        prev = Some(t);
    }

    Ok(())
}

impl PixelTile {
    fn from_layer(
        x0: usize,
        y0: usize,
        quality: QualitySettings,
        layer: &[u8],
        width: usize,
        height: usize
    ) -> PixelTile {
        let mut pixels = [0; TILE_SZ2];
        for x in 0..TILE_SZ {
            for y in 0..TILE_SZ {
                let src_x = x0 + x;
                let src_y = y0 + y;
                pixels[y * TILE_SZ + x] = 
                    if src_x < width && src_y < height { 
                        layer[src_y * width + src_x] as i16
                    } else {
                        0
                    };
            }
        }

        return PixelTile { quality, pixels };
    }
}

impl CoefTile {
    fn from_pixel_tile(pt: &PixelTile) -> CoefTile {
        let mut coefs = pt.pixels.clone();

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

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

        return CoefTile { quality: pt.quality, coefs }
    }
}

impl QuantTile {
    fn from_coef_tile(pt: &CoefTile) -> QuantTile {
        QuantTile {
            quality: pt.quality,
            quants: quantization::to_quantized(pt.quality, pt.coefs)
        }
    }

    fn write_zero<W: Write>(&self, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
        if self.quality.values[quality_settings::THUMBNAIL_IS_WIDE] != 0 {
            writer.write_i16_wide(self.quants[0])?;
        } else {
            writer.write_i16_packed(self.quants[0])?;
        }
        Ok(())
    }

    fn write_rest<W: Write>(&self, prev: Option<QuantTile>, writer: &mut ProtocolWriter<W>) -> ProtocolWriterResult<()> {
        if let Some(p) = prev {
            if self.quants[1..] == p.quants[1..] {
                // 255 zeroes can also encode "copy the rest from the previous tile"
                writer.write_u8_wide(0xff)?;  
                return Ok(());
            }
        }

        let mut i = 0;
        loop {
            let mut zeros: u8 = 0;
            while i < TILE_SZ2 && self.quants[i] == 0 {
                zeros += 1;
                i += 1;
            }
            writer.write_u8_wide(zeros)?;

            if i >= TILE_SZ2 { break; }

            writer.write_i16_packed(self.quants[i])?;
            i += 1;
        }
        Ok(())
    }
}