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

use crate::{jpeg::decoder::read_u8, jpeg::error::Error};

pub const LUT_BITS: u8 = 9;

#[derive(Debug, Clone, Copy)]
pub struct LutEntry {
    pub symbol: u8,
    pub length: u8,
}

#[derive(Debug)]
struct Lut<const BITS: usize>(Box<[LutEntry; 1 << BITS]>)
where
    [LutEntry; 1 << BITS]: Sized;

#[derive(Debug)]
pub struct HuffmanTable {
    pub(crate) bits: [u8; 16],
    pub(crate) huffval: [u8; 256],

    lut: Lut<{ LUT_BITS as usize }>,
}

impl<const BITS: usize> Lut<BITS>
where
    [Option<LutEntry>; 1 << BITS]: Sized,
{
    pub fn new() -> Self {
        Self(Box::new(
            [LutEntry {
                length: 0,
                symbol: 0,
            }; 1 << BITS],
        ))
    }
}

#[derive(Default)]
pub struct HuffmanDecoder {
    buffer: u64,
    bit_count: usize,
    pub marker: Option<u8>,
}

impl HuffmanTable {
    pub fn new(
        counts: [u8; 16],
        values: [u8; 256],
        is_dc: bool,
        is_progressive: bool,
    ) -> Result<Self, Error> {
        // TODO handle progressive decoding
        _ = is_dc;
        _ = is_progressive;

        // Build huffman code table
        let (huffcode, huffsize) = Self::derive_huffman_codes(&counts)?;
        let mut lut = Lut::new();

        for (i, &size) in huffsize
            .iter()
            .enumerate()
            .filter(|&(_, &size)| size <= LUT_BITS)
        {
            let extra_bits = LUT_BITS - size;
            let start = (huffcode[i] as usize) << extra_bits;

            for j in 0..1 << extra_bits {
                lut.0[start + j] = LutEntry {
                    symbol: values[i],
                    length: size,
                };
            }
        }

        let mut p = 0;
        let mut si = huffsize[0];
        let mut code = 0;
        let mut maxcode = [0; 18];
        while p < huffsize.len() && huffsize[p] != 0 {
            while p < huffsize.len() && huffsize[p] == si {
                code += 1;
                p += 1;
            }
            maxcode[si as usize] = code << (16 - si);

            code <<= 1;
            si += 1;
        }

        p = 0;
        let mut offset = [0; 18];
        for l in 0..16 {
            if counts[l] == 0 {
                maxcode[l] = -1;
            } else {
                offset[l] = (p as i32) - (huffcode[p] as i32);
                p += usize::from(counts[l]);
            }
        }

        offset[17] = 0;
        offset[17] = 0x000F_FFFF;

        Ok(Self {
            lut,
            bits: counts,
            huffval: values,
        })
    }

    #[inline]
    pub fn lookup(&self, index: usize) -> LutEntry {
        self.lut.0[index]
    }

    fn derive_huffman_codes(counts: &[u8; 16]) -> Result<(Vec<u16>, Vec<u8>), Error> {
        let huffsize = counts
            .iter()
            .enumerate()
            .fold(Vec::new(), |mut acc, (i, &value)| {
                acc.extend(iter::repeat_n((i + 1) as u8, value as usize));
                acc
            });
        let mut huffcode = vec![0u16; huffsize.len()];
        let mut code_size = huffsize[0];
        let mut code = 0;

        for (i, &size) in huffsize.iter().enumerate() {
            while code_size < size {
                code <<= 1;
                code_size += 1;
            }

            if code >= 1 << size {
                return Err(Error::InvalidHuffmanTable);
            }

            huffcode[i] = code as u16;
            code += 1;
        }

        Ok((huffcode, huffsize))
    }
}

impl HuffmanDecoder {
    pub fn decode_symbol<R: Read>(
        &mut self,
        reader: &mut R,
        table: &HuffmanTable,
    ) -> Result<u8, Error> {
        if self.bit_count < 16 {
            self.read_bits(reader)?;
        }

        let bits = self.peek_bits(LUT_BITS);
        let lut_entry = table.lookup(bits as usize);
        if lut_entry.length != 0 {
            self.consume_bits(lut_entry.length);
            return Ok(lut_entry.symbol);
        }

        let mut code = 0;
        let mut first = 0;
        let mut index = 0;
        for len in 1..=16 {
            code = (code << 1) | self.get_bits(reader, 1)?;
            let count = table.bits[len - 1] as usize;
            if code as usize >= first && code as usize - first < count {
                let symbol = table.huffval[index + (code as usize - first)];
                return Ok(symbol);
            }
            index += count;
            first = (first + count) << 1;
        }

        Err(Error::MalformedHuffmanCode(code))
    }

    fn consume_bits(&mut self, count: u8) {
        self.buffer <<= count as usize;
        self.bit_count -= count as usize;
    }

    pub fn get_bits<R: Read>(&mut self, reader: &mut R, count: u8) -> Result<u16, Error> {
        if self.bit_count < 16 {
            self.read_bits(reader)?;
        }
        let bits = self.peek_bits(count);
        self.consume_bits(count);
        Ok(bits)
    }

    fn peek_bits(&mut self, count: u8) -> u16 {
        assert!(count <= 16);
        assert!(self.bit_count >= count as usize);
        ((self.buffer.wrapping_shr(64 - count as u32)) & ((1 << count) - 1)) as u16
    }

    fn read_bits<R: Read>(&mut self, reader: &mut R) -> Result<(), Error> {
        while self.bit_count <= 56 {
            let byte = match self.marker {
                Some(_) => 0,
                None => read_u8(reader)?,
            };

            if byte == 0xFF {
                let mut next_byte = read_u8(reader)?;

                if next_byte != 0x00 {
                    while next_byte == 0xFF {
                        next_byte = read_u8(reader)?;
                    }

                    match next_byte {
                        0x00 => todo!(),
                        _ => self.marker = Some(next_byte),
                    }

                    continue;
                }
            }

            self.buffer |= (byte as u64) << (56 - self.bit_count);
            self.bit_count += 8;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::jpeg::huffman::HuffmanDecoder;

    #[test]
    fn test_huffman_bit_read() {
        let mut huffman = HuffmanDecoder::default();
        let bitstring = b"\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34\x56\x78\x12\x34\x56\x78";
        let mut reader = &bitstring[..];
        let bits = huffman.get_bits(&mut reader, 3).unwrap();
        assert_eq!(bits, 0);
        let bits = huffman.get_bits(&mut reader, 2).unwrap();
        assert_eq!(bits, 0b10);
        let bits = huffman.get_bits(&mut reader, 2).unwrap();
        assert_eq!(bits, 0b01);
        let bits = huffman.get_bits(&mut reader, 1).unwrap();
        assert_eq!(bits, 0);
    }
}
