use byteorder::{LittleEndian, ReadBytesExt};
use image::{
    buffer, ColorType, DynamicImage, GenericImage, GenericImageView, RgbImage, Rgba, RgbaImage,
};
use std::{
    collections::BTreeMap,
    io::{Read, Seek, Write},
};

use crate::binio::BitIO;
use crate::common::{CzError, CzHeader};

#[derive(Debug, Clone, Copy)]
pub struct ChunkInfo {
    pub size_compressed: usize,
    pub size_raw: usize,
}

#[derive(Debug, Clone)]
pub struct CompressionInfo {
    pub chunk_count: usize,
    pub total_size_compressed: usize,
    pub total_size_raw: usize,
    pub chunks: Vec<ChunkInfo>,

    /// Length of the compression chunk info
    pub length: usize,
}

/// Get info about the compression chunks
pub fn parse_chunk_info<T: Seek + ReadBytesExt + Read>(
    bytes: &mut T,
) -> Result<CompressionInfo, CzError> {
    let parts_count = bytes.read_u32::<LittleEndian>()?;

    let mut part_sizes = vec![];
    let mut total_size = 0;
    let mut total_size_raw = 0;

    // Loop over the compressed bytes
    for _ in 0..parts_count {
        let compressed_size = bytes.read_u32::<LittleEndian>()?;
        total_size += compressed_size;

        let raw_size = bytes.read_u32::<LittleEndian>()?;
        total_size_raw += raw_size;

        part_sizes.push(ChunkInfo {
            size_compressed: compressed_size as usize,
            size_raw: raw_size as usize,
        });
    }

    Ok(CompressionInfo {
        chunk_count: parts_count as usize,
        total_size_compressed: total_size as usize,
        total_size_raw: total_size_raw as usize,
        chunks: part_sizes,
        length: bytes.stream_position()? as usize,
    })
}

/// Decompress an LZW compressed stream like CZ1
pub fn decompress<T: Seek + ReadBytesExt + Read>(
    input: &mut T,
    chunk_info: &CompressionInfo,
) -> Result<Vec<u8>, CzError> {
    let mut m_dst = 0;
    let mut bitmap = vec![0; chunk_info.total_size_raw];
    for chunk in &chunk_info.chunks {
        let mut part = vec![0u8; chunk.size_compressed * 2];
        input.read_exact(&mut part)?;

        for j in (0..part.len()).step_by(2) {
            let ctl = part[j + 1];

            if ctl == 0 {
                bitmap[m_dst] = part[j];
                m_dst += 1;
            } else {
                m_dst += copy_range(&mut bitmap, &part, get_offset(&part, j), m_dst);
            }
        }
    }

    bitmap.truncate(chunk_info.total_size_raw);

    Ok(bitmap)
}

/// Decompress an LZW compressed stream like CZ2
pub fn decompress_2<T: Seek + ReadBytesExt + Read>(
    input: &mut T,
    chunk_info: &CompressionInfo,
) -> Result<Vec<u8>, CzError> {
    let mut output_buf: Vec<u8> = vec![];

    for block in &chunk_info.chunks {
        let mut buffer = vec![0u8; block.size_compressed];
        input.read_exact(&mut buffer).unwrap();

        let raw_buf = decompress_lzw2(&buffer, block.size_raw);

        output_buf.write_all(&raw_buf).unwrap();
    }

    Ok(output_buf)
}

pub fn decompress_lzw2(input_data: &[u8], size: usize) -> Vec<u8> {
    let mut data = input_data.to_vec();
    data[0..2].copy_from_slice(&[0, 0]);
    let mut dictionary = BTreeMap::new();
    for i in 0..256 {
        dictionary.insert(i as u64, vec![i as u8]);
    }
    let mut dictionary_count = dictionary.len() as u64;
    let mut result = Vec::with_capacity(size);

    let data_size = input_data.len();
    data.extend_from_slice(&[0, 0]);
    let mut bit_io = BitIO::new(data);
    let mut w = dictionary.get(&0).unwrap().clone();

    let mut element;
    loop {
        let flag = bit_io.read_bit(1);
        if flag == 0 {
            element = bit_io.read_bit(15);
        } else {
            element = bit_io.read_bit(18);
        }

        if bit_io.byte_offset() > data_size {
            break;
        }

        let mut entry;
        if let Some(x) = dictionary.get(&element) {
            // If the element was already in the dict, get it
            entry = x.clone()
        } else if element == dictionary_count {
            entry = w.clone();
            entry.push(w[0])
        } else {
            panic!("Bad compressed element: {}", element)
        }

        //println!("{}", element);

        result.write(&entry).unwrap();
        w.push(entry[0]);
        dictionary.insert(dictionary_count, w.clone());
        dictionary_count += 1;
        w = entry.clone();
    }
    result
}

fn get_offset(input: &[u8], src: usize) -> usize {
    (((input[src] as usize) | (input[src + 1] as usize) << 8) - 0x101) * 2
}

fn copy_range(bitmap: &mut Vec<u8>, input: &[u8], src: usize, dst: usize) -> usize {
    let mut dst = dst;
    let start_pos = dst;

    if input[src + 1] == 0 {
        bitmap[dst] = input[src];
        dst += 1;
    } else if get_offset(input, src) == src {
        bitmap[dst] = 0;
        dst += 1;
    } else {
        dst += copy_range(bitmap, input, get_offset(input, src), dst);
    }

    if input[src + 3] == 0 {
        bitmap[dst] = input[src + 2];
        dst += 1;
    } else if get_offset(input, src + 2) == src {
        bitmap[dst] = bitmap[start_pos];
        dst += 1;
    } else {
        bitmap[dst] = copy_one(input, get_offset(input, src + 2));
        dst += 1;
    }

    dst - start_pos
}

fn copy_one(input: &[u8], src: usize) -> u8 {
    if input[src + 1] == 0 {
        input[src]
    } else if get_offset(input, src) == src {
        0
    } else {
        copy_one(input, get_offset(input, src))
    }
}

pub fn line_diff<T: CzHeader>(header: &T, data: &[u8]) -> Vec<u8> {
    let width = header.width() as u32;
    let height = header.height() as u32;
    let mut output_buf = data.to_vec();

    let block_height =
        (f32::ceil(height as f32 / header.color_block() as f32) as u16) as usize;
    let pixel_byte_count = header.depth() >> 3;
    let line_byte_count = (width * pixel_byte_count as u32) as usize;

    let mut curr_line: Vec<u8>;
    let mut prev_line: Vec<u8> = Vec::with_capacity(line_byte_count);

    let mut i = 0;
    for y in 0..height {
        curr_line = data[i..i + line_byte_count].to_vec();

        if y % block_height as u32 != 0 {
            for x in 0..line_byte_count {
                curr_line[x] += prev_line[x]
            }
        }

        prev_line = curr_line.clone();
        if pixel_byte_count == 4 {
            output_buf[i..i + line_byte_count].copy_from_slice(&curr_line);
        } else if pixel_byte_count == 3 {
            for x in 0..line_byte_count {
                let loc = ((y * width) as usize + x) * 4;

                output_buf[loc..loc + 4].copy_from_slice(&[
                    curr_line[x + 0],
                    curr_line[x + 1],
                    curr_line[x + 2],
                    0xFF
                ])
            }
        }

        i += line_byte_count;
    }

    output_buf
}

pub fn line_diff_cz4<T: CzHeader>(header: &T, data: &[u8]) -> Vec<u8> {
    let width = header.width() as u32;
    let height = header.height() as u32;
    let block_height = (f32::ceil(height as f32 / 3.0) as u16) as u32;

    //let pixel_byte_count = (header.depth() >> 3) as usize;
    let pixel_byte_count = 3;

    let mut output_buf = data.to_vec();

    let mut curr_line;
    let mut prev_line = vec![0u8; width as usize * pixel_byte_count];

    let mut i = 0;
    for y in 0..height {
        curr_line = data[i..i + width as usize * pixel_byte_count].to_vec();

        if y % block_height != 0 {
            for x in 0..(width as usize * pixel_byte_count) {
                curr_line[x] = u8::wrapping_add(curr_line[x], prev_line[x])
            }
        }

        for x in 0..width as usize {
            let loc = ((y * width) as usize + x) * 4;

            if pixel_byte_count == 1 {
                output_buf[loc + 4] = curr_line[x];
            } else if pixel_byte_count == 4 {
                output_buf[loc..loc + 4].copy_from_slice(&[
                    curr_line[x * pixel_byte_count + 0],
                    curr_line[x * pixel_byte_count + 1],
                    curr_line[x * pixel_byte_count + 2],
                    curr_line[x * pixel_byte_count + 3],
                ]);
            } else if pixel_byte_count == 3 {
                output_buf[loc..loc + 4].copy_from_slice(&[
                    curr_line[x * pixel_byte_count + 0],
                    curr_line[x * pixel_byte_count + 1],
                    curr_line[x * pixel_byte_count + 2],
                    0xFF,
                ]);
            }
        }

        prev_line = curr_line.clone();
        i += width as usize * pixel_byte_count;
    }

    output_buf
}