use core::{cell::UnsafeCell, mem::MaybeUninit};

use alloc::sync::Arc;
use device_api::device::{Device, DeviceInitContext};
use device_tree::driver::{device_tree_driver, Node, ProbeContext};
use kernel_arch_aarch64::mem::table::L3;
use libk::device::{
    display::{
        DisplayDevice, DisplayMode, DisplayOwner, DriverFlags, FramebufferInfo, PixelFormat,
    },
    manager::DEVICE_REGISTRY,
};
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    device::DeviceMemoryIo,
    table::{EntryLevel, MapAttributes},
    OnDemandPage, PageBox, PageProvider, VirtualPage,
};
use libk_util::sync::IrqSafeSpinlock;
use tock_registers::{
    interfaces::{Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite, WriteOnly},
};
use yggdrasil_abi::error::Error;

register_bitfields! {
    u32,
    MboxValue [
        ADDRESS OFFSET(4) NUMBITS(28) [],
        CHANNEL OFFSET(0) NUMBITS(4) [
            PowerManagement = 0,
            Framebuffer = 1,
            VirtualUart = 2,
            Vchiq = 3,
            Led = 4,
            Button = 5,
            Touch = 6,
            PropertyArmToVc = 8,
            PropertyVcToArm = 9
        ]
    ],
    MboxStatus [
        FULL    OFFSET(31) NUMBITS(1)   [],
        EMPTY   OFFSET(30) NUMBITS(1)   [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    Regs {
        (0x00 => READ: ReadOnly<u32, MboxValue::Register>),
        (0x04 => _0),
        (0x10 => POLL: ReadOnly<u32>),
        (0x14 => SENDER: ReadOnly<u32>),
        (0x18 => STATUS: ReadOnly<u32, MboxStatus::Register>),
        (0x1C => CONFIG: ReadWrite<u32>),
        (0x20 => WRITE: WriteOnly<u32, MboxValue::Register>),
        (0x24 => _1),
        (0x40 => @END),
    }
}

struct Bcm2835Mbox {
    regs: IrqSafeSpinlock<DeviceMemoryIo<'static, Regs>>,
}

struct Bcm2835Framebuffer {
    mode: DisplayMode,
    base: PhysicalAddress,
    size: usize,
    pitch: usize,
}

struct FramebufferResponse {
    physical_address: u32,
    size: u32,
    pitch: u32,
}

unsafe impl Send for Bcm2835Mbox {}
unsafe impl Sync for Bcm2835Mbox {}

impl Device for Bcm2835Mbox {
    fn display_name(&self) -> &str {
        "bcm2835-mbox"
    }
}

impl Bcm2835Mbox {
    const RAM_TO_GPU_OFFSET: u32 = 0x40000000;

    pub unsafe fn call_raw(&self, buffer: PhysicalAddress, channel: u32) -> Result<(), Error> {
        log::info!("bcm2835-mbox: mbox call buffer={buffer:#x}, channel={channel}");

        let address = buffer
            .try_into_u32()
            .inspect_err(|_| log::error!("bcm2835-mbox: invalid physical address {buffer:#x}"))
            .map_err(|_| Error::InvalidArgument)?;
        if address >= Self::RAM_TO_GPU_OFFSET {
            log::error!("bcm2835-mbox: RAM address {buffer:#x} is too high");
            return Err(Error::InvalidArgument);
        }
        let address = address + Self::RAM_TO_GPU_OFFSET;

        let regs = self.regs.lock();
        while regs.STATUS.matches_all(MboxStatus::FULL::SET) {
            core::hint::spin_loop();
        }

        let val = MboxValue::ADDRESS.val(address >> 4) + MboxValue::CHANNEL.val(channel);

        regs.WRITE.write(val);

        loop {
            while regs.STATUS.matches_all(MboxStatus::EMPTY::SET) {
                core::hint::spin_loop();
            }

            let val = regs.READ.extract();

            if val.matches_all(MboxValue::CHANNEL.val(channel)) {
                return Ok(());
            }
        }
    }

    pub fn configure_framebuffer(
        &self,
        mode_info: &DisplayMode,
    ) -> Result<FramebufferResponse, Error> {
        #[derive(Debug, Clone, Copy)]
        #[repr(C)]
        struct FramebufferRequest {
            physical_width: u32,
            physical_height: u32,
            virtual_width: u32,
            virtual_height: u32,
            pitch: u32,
            bpp: u32,
            x_offset: u32,
            y_offset: u32,
            address: u32,
            size: u32,
        }

        let mut buffer = PageBox::new(UnsafeCell::new(FramebufferRequest {
            physical_width: mode_info.width,
            physical_height: mode_info.height,
            virtual_width: mode_info.width,
            virtual_height: mode_info.height,
            pitch: 0,
            bpp: 32,
            x_offset: 0,
            y_offset: 0,
            address: 0,
            size: 0,
        }))?;

        // TODO flush cache for buffer
        unsafe { self.call_raw(buffer.as_physical_address(), 0x01) }?;

        let buffer = buffer.get_mut();

        log::info!("bcm2835-fb: mode set response: {buffer:#?}");

        if buffer.address == 0 || buffer.size == 0 || buffer.pitch == 0 {
            // Request failed
            log::error!("bcm2835-fb: mode set failed");
            return Err(Error::InvalidArgument);
        }

        let address = buffer.address & (Bcm2835Mbox::RAM_TO_GPU_OFFSET - 1);

        Ok(FramebufferResponse {
            physical_address: address,
            size: buffer.size,
            pitch: buffer.pitch,
        })
    }
}

impl Bcm2835Framebuffer {
    pub fn setup(mbox: &Arc<Bcm2835Mbox>) -> Result<(), Error> {
        const WIDTH: u32 = 1024;
        const HEIGHT: u32 = 768;

        let mode = DisplayMode {
            index: 0,
            width: WIDTH,
            height: HEIGHT,
            frames_per_second: 60,
            pixel_format: PixelFormat::R8G8B8A8,
        };

        let framebuffer = mbox.configure_framebuffer(&mode)?;

        let fb = Arc::new(Self {
            mode,
            base: PhysicalAddress::from_u32(framebuffer.physical_address),
            size: framebuffer.size as usize,
            pitch: framebuffer.pitch as usize,
        });

        DEVICE_REGISTRY.display.register(fb, false)?;

        Ok(())
    }
}

impl Device for Bcm2835Framebuffer {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        unreachable!()
    }

    fn display_name(&self) -> &str {
        "bcm2835-fb"
    }
}

impl PageProvider for Bcm2835Framebuffer {
    fn get_page(&self, offset: u64) -> Result<VirtualPage, Error> {
        let offset = offset as usize;
        if offset + L3::SIZE > self.size {
            log::warn!(
                "Tried to map offset {:#x}, but size is {:#x}",
                offset,
                self.size
            );
            Err(Error::InvalidMemoryOperation)
        } else {
            let page = self.base.add(offset);
            Ok(VirtualPage::Immediate(page))
        }
    }

    fn clone_page(
        &self,
        _offset: u64,
        _src_phys: PhysicalAddress,
        _src_attrs: MapAttributes,
    ) -> Result<PhysicalAddress, Error> {
        unimplemented!()
    }

    fn release_page(
        &self,
        _offset: u64,
        _phys: PhysicalAddress,
        _dirty: bool,
    ) -> Result<(), Error> {
        todo!()
    }

    fn ondemand_fetch(&self, _offset: u64) -> Result<OnDemandPage, Error> {
        unimplemented!()
    }
}

impl DisplayDevice for Bcm2835Framebuffer {
    fn lock(&self, _owner: DisplayOwner) -> Result<(), Error> {
        Ok(())
    }

    fn unlock(&self, _owner: DisplayOwner) -> Result<(), Error> {
        Ok(())
    }

    fn synchronize(&self) -> Result<(), Error> {
        Ok(())
    }

    fn driver_flags(&self) -> DriverFlags {
        DriverFlags::empty()
    }

    fn set_display_mode(&self, _index: u32, _double_buffer: bool) -> Result<(), Error> {
        Err(Error::NotImplemented)
    }

    fn active_framebuffer(&self) -> Result<FramebufferInfo, usize> {
        Ok(FramebufferInfo {
            base: self.base,
            kernel_base: None,
            stride: self.pitch,
            size: self.size,
        })
    }

    fn active_display_mode(&self) -> Option<DisplayMode> {
        Some(self.mode)
    }

    fn query_display_modes(&self, modes: &mut [MaybeUninit<DisplayMode>]) -> Result<usize, Error> {
        if modes.is_empty() {
            return Err(Error::BufferTooSmall);
        }
        modes[0].write(self.mode);
        Ok(1)
    }

    fn active_framebuffers(
        &self,
        output: &mut [MaybeUninit<FramebufferInfo>],
    ) -> Result<usize, usize> {
        if output.is_empty() {
            return Err(1);
        }
        output[0].write(FramebufferInfo {
            base: self.base,
            kernel_base: None,
            stride: self.pitch,
            size: self.size,
        });
        Ok(1)
    }
}

device_tree_driver! {
    compatible: ["brcm,bcm2835-mbox"],
    driver: {
        fn probe(&self, node: &Arc<Node>, context: &mut ProbeContext) -> Option<Arc<dyn Device>> {
            let base = node.map_base(context, 0)?;
            // TODO interrupts
            let regs = unsafe { DeviceMemoryIo::map(base, Default::default()) }
                .inspect_err(|e| log::error!("bcm2835-mbox: mapping failed: {e:?}"))
                .ok()?;
            let mbox = Arc::new(Bcm2835Mbox {
                regs: IrqSafeSpinlock::new(regs)
            });
            // TODO there's no fdt node for a framebuffer, so it's hardcoded as a child of the mbox
            if let Err(error) = Bcm2835Framebuffer::setup(&mbox) {
                log::error!("bcm2835-mbox: framebuffer setup error: {error:?}");
            }

            Some(mbox)
        }
    }
}
