use core::mem::MaybeUninit;

use bytemuck::{Pod, Zeroable};
use device_api::dma::DmaAllocator;
use libk::{
    device::display::PixelFormat,
    dma::{BusAddress, DmaBuffer},
    error::Error,
};
use libk_util::sync::IrqSafeSpinlock;
use ygg_driver_virtio_core::{
    queue::{VirtQueue, VqManualNotification},
    transport::Transport,
};

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ControlHeader {
    pub ty: u32,
    pub flags: u32,
    pub fence_id: u64,
    pub ctx_id: u32,
    pub ring_idx: u8,
    _0: [u8; 3],
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct Rect {
    pub x: u32,
    pub y: u32,
    pub width: u32,
    pub height: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ScanoutInfo {
    pub r: Rect,
    pub enabled: u32,
    pub flags: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ResourceCreate2d {
    pub header: ControlHeader,
    pub resource_id: u32,
    pub format: u32,
    pub width: u32,
    pub height: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct SetScanout {
    pub header: ControlHeader,
    pub r: Rect,
    pub scanout_id: u32,
    pub resource_id: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct GpuMemEntry {
    pub addr: u64,
    pub length: u32,
    pub _0: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ResourceAttachBacking {
    pub header: ControlHeader,
    pub resource_id: u32,
    pub nr_entries: u32,
    pub entry: GpuMemEntry,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct ResourceFlush {
    pub header: ControlHeader,
    pub r: Rect,
    pub resource_id: u32,
    pub _0: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
pub struct TransferToHost2d {
    pub header: ControlHeader,
    pub r: Rect,
    pub offset: u64,
    pub resource_id: u32,
    pub _0: u32,
}

pub struct CommandExecution<'a, T: Transport> {
    pub(super) transport: &'a IrqSafeSpinlock<T>,
    pub(super) control: &'a VirtQueue<VqManualNotification>,
}

impl<'a, T: Transport> CommandExecution<'a, T> {
    fn send_recv<'r, Req: Pod>(
        &mut self,
        dma: &dyn DmaAllocator,
        req: &Req,
        buffer: &'r mut DmaBuffer<[MaybeUninit<u8>]>,
    ) -> Result<(&'r ControlHeader, &'r [u8]), Error> {
        let request = DmaBuffer::new_uninit_slice(dma, size_of::<Req>())?;
        let mut request = unsafe { DmaBuffer::assume_init_slice(request) };
        request.copy_from_slice(bytemuck::bytes_of(req));

        let len = self.control.enqueue_blocking(
            &[request.slice(0..size_of::<Req>())],
            &[buffer.slice_mut(0..buffer.len())],
            || self.transport.lock().notify(0),
        )? as usize;

        if len < size_of::<ControlHeader>() {
            log::warn!("virtio-gpu: invalid device response length: {len}");
            return Err(Error::InvalidArgument);
        }

        let payload = unsafe { buffer[..len].assume_init_ref() };
        let header = bytemuck::from_bytes(&payload[..size_of::<ControlHeader>()]);
        let data = &payload[size_of::<ControlHeader>()..len];

        Ok((header, data))
    }

    fn send_recv_no_data<Req: Pod>(
        &mut self,
        dma: &dyn DmaAllocator,
        req: &Req,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
    ) -> Result<(), Error> {
        let (response, _) = self.send_recv(dma, req, buffer)?;
        if response.ty == 0x1100 {
            Ok(())
        } else {
            Err(Error::InvalidArgument)
        }
    }

    pub fn query_scanouts<'r>(
        &mut self,
        dma: &dyn DmaAllocator,
        max_scanouts: usize,
        buffer: &'r mut DmaBuffer<[MaybeUninit<u8>]>,
    ) -> Result<&'r [ScanoutInfo], Error> {
        let request = ControlHeader {
            ty: 0x0100,
            flags: 0,
            fence_id: 0,
            ctx_id: 0,
            ring_idx: 0,
            _0: [0; 3],
        };

        let (response, data) = self.send_recv(dma, &request, buffer)?;
        if response.ty != 0x1101 {
            log::warn!("virtio-gpu: query_scanouts returned {:#x}", response.ty);
            return Err(Error::InvalidArgument);
        }
        let limit = size_of::<ScanoutInfo>() * max_scanouts;

        let scanouts =
            bytemuck::try_cast_slice(&data[..limit]).map_err(|_| Error::InvalidArgument)?;

        Ok(scanouts)
    }

    pub fn create_resource_2d(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        width: u32,
        height: u32,
        pixel_format: PixelFormat,
    ) -> Result<u32, Error> {
        let format = match pixel_format {
            PixelFormat::R8G8B8A8 => 67,
            _ => todo!(),
        };
        let request = ResourceCreate2d {
            header: ControlHeader {
                ty: 0x0101,
                ..ControlHeader::zeroed()
            },
            resource_id: 1,
            width,
            height,
            format,
        };

        self.send_recv_no_data(dma, &request, buffer)?;

        Ok(1)
    }

    pub fn attach_backing(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        resource_id: u32,
        base: BusAddress,
        length: u32,
    ) -> Result<(), Error> {
        let request = ResourceAttachBacking {
            header: ControlHeader {
                ty: 0x0106,
                ..ControlHeader::zeroed()
            },
            resource_id,
            nr_entries: 1,
            entry: GpuMemEntry {
                addr: base.into_u64(),
                length,
                _0: 0,
            },
        };

        self.send_recv_no_data(dma, &request, buffer)
    }

    pub fn set_scanout(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        scanout_id: u32,
        resource_id: u32,
        width: u32,
        height: u32,
    ) -> Result<(), Error> {
        let request = SetScanout {
            header: ControlHeader {
                ty: 0x0103,
                ..ControlHeader::zeroed()
            },
            r: Rect {
                x: 0,
                y: 0,
                width,
                height,
            },
            scanout_id,
            resource_id,
        };

        self.send_recv_no_data(dma, &request, buffer)
    }

    pub fn transfer_to_host_2d(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        resource_id: u32,
        r: Rect,
    ) -> Result<(), Error> {
        let request = TransferToHost2d {
            header: ControlHeader {
                ty: 0x0105,
                ..ControlHeader::zeroed()
            },
            r,
            offset: 0,
            resource_id,
            _0: 0,
        };

        self.send_recv_no_data(dma, &request, buffer)
    }

    pub fn resource_flush(
        &mut self,
        dma: &dyn DmaAllocator,
        buffer: &mut DmaBuffer<[MaybeUninit<u8>]>,
        resource_id: u32,
        r: Rect,
    ) -> Result<(), Error> {
        let request = ResourceFlush {
            header: ControlHeader {
                ty: 0x0104,
                ..ControlHeader::zeroed()
            },
            r,
            resource_id,
            _0: 0,
        };

        self.send_recv_no_data(dma, &request, buffer)
    }
}
