use std::collections::{self, HashMap};

use byte_encode_derive::ByteEncode;

use crate::{byte_encode::ByteEncode, crc::update_crc, image::{ColorType, Image, ImageBase}, reader::FileReader, zlib::{zlib_decode, ByteBuffer}};

#[derive(Default)]
pub struct PNGImage
{
	IHDR: Option<IHDR>,
	IDAT: Option<IDAT>,
	IEND: Option<IEND>,
	tEXt: Option<tEXt>,
}

trait Mergable: Sized
{
	fn merge(a: Self, b: Self) -> Self;
}

impl<T: Mergable + Sized + Clone> Mergable for Option<T>
{
	fn merge(a: Option<T>, b: Option<T>) -> Option<T>
	{
		if let Some(ref a_content) = a
		{
			if let Some(b_content) = b
			{
				Some(T::merge(a_content.clone(), b_content))
			}
			else
			{
				a
			}
		}
		else
		{
			b
		}
	}
}

trait PNGChunk: Mergable
{
	fn read(reader: &mut FileReader, length: usize) -> Self;
	const CHUNK_TYPE: [char; 4];
}

#[derive(ByteEncode, Clone)]
struct IHDR
{
	pub width: u32,
	pub height: u32,
	pub bit_depth: u8,
	pub color_type: u8,
	pub compression_method: u8,
	pub filter_method: u8,
	pub interlace_method: u8
}
impl PNGChunk for IHDR
{
	const CHUNK_TYPE: [char; 4] = ['I', 'H', 'D', 'R'];
    fn read(reader: &mut FileReader, length: usize) -> Self {
		reader.read()
    }
}
impl Mergable for IHDR {
    fn merge(a: Self, b: Self) -> Self {
        todo!()
    }
}
#[derive(ByteEncode, Clone)]
struct IEND
{
}
impl PNGChunk for IEND
{
	const CHUNK_TYPE: [char; 4] = ['I', 'E', 'N', 'D'];
    fn read(reader: &mut FileReader, length: usize) -> Self {
		Self {}
    }
}
impl Mergable for IEND {
    fn merge(a: Self, b: Self) -> Self {
        todo!()
    }
}

#[derive(Clone)]
struct tEXt
{
	keywords: HashMap<String, String>
}
impl PNGChunk for tEXt
{
	const CHUNK_TYPE: [char; 4] = ['t', 'E', 'X', 't'];
	fn read(reader: &mut FileReader, length: usize) -> Self
	{
		let mut i = 0;
		let mut keyword = String::new();
		let mut c: u8 = reader.read();
		i += 1;
		while c != 0
		{
			keyword.push(c as char);
			c = reader.read::<1, u8>();
			i += 1;
		}

		let mut text_string = String::new();
		while i < length
		{
			text_string.push(reader.read::<1, u8>() as char);
			i += 1;
		}

		let mut keywords = HashMap::new();
		keywords.insert(keyword, text_string);
		
		Self
		{
			keywords
		}
	}
}
impl Mergable for tEXt {
    fn merge(a: Self, b: Self) -> Self {
        let mut keywords = HashMap::new();
		keywords.extend(a.keywords);
		keywords.extend(b.keywords);
		Self
		{
			keywords
		}
    }
}

#[derive(Clone)]
struct IDAT
{
	data: Vec<u8>
}
impl PNGChunk for IDAT
{
    fn read(reader: &mut FileReader, length: usize) -> Self {
		let data = reader.get_bytes(length);
		Self
		{
			data
		}
    }

    const CHUNK_TYPE: [char; 4] = ['I', 'D', 'A', 'T'];
}
impl Mergable for IDAT
{
	fn merge(a: Self, b: Self) -> Self {
        let mut data = Vec::new();
		data.extend(a.data);
		data.extend(b.data);
		Self
		{
			data
		}
    }
}

impl<T: ColorType> Image<T> for PNGImage
{
    fn read_image(reader: &mut FileReader) -> Result<ImageBase<T>, String> {
		let magic: [u8; 8] = reader.read_array();
		if magic != [0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]
		{
			return Err("Not a PNG!".to_owned());
		}

		reader.set_endianness(false);
		reader.set_checksum_handler(update_crc, 0xffffffffu32);

		let mut p = PNGImage::default();

		while p.IEND.is_none()
		{
			let length: usize = reader.read::<4, u32>() as usize;
			reader.reset_checksum();
			let chunk_type: [char; 4] = reader.read_array::<1, u8, 4>().map(|x| x as char);
			let chunk_flags = chunk_type.map(|x| x.is_lowercase());
			println!("type: {chunk_type:?}");
			println!("length: {length}");
			match chunk_type
			{
				IHDR::CHUNK_TYPE => { p.IHDR = Option::<IHDR>::merge(p.IHDR, Some(IHDR::read(reader, length))) }
				IDAT::CHUNK_TYPE => { p.IDAT = Option::<IDAT>::merge(p.IDAT, Some(IDAT::read(reader, length))) }
				IEND::CHUNK_TYPE => { p.IEND = Option::<IEND>::merge(p.IEND, Some(IEND::read(reader, length))) }
				tEXt::CHUNK_TYPE => { p.tEXt = Option::<tEXt>::merge(p.tEXt, Some(tEXt::read(reader, length))) }
				_ => {
					if chunk_flags[0]
					{
						println!("Skipping ancillary chunk");
						println!("{}", if chunk_flags[1] { "Private (can find definition)" } else { "Public" })
					}
					else
					{
						println!("Skipping important chunk!!!");
					}
					reader.skip(length)
				}
			}
			let mut calculated_crc = reader.get_checksum();
			let crc: u32 = reader.read();
			if crc ^ calculated_crc != 0xffffffffu32 && !chunk_flags[0]
			{
				calculated_crc ^= 0xffffffffu32;
				println!("Bad crc: {crc} {calculated_crc}");
			}
		}

		if let Some(ref tEXt) = p.tEXt
		{
			for x in tEXt.keywords.iter()
			{
				println!("{0}: {1}", x.0, x.1);
			}
		}

		if p.IEND.is_none()
		{
			return Err("IEND not present".to_owned());
		}

		let Some(IHDR) = p.IHDR else { return Err("IHDR not present".to_owned()); };
		let Some(IDAT) = p.IDAT else { return Err("IDAT not present".to_owned()); };

		let colors_per_pixel: u8 = if IHDR.color_type == 4 || IHDR.color_type == 6 { 4 } else { 3 };
		let bpp: u8 = IHDR.bit_depth * colors_per_pixel;
		let width: u32 = IHDR.width;
		let height: u32 = IHDR.height;
		println!("{bpp}, {width}, {height}");

		let compression_method: u8 = IDAT.data[0];
		let additional_flags: u8 = IDAT.data[1];

		let decoded = zlib_decode(ByteBuffer::new(&IDAT.data[2..]));
		
        todo!("PNG not finished")
    }

    fn write_image(image: &crate::image::ImageBase<T>, writer: &mut crate::writer::FileWriter)-> () {
        todo!()
    }
}