const MAX_BITS: usize = 15;
const CODE_TREE_LEN: usize = 288;
const DIST_TREE_LEN: usize = 32;

const LEN_START: [usize; 29] = [3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195, 227, 258];
const LEN_EXTRA: [usize; 29] = [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0];

const DIST_START: [usize; 30] = [1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577];
const DIST_EXTRA: [usize; 30] = [0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13];

#[derive(Clone)]
struct HuffmanTree<const LENGTH: usize>
{
	pub len_counts: [u8; MAX_BITS + 1],
	pub ordered_codes: [u16; LENGTH]
}

impl<const LENGTH: usize> HuffmanTree<LENGTH>
{
    const fn new(code_lens: [u8; LENGTH]) -> Self
	{
		let mut len_counts = [0u8; MAX_BITS + 1];
		let mut ordered_codes = [0u16; LENGTH];

		let mut code: usize = 0;
		while code < LENGTH
		{
			len_counts[code_lens[code] as usize] += 1;
			code += 1
		}
		
		let mut offsets = [0usize; MAX_BITS + 1];
		let mut len: usize = 1;
		while len < MAX_BITS
		{
			offsets[len + 1]  = offsets[len] + len_counts[len] as usize;
			len += 1;
		}

		let mut code = 0;
		while code < LENGTH
		{
			if code_lens[code] != 0
			{
				ordered_codes[offsets[code_lens[code] as usize]] = code as u16;
				offsets[code_lens[code] as usize] += 1;
			}
			code += 1;
		}
		
		Self
		{
			len_counts,
			ordered_codes
		}
    }

	

	fn decode<T: Bitstream>(&self, bitstream: &mut T) -> u16
	{
		let mut code: u16 = 0;
		let mut first: u16 = 0;
		let mut index: usize = 0;
		
		for len in 1..MAX_BITS
		{
			code |= bitstream.next_bit() as u16;
			let count: u16 = self.len_counts[len] as u16;
			if code - count < first
			{
				return self.ordered_codes[index + ((code - first) as usize)];
			}
			index += count as usize;
			first += count;
			first <<= 1;
			code <<= 1;
		}
		
		panic!("No codes left")
	}
	
	const fn get_static_code_lens() -> [u8; CODE_TREE_LEN]
	{
		let mut code_lens = [0u8; CODE_TREE_LEN];
		let mut i = 0;
		while i < CODE_TREE_LEN
		{
			code_lens[i] = match i
			{
				0..144 => 8,
				144..256 => 9,
				256..280 => 7,
				280..CODE_TREE_LEN => 8,
				_ => panic!()
			};
			i += 1
		}
		code_lens
	}
	const fn get_static_dist_code_lens() -> [u8; DIST_TREE_LEN]
	{
		[5u8; DIST_TREE_LEN]
	}
}

static STATIC_HUFFMAN_TREE: HuffmanTree<CODE_TREE_LEN> = HuffmanTree::new(HuffmanTree::<CODE_TREE_LEN>::get_static_code_lens());
static STATIC_HUFFMAN_DIST_TREE: HuffmanTree<DIST_TREE_LEN> = HuffmanTree::new(HuffmanTree::<DIST_TREE_LEN>::get_static_dist_code_lens());

#[derive(PartialEq)]
enum CompressionType
{
	None = 0,
	StaticCodes = 1,
	DynamicCodes = 2,
	Reserved = 3,
}

impl CompressionType
{
	fn new(x: u8) -> Self
	{
		match x
		{
			0 => Self::None,
			1 => Self::StaticCodes,
			2 => Self::DynamicCodes,
			3 => Self::Reserved,
			_ => panic!("Compression type should not be greater than 4")
		}
	}
}

pub struct ByteBuffer<'a>
{
	buffer: &'a [u8],
	pos: usize,
}

impl<'a> ByteBuffer<'a>
{
	pub fn new(buffer: &'a [u8]) -> Self
	{
		Self
		{
			buffer,
			pos: 0
		}
	}
}

pub trait Bitstream {
	fn next_bit(&mut self) -> bool;

	fn next_bits(&mut self, c: usize) -> u16
	{
		if c > 16
		{
			panic!("Cannot fetch more than 16 bits at a time");
		}
		let mut bits = 0u16;

		for i in 0..c
		{
			bits |= (self.next_bit() as u16) << i;
		}

		bits
	}

	fn skip_to_byte_boundary(&mut self);
}

impl<'a> Bitstream for ByteBuffer<'a>
{
	fn next_bit(&mut self) -> bool
	{
		let bit = (self.buffer[self.pos/8 as usize] & 1 << self.pos%8) != 0;
		self.pos += 1;
		bit
	}

	fn skip_to_byte_boundary(&mut self)
	{
		// println!("{}", (8 - self.pos%8)%8);
		self.pos += (8 - self.pos%8)%8;
    }
}

pub fn zlib_decode<T: Bitstream>(mut instream: T) -> Vec<u8>
{
	let mut output_buffer: Vec<u8> = Vec::new();
	// let mut output_pos: usize = 0;

	let mut final_chunk = false;
	while !final_chunk
	{		
		final_chunk = instream.next_bit();
		let compression_type = CompressionType::new(instream.next_bits(2) as u8);
		if compression_type == CompressionType::Reserved
		{
			panic!("Reserved compression type should not be used");
		}
		else if compression_type == CompressionType::None
		{
			instream.skip_to_byte_boundary();
			let len = instream.next_bits(16);
			let nlen = instream.next_bits(16);
			if len != !nlen
			{
				panic!("len and nlen don't match");
			}
			for _ in 0..len
			{
				output_buffer.push(instream.next_bits(8) as u8);
			}
			// todo!("No compression");
		}
		else
		{
			let (code_tree, dist_tree) =
				if compression_type == CompressionType::StaticCodes
			{
				(STATIC_HUFFMAN_TREE.clone(), STATIC_HUFFMAN_DIST_TREE.clone())
			}
			else // Dynamic codes
			{
				todo!("Dynamic codes")
			};
			

			let mut code = 0u16;
			while code != 256
			{
				code = code_tree.decode(&mut instream);
				if code < 256
				{
					output_buffer.push(code as u8);
				}
				else if code > 256
				{
					let len = LEN_START[code as usize - 257] + instream.next_bits(LEN_EXTRA[code as usize - 257]) as usize;

					let dist_code = dist_tree.decode(&mut instream);
					let dist = DIST_START[dist_code as usize] + instream.next_bits(DIST_EXTRA[dist_code as usize]) as usize;

					let out_pos = output_buffer.len() - 1;
					for _ in 0..len
					{
						output_buffer.push(output_buffer[out_pos - dist]);
					}
				}
			}
		}
	}
	
	output_buffer
}