#![allow(dead_code)]
use std::{
    collections::HashMap,
    sync::{Arc, Weak},
};

use ash::vk::Extent2D;
use cosmic_text::{CacheKey, FontSystem, PhysicalGlyph, SwashCache};
use glam::IVec2;
use guillotiere::size2;
#[cfg(test)]
use image::{GenericImage, GenericImageView};

use crate::{
    def_monotonic_id,
    util::{self, F32, Rect2D},
};

const ROBOTO_BYTES: &[u8] =
    include_bytes!("../../../assets/fonts/Roboto/Roboto-VariableFont_wdth,wght.ttf");
const NOTO_SANS_HAN_BYTES: &[u8] =
    include_bytes!("../../../assets/fonts/Noto_Sans_SC/NotoSansSC-VariableFont_wght.ttf");

def_monotonic_id!(pub FontId);

type FontData = Arc<dyn AsRef<[u8]> + Send + Sync>;
struct FontStore {
    fonts: HashMap<FontId, FontData>,
}

impl FontStore {
    fn new() -> Self {
        Self {
            fonts: HashMap::new(),
        }
    }

    fn add_font_bytes(&mut self, bytes: FontData) -> FontId {
        let id = FontId::new();
        self.fonts.insert(id, bytes);
        id
    }

    fn as_source(&self, id: FontId) -> Option<cosmic_text::fontdb::Source> {
        self.fonts
            .get(&id)
            .map(|bytes| cosmic_text::fontdb::Source::Binary(Arc::clone(bytes)))
    }
}

#[derive(Debug, Default)]
struct FontAtlasSets {
    sets: HashMap<FontId, FontAtlasSet>,
}

impl FontAtlasSets {
    fn new() -> Self {
        Self {
            sets: HashMap::new(),
        }
    }
    fn get(&self, key: &FontId) -> Option<&FontAtlasSet> {
        self.sets.get(key)
    }
    fn get_mut(&mut self, key: &FontId) -> Option<&mut FontAtlasSet> {
        self.sets.get_mut(key)
    }
}

/// Per-Font Font Atlas Set
#[derive(Debug, Default)]
struct FontAtlasSet {
    // TODO: add proper plural of atlas to english language.
    atlantes: HashMap<util::F32, Vec<FontAtlas>>,
}

impl FontAtlasSet {
    fn new() -> Self {
        Self {
            atlantes: HashMap::new(),
        }
    }
    fn get_glyph_info(&self, key: CacheKey) -> Option<AtlasGlyphInfo> {
        self.atlantes
            .get(&F32::from_bits(key.font_size_bits))?
            .iter()
            .find_map(|atlas| {
                atlas
                    .glyphs
                    .get(&key)
                    .map(|&(rect, offset)| (Arc::downgrade(&atlas.image), rect, offset))
            })
            .map(|(image, rect, offset)| AtlasGlyphInfo {
                image,
                rect,
                offset,
            })
    }
    fn add_glyph(
        &mut self,
        font_system: &mut FontSystem,
        swash_cache: &mut SwashCache,
        physical: PhysicalGlyph,
    ) -> Result<AtlasGlyphInfo, Error> {
        let key = physical.cache_key;

        let atlantes = self
            .atlantes
            .entry(F32::from_bits(physical.cache_key.font_size_bits))
            .or_insert_with(|| vec![FontAtlas::new(512)]);

        let (data, size, offset) = get_outlined_glyph_texture(font_system, swash_cache, physical)?;

        if !atlantes
            .iter_mut()
            .any(|atlas| atlas.add_glyph(key, &data, size, offset).is_some())
        {
            let max_size = size.height.max(size.height);
            let x2_or_512 = (1u32 << (32 - max_size.leading_zeros())).max(512);
            atlantes.push(FontAtlas::new(x2_or_512));
            atlantes
                .last_mut()
                .unwrap()
                .add_glyph(key, &data, size, offset)
                .ok_or(Error::FailedToRasterizeGlyph(key))?;
        }

        Ok(self.get_glyph_info(key).unwrap())
    }
}

#[derive(Debug)]
struct Image(parking_lot::RwLock<ImageInner>);

#[derive(Debug, Clone)]
struct ImageInner {
    image: Vec<u8>,
    image_size: Extent2D,
}

impl ImageInner {
    fn bytes(&self) -> &[u8] {
        self.image.as_slice()
    }
    fn bytes_mut(&mut self) -> &mut [u8] {
        self.image.as_mut_slice()
    }
    fn width(&self) -> u32 {
        self.image_size.width
    }
    fn height(&self) -> u32 {
        self.image_size.height
    }
}

#[derive(Debug)]
struct AtlasGlyphInfo {
    image: Weak<Image>,
    rect: Rect2D,
    offset: IVec2,
}

struct FontAtlas {
    // TODO: this image will be host-coherent and host-visible, so that it can
    // be updated with new glyphs.
    // that begs the question:
    // - should it be staged to a device-local image when it's used?
    // - does it make sense to use VK_EXT_external_memory_host here? It's
    // supported on virtually all Windows and Linux drivers.
    // - how to sync this? I need to make sure that when this image is written
    // to, it isn't also being read from. Usually, these images will be
    // write-only, except when rendering. Since currently my rendering may
    // happen on any thread at any time, that's problematic.
    //
    // In fact, this is an awful type to use here because of the unique access
    // requirement for mapping.
    image: Arc<Image>,
    // stores sub-rect of image and placement offset of glyph
    glyphs: HashMap<CacheKey, (Rect2D, IVec2)>,
    allocator: guillotiere::AtlasAllocator,
}

impl std::fmt::Debug for FontAtlas {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FontAtlas")
            .field("image", &self.image)
            .field("glyphs", &self.glyphs)
            .finish_non_exhaustive()
    }
}

impl FontAtlas {
    fn new(size: u32) -> Self {
        let num_bytes = size * size * 4;
        Self {
            image: Arc::new(Image(parking_lot::RwLock::new(ImageInner {
                image: vec![0; num_bytes as usize],
                image_size: Extent2D {
                    width: size,
                    height: size,
                },
            }))),
            glyphs: HashMap::new(),
            allocator: guillotiere::AtlasAllocator::new(size2(
                size.try_into().unwrap(),
                size.try_into().unwrap(),
            )),
        }
    }
    fn has_glyph(&self, key: CacheKey) -> bool {
        self.glyphs.contains_key(&key)
    }
    fn add_glyph(
        &mut self,
        key: CacheKey,
        data: &[u8],
        size: Extent2D,
        offset: IVec2,
    ) -> Option<AtlasGlyphInfo> {
        let allocation = self.allocator.allocate(guillotiere::size2(
            (size.width + 1).try_into().unwrap(),
            (size.height + 1).try_into().unwrap(),
        ))?;

        let rect = allocation.rectangle;
        let x = rect.min.x;
        let y = rect.min.y;
        let width = rect.width() - 1;
        let height = rect.height() - 1;
        let rect = Rect2D::new_from_size(IVec2::new(x, y), IVec2::new(width, height));
        self.glyphs.insert(key, (rect, offset));

        // put data into image array
        let mut image = self.image.0.write();
        for line_y in 0..height {
            let y = y + line_y;
            let image_offset = 4 * y as u32 * image.width() + 4 * x as u32;
            let glyph_offset = 4 * line_y * width;
            let len = 4 * width as usize;
            let dst = &mut image.bytes_mut()[image_offset as usize..][..len];
            let src = &data[glyph_offset as usize..][..len];
            dst.copy_from_slice(src);
        }

        Some(AtlasGlyphInfo {
            image: Arc::downgrade(&self.image),
            rect,
            offset,
        })
    }
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Failed to rasterize glyph {0:?}.")]
    FailedToRasterizeGlyph(CacheKey),
}

fn get_outlined_glyph_texture(
    font_system: &mut FontSystem,
    swash_cache: &mut SwashCache,
    glyph: PhysicalGlyph,
) -> Result<(Vec<u8>, Extent2D, IVec2), Error> {
    let image = swash_cache
        .get_image_uncached(font_system, glyph.cache_key)
        .ok_or(Error::FailedToRasterizeGlyph(glyph.cache_key))?;

    let cosmic_text::Placement {
        left,
        top,
        width,
        height,
    } = image.placement;

    let data = match image.content {
        cosmic_text::SwashContent::Mask => image
            .data
            .iter()
            .flat_map(|a| [255, 255, 255, *a])
            .collect(),
        cosmic_text::SwashContent::Color => image.data,
        cosmic_text::SwashContent::SubpixelMask => {
            // TODO: implement
            todo!()
        }
    };

    let extent = Extent2D { width, height };

    Ok((data, extent, IVec2::new(left, top)))
}

#[cfg(test)]
mod tests {
    use super::*;
    use cosmic_text::{Attrs, Buffer, Family, Metrics};
    use glam::Vec2;

    #[test]
    fn test() {
        let mut font_store = FontStore::new();
        let mut db = cosmic_text::fontdb::Database::new();
        let mut font_id_map = HashMap::<FontId, cosmic_text::fontdb::ID>::new();
        let mut reverse_font_id_map = HashMap::<cosmic_text::fontdb::ID, FontId>::new();

        let roboto = font_store.add_font_bytes(Arc::new(ROBOTO_BYTES));
        let noto_han = font_store.add_font_bytes(Arc::new(NOTO_SANS_HAN_BYTES));
        let id = *db
            .load_font_source(font_store.as_source(roboto).unwrap())
            .last()
            .unwrap();
        font_id_map.insert(roboto, id);
        reverse_font_id_map.insert(id, roboto);

        let id = *db
            .load_font_source(font_store.as_source(noto_han).unwrap())
            .last()
            .unwrap();
        font_id_map.insert(noto_han, id);
        reverse_font_id_map.insert(id, noto_han);

        db.set_sans_serif_family("Roboto");
        // db.load_system_fonts();

        let locale = sys_locale::get_locale().unwrap_or("en-DK".to_string());
        let mut font_system = FontSystem::new_with_locale_and_db(locale, db);
        let mut swash = SwashCache::new();

        let mut buffer = Buffer::new_empty(Metrics::new(14.0, 20.0));

        let mut buf = buffer.borrow_with(&mut font_system);

        let path = format!("../../assets/testing/hello.txt");
        let attrs = Attrs::new()
            .family(Family::SansSerif)
            .metrics(Metrics::new(48.0, 56.0));
        let _text = std::fs::read_to_string(path).expect("hello.txt");

        buf.set_text(
            "Hello, World! 你好! 안녕하세요",
            attrs,
            cosmic_text::Shaping::Advanced,
        );
        //buf.set_size(Some(400.0), Some(400.0));
        buf.set_wrap(cosmic_text::Wrap::Word);

        for line in buf.lines.iter_mut() {
            line.set_align(Some(cosmic_text::Align::Left));
        }

        buf.shape_until_scroll(false);

        let mut font_atlantes = FontAtlasSets::new();

        let mut glyphs = Vec::new();
        let mut size = Vec2::new(0.0, 0.0);

        _ = buffer
            .layout_runs()
            .flat_map(|run| {
                size.x = size.x.max(run.line_w);
                size.y = size.y + run.line_height;
                run.glyphs.iter().map(move |glyph| (glyph, run.line_y))
            })
            .try_for_each(|(glyph, line_y)| -> Result<(), Error> {
                let font_id = *reverse_font_id_map.get(&glyph.font_id).unwrap();
                let set = font_atlantes.sets.entry(font_id).or_default();

                let physical = glyph.physical((0.0, 0.0), 1.0);

                let glyph_info = set
                    .get_glyph_info(physical.cache_key)
                    .map(Ok)
                    .unwrap_or_else(|| {
                        set.add_glyph(&mut font_system, &mut swash, physical.clone())
                    })?;

                let pos = {
                    let x = glyph_info.offset.x as f32 + physical.x as f32;
                    let y = line_y.round() + physical.y as f32 - glyph_info.offset.y as f32;

                    Vec2::new(x, y)
                };

                let size = glyph_info.rect.size();
                glyphs.push((glyph_info, pos, size));

                Ok(())
            });

        let (width, height) = {
            let tmp = size.ceil().as_uvec2();
            (tmp.x, tmp.y)
        };

        let mut image = image::RgbaImage::from_pixel(width, height, image::Rgba([0, 0, 0, 255]));

        eprintln!("glyphs: {glyphs:#?}");
        eprintln!("image: {width}x{height}");

        for (info, pos, _size) in glyphs {
            let atlas_image = info.image.upgrade().unwrap();
            let glyph = atlas_image.0.read();

            let atlas_image = image::ImageBuffer::<image::Rgba<u8>, _>::from_raw(
                glyph.width(),
                glyph.height(),
                glyph.bytes(),
            )
            .unwrap();

            let glyph = atlas_image.view(
                info.rect.top_left().x as u32,
                info.rect.top_left().y as u32,
                info.rect.width() as u32,
                info.rect.height() as u32,
            );

            eprintln!("rect: {:?}", info.rect);
            eprintln!("image_size: {:?}", image.dimensions());

            image
                .copy_from(&*glyph, pos.x as u32, pos.y as u32)
                .unwrap();
        }

        image.save("rendered.png").unwrap();
    }
}