tato_pipe 0.1.3

use std::{
    any::type_name,
    collections::HashMap,
    fs::write,
    marker::PhantomData
};

use crate::*;

/// Facilitates creating tilesets with consisten color palettes and resources like  Animations, Fonts and Groups.
/// After creating the AtlasBuilder and inserting each resource (normally using a .png file as input) you can save individual tilesets as '.tile' binary files.
// #[derive(Debug)]
pub struct AtlasBuilder<T, P, G>
where T:TilesetEnum, P:PaletteEnum, G:GroupEnum,
{
    tilesets: Vec<Option<TilesetBuilder>>,
    palettes: Vec<Option<Palette>>,
    palette_hash: Vec<HashMap<Color32, u8>>,
    specs:Specs,
    tileset_marker: PhantomData<T>,
    palette_marker: PhantomData<P>,
    group_marker: PhantomData<G>,
}

impl<T, P, G> AtlasBuilder<T, P, G>
where T:TilesetEnum, P:PaletteEnum, G:GroupEnum,
{
    
    pub fn new(specs:Specs) -> Self {
        
        Self {
            palette_hash: (0 .. P::count()).map(|_| HashMap::new() ).collect(),
            palettes: (0 .. P::count()).map(|_| None ).collect(),
            tilesets: (0 .. T::count()).map(|_| None ).collect(),
            specs,
            tileset_marker: Default::default(),
            palette_marker: Default::default(),
            group_marker: Default::default()
        }
    }


    pub fn init_tileset(&mut self, tileset_id:T, palette_id:P) {
        println!("cargo:warning=Init tileset '{}'", type_name::<T>());
        let palette_id:u8 = palette_id.into();
        let tileset_id:u8 = tileset_id.into();

        if tileset_id as usize > self.tilesets.len() {
            panic!("AtlasBuilder: Error, tileset index {} above count of '{}'", tileset_id, type_name::<T>())
        }

        if let Some(entry) = self.tilesets.get(tileset_id as usize){
            if entry.is_none(){
                println!("cargo:warning=AtlasBuilder: initializing tileset at index {}.", tileset_id);
                self.tilesets[tileset_id as usize] = Some(
                    TilesetBuilder::new(self.specs, P::from(palette_id))
                );
            } else {
                panic!("AtlasBuilder: Error, tileset {} already initialized", tileset_id)
            }
        } else {
            panic!("AtlasBuilder: Error, invalid tileset index {}.", tileset_id)
        }
        
        if palette_id as usize > self.palettes.len() {
            panic!("AtlasBuilder: Error, palette index {} above capacity of {}", tileset_id, type_name::<P>())
        }

        if let Some(entry) = self.palettes.get(tileset_id as usize){
            if entry.is_none() {
                self.palettes[palette_id as usize] = Some(
                    Palette::new(self.specs, palette_id)
                )
            }
        }   
    }


    // Helper function
    fn get_data(&mut self, tileset_id:impl TilesetEnum) -> Option<(&mut TilesetBuilder, &mut Palette, &mut HashMap<Color32, u8>)> {
        let tileset_id:u8 = tileset_id.into();
        if let Some(tileset) = self.tilesets.get_mut(tileset_id as usize)? {
            let palette_id:u8 = tileset.palette_id;
            let palette_option = self.palettes.get_mut(palette_id as usize)?;
            let Some(palette) = palette_option.as_mut() else { panic!("Invalid palette for tileset {}: not initialized.", tileset_id) };
            let hash = &mut self.palette_hash[palette_id as usize];
            Some((tileset, palette, hash))
        } else {
            None
        }
    }


    pub fn init_group(&mut self, path:&str, tileset_key:impl TilesetEnum, group_id:impl GroupEnum, collider:bool) {
        let specs = self.specs;
        let Some((tileset, palette, palette_hash)) = self.get_data(tileset_key) else { return };
        let img = ImageBuilder::from_image(specs, path, None, palette, palette_hash);
        let group_id:u8 = group_id.into();
        tileset.add_tiles(&img, group_id, collider);
    }


    pub fn init_font(&mut self, path:&str, tileset_id:impl TilesetEnum, group_id:impl GroupEnum) {
        let specs = self.specs;
        let group_id:u8 = group_id.into();

        let Some((tileset, palette, palette_hash)) = self.get_data(tileset_id) else {
            let tileset_id:u8 = tileset_id.into();
            panic!("AtlasBuilder Error: Invalid tileset ID {}.", tileset_id)
        };
        
        if tileset.anims.len() == u8::MAX as usize {
            panic!("AtlasBuilder Error: Max anim capacity of {} exceeded.", u8::MAX)
        }

        let font_id:u8 = tileset.fonts.len().try_into().unwrap();
        let start_index = tileset.next_tile;
        let img = ImageBuilder::from_image(specs, path, None, palette, palette_hash); //TODO: Move palette hash to atlas builder?
        tileset.add_tiles(&img, group_id, false);
        let len = tileset.tile_count;

        tileset.fonts.push( Font {
            start_index,
            len,
            id: font_id,
            tileset_id: Default::default(),
        });

        tileset.font_names.push( strip_path_name(path) );
    }
    

    pub fn init_anim(&mut self, path:&str, fps:u8, frames_h:u8, frames_v:u8, tileset_id:impl TilesetEnum, group_id:impl GroupEnum ) {
        let specs = self.specs;
        
        let Some((tileset, palette, palette_hash)) = self.get_data(tileset_id) else {
            let tileset_id:u8 = tileset_id.into();
            panic!("AtlasBuilder Error: Invalid tileset ID {}.", tileset_id)
        };
        
        if tileset.anims.len() == u8::MAX as usize {
            panic!("AtlasBuilder Error: Max anim capacity of {} exceeded.", u8::MAX)
        }

        let anim_id:u8 = tileset.anims.len().try_into().unwrap();
        let group_id:u8 = group_id.into();
        let palette_id:u8 = tileset.palette_id;

        let group:u8 = group_id;
        let img = ImageBuilder::from_image(specs, path, Some((frames_h, frames_v)), palette, palette_hash);
        let tiles = tileset.add_tiles(&img, group, false);
        let frame_len = img.cols_per_frame as usize * img.rows_per_frame as usize;
        let frame_count = u8::try_from(tiles.len() / frame_len).unwrap();
        let len = u8::try_from(tiles.len() / frame_len).ok().unwrap();

        tileset.anims.push( Anim {
            frames: core::array::from_fn(|n|{
                let index = n * frame_len;
                if n < frame_count as usize {
                    Frame::from_slice(&tiles[index .. index+frame_len], img.cols_per_frame, img.rows_per_frame)
                } else {
                    Frame::default() 
                }
            }),
            len,
            group,
            fps,
            id: anim_id,
            tileset: tileset_id.into(),
            palette: palette_id
        });

        tileset.anim_names.push( strip_path_name(path) );

    }


    pub fn init_tilemap(&mut self, path:&str, tileset_id:impl TilesetEnum, group_id:impl GroupEnum) {
        let specs = self.specs;
        

        let Some((tileset, palette, palette_hash)) = self.get_data(tileset_id) else {
            let tileset_id:u8 = tileset_id.into();
            panic!("AtlasBuilder Error: Invalid tileset ID {}.", tileset_id)
        };

        if tileset.tilemaps.len() == u8::MAX as usize {
            panic!("AtlasBuilder Error: Max tilemap capacity of {} exceeded.", u8::MAX)
        }

        let img = ImageBuilder::from_image(specs, path, None, palette, palette_hash);
        let tiles = tileset.add_tiles(&img, group_id.into(), false);
        let cols = u16::try_from(img.width / specs.tile_width as usize).unwrap();
        let rows = u16::try_from(img.height / specs.tile_height as usize).unwrap();
        let map_id:u8 = tileset.tilemaps.len().try_into().unwrap();

        tileset.tilemaps.push( Tilemap{
            tiles: core::array::from_fn(|i|{
                *tiles.get(i).unwrap_or( &Tile::default() )
            }),
            id: map_id,
            cols,
            rows,
            tileset: tileset_id.into(),
            palette: tileset.palette_id,
            bg_buffers: Default::default(),
        });

        tileset.tilemap_names.push( strip_path_name(path) );
    }


    pub fn save(&self, path:&str) {
        println!("cargo:warning=Saving Atlas");
        let mut data:Vec<u8> = vec![];
        let mut anim_toc:HashMap<String, u8> = HashMap::new();
        let mut font_toc:HashMap<String, u8> = HashMap::new();
        let mut tilemap_toc:HashMap<String, u8> = HashMap::new();

        // Header text
        for letter in ATLAS_HEADER_TEXT.as_bytes() {
            data.push(*letter)
        }

        // Header data
        data.push( self.specs.tile_width );                             // Tile width
        data.push( self.specs.tile_height );                            // Tile Height
        data.push( u8::try_from(self.palettes.len()).ok().unwrap() );   // Palette Count
        data.push( u8::try_from(self.tilesets.len()).ok().unwrap() );   // Tileset Count

        // Palettes
        for (i,palette) in self.palettes.iter().enumerate() {
            println!("cargo:warning=    Saving palette {}...", i);
            let Some(ref palette) = palette else {
                panic!("AtlasBuilder Error: Palette {} not initialized", i)
            };

            let palette_id = u8::try_from(i).unwrap();
            data.push(palette_id);                           // Palette ID
            for color in palette.colors() {                  // Each Color24
                let color_data = color.serialize();
                data.extend_from_slice(&color_data);
            }
        }

        // Tilesets
        for (t,tileset) in self.tilesets.iter().enumerate() {
            println!("cargo:warning=    Saving tileset {}...", t);
            let Some(ref tileset) = tileset else {
                panic!("AtlasBuilder Error: Tileset {} not initialized", t)
            };

            // Header text
            for letter in TILESET_HEADER_TEXT.chars() { 
                data.push( letter as u8 )
            }

            // Counts
            data.push( self.specs.atlas_width.to_ne_bytes()[0]);
            data.push( self.specs.atlas_width.to_ne_bytes()[1]);                // Atlas Width
            data.push( tileset.pixels.len().to_ne_bytes()[0]);                  
            data.push( tileset.pixels.len().to_ne_bytes()[1] );                 // Pixel Count
            data.push( tileset.tile_count );                                    // Tile Count
            data.push( u8::try_from(tileset.fonts.len()).ok().unwrap() );       // Font Count
            data.push( u8::try_from(tileset.anims.len()).ok().unwrap() );       // Anim Count
            data.push( u8::try_from(tileset.tilemaps.len()).ok().unwrap() );    // Maps Count

            // Debug palette
            data.push( tileset.palette_id );                            // Palette

            // Pixels
            println!("cargo:warning=        Saving {} pixels...", tileset.pixels.len());
            for pixel in tileset.pixels.iter(){
                // print!("cargo:warning= {}, ", *pixel);
                data.push(*pixel);
            }
            // println!("cargo:warning=        ");

            // Fonts
            data.extend_from_slice("fonts".as_bytes());
            for (i,font) in tileset.fonts.iter().enumerate() {
                println!("cargo:warning=        Saving font {}...", i);
                let font_data = font.serialize();
                data.extend_from_slice(&font_data);
                font_toc.insert( tileset.font_names[i].clone(), font.id );
            }

            // Anims
            data.extend_from_slice("anims".as_bytes());
            for (i,anim) in tileset.anims.iter().enumerate() {
                println!("cargo:warning=        Saving anim {}... ", i);
                let anim_data = anim.serialize();
                data.extend_from_slice(&anim_data);
                anim_toc.insert( tileset.anim_names[i].clone(), anim.id );
            }

            // Tilemaps
            data.extend_from_slice("tilemaps".as_bytes());
            for (i,tilemap) in tileset.tilemaps.iter().enumerate() {
                println!("cargo:warning=        Saving tilemap {}...", i);
                let map_data = tilemap.serialize();
                data.extend_from_slice(&map_data);
                tilemap_toc.insert( tileset.tilemap_names[i].clone(), tilemap.id );
            }
        }

        // Write data
        println!("cargo:warning=    Finished atlas at pos={}!", data.len());
        write(path, data.as_slice()).unwrap();

        fn write_toc(path:&str, toc:HashMap<String, u8>){
            let mut toc_data:Vec<u8> = vec![];
            for (key, value) in toc.iter() {
                for byte in key.as_bytes(){
                    toc_data.push(*byte);
                }
                toc_data.push(*value);
            }
            write(path, toc_data.as_slice()).unwrap();
        }

        write_toc( format!("{}{}", path, ".anims").as_str(), anim_toc );
        write_toc( format!("{}{}", path, ".fonts").as_str(), font_toc );
        write_toc( format!("{}{}", path, ".maps").as_str(), tilemap_toc );
    }

}



fn strip_path_name(path:&str) -> String {
    let split = path.split('/');
    let file_name = split.last().unwrap();
    let mut file_name_split = file_name.split('.');
    file_name_split.next().unwrap().to_string()
}