use egui::Vec2;

use crate::{MemoryBank, Page, backend::memory::PhysAddr};

use super::Component;

pub struct FrameBuffer {
    width: usize,
    height: usize,
    addr: PhysAddr,
    mem: MemoryBank,
    pub buffer: Vec<u32>, // RGBA pixels
    visible: bool,

    texture: Option<egui::TextureHandle>,
}

impl FrameBuffer {
    pub fn new(width: u32, height: u32, addr: PhysAddr, mem: MemoryBank) -> Self {
        Self {
            width: width as usize,
            height: height as usize,
            addr,
            mem,
            buffer: vec![0u32; width as usize * height as usize],
            visible: true,

            texture: None,
        }
    }

    pub fn width(&self) -> usize {
        self.width
    }
    pub fn height(&self) -> usize {
        self.height
    }
    pub fn visible(&mut self) -> &mut bool {
        &mut self.visible
    }

    fn internal_read(&self) -> Vec<u32> {
        let byte_size = self.width * self.height * 4;
        let page_count = byte_size.div_ceil(Page::SIZE);

        (0..page_count)
            .map(|i| {
                let page = self.mem.read_page(self.addr + i as u32 * 4096);
                page.as_slice().to_owned()
            })
            .flatten()
            .take(byte_size)
            .collect::<Vec<u8>>()
            .chunks_exact(4)
            .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap()))
            .collect()
    }

    pub fn changed(&mut self) -> bool {
        let temp = self.internal_read();
        if temp != self.buffer {
            self.buffer = temp;
            return true;
        }
        false
    }

    pub fn read(&self) -> &[u32] {
        &self.buffer
    }

    /// Convert buffer to RGBA bytes for use with egui/wgpu textures
    pub fn as_rgba_bytes(&self) -> Vec<u8> {
        self.buffer
            .iter()
            .flat_map(|&pixel| pixel.to_le_bytes())
            .collect()
    }
}

impl Component for FrameBuffer {
    fn title(&self) -> &str {
        "Framebuffer"
    }

    fn visible(&mut self) -> &mut bool {
        &mut self.visible
    }

    fn ui(&mut self, ui: &mut egui::Ui, ctx: &egui::Context) {
        let changed = self.changed();

        // get or create the texture
        let texture = self.texture.get_or_insert_with(|| {
            ctx.load_texture(
                "framebuffer",
                egui::ColorImage {
                    size: [self.width, self.height],
                    source_size: Vec2::from([self.width as f32, self.height as f32]),
                    pixels: vec![egui::Color32::BLACK; self.width * self.height],
                },
                egui::TextureOptions::NEAREST,
            )
        });

        if changed {
            let pixels: Vec<egui::Color32> = self
                .buffer
                .iter()
                .map(|&p| {
                    let bytes = p.to_le_bytes();
                    egui::Color32::from_rgba_premultiplied(bytes[0], bytes[1], bytes[2], bytes[3])
                })
                .collect();

            texture.set(
                egui::ColorImage {
                    size: [self.width, self.height],
                    source_size: Vec2::from([self.width as f32, self.height as f32]),
                    pixels,
                },
                egui::TextureOptions::NEAREST,
            );
        }

        // scale to fit available space while preserving aspect ratio
        let available = ui.available_size();
        let aspect = self.width as f32 / self.height as f32;
        let size = if available.x / aspect <= available.y {
            egui::vec2(available.x, available.x / aspect)
        } else {
            egui::vec2(available.y * aspect, available.y)
        };

        ui.image(egui::load::SizedTexture::new(texture.id(), size));
    }
}