#![feature(maybe_uninit_slice)]
#![no_std]

extern crate alloc;

use core::mem::MaybeUninit;

use alloc::{sync::Arc, vec::Vec};
use command::{CommandExecution, ScanoutInfo};
use device_api::{
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
};
use libk::{
    device::{
        display::{
            DisplayDevice, DisplayMode, DisplayOwner, DriverFlags, FramebufferInfo, PixelFormat,
        },
        manager::DEVICE_REGISTRY,
    },
    dma::DmaBuffer,
};
use libk_mm::{
    address::{AsPhysicalAddress, PhysicalAddress},
    table::MapAttributes,
    OnDemandPage, PageProvider, VirtualPage, L3_PAGE_SIZE,
};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use ygg_driver_pci::{device::PciDeviceInfo, macros::pci_driver};
use ygg_driver_virtio_core::{
    queue::{VirtQueue, VqManualNotification},
    transport::{pci::PciTransport, Transport},
    DeviceStatus,
};
use yggdrasil_abi::error::Error;

mod command;

struct Queues {
    control: VirtQueue<VqManualNotification>,
}

struct Framebuffer {
    scanout_index: usize,
    resource_id: u32,

    double: bool,
    dma_buffer: DmaBuffer<[MaybeUninit<u8>]>,
    stride: usize,
    size: usize,
}

struct Config {
    scanouts: Vec<ScanoutInfo>,
    framebuffer: Option<Framebuffer>,

    owner: DisplayOwner,
    response: DmaBuffer<[MaybeUninit<u8>]>,
}

pub struct VirtioGpu<T: Transport> {
    transport: IrqSafeSpinlock<T>,
    #[allow(unused)]
    pci_device_info: PciDeviceInfo,

    queues: OneTimeInit<Queues>,
    config: IrqSafeRwLock<Config>,

    dma: Arc<dyn DmaAllocator>,

    num_scanouts: usize,
}

impl<T: Transport + 'static> VirtioGpu<T> {
    pub fn new(
        dma: Arc<dyn DmaAllocator>,
        transport: T,
        info: PciDeviceInfo,
    ) -> Result<Self, Error> {
        // Read num-scanouts from device config
        let Some(device_cfg) = transport.device_cfg() else {
            log::error!("virtio-gpu must have device-specific configuration section");
            return Err(Error::InvalidArgument);
        };

        let mut num_scanouts = [0; size_of::<u32>()];
        num_scanouts.copy_from_slice(&device_cfg[8..12]);
        let num_scanouts = u32::from_le_bytes(num_scanouts);

        if !(1..=16).contains(&num_scanouts) {
            log::error!("virtio-gpu has invalid num_scanouts: {num_scanouts}");
            return Err(Error::InvalidArgument);
        }

        Ok(Self {
            transport: IrqSafeSpinlock::new(transport),
            pci_device_info: info,

            queues: OneTimeInit::new(),
            config: IrqSafeRwLock::new(Config {
                scanouts: Vec::new(),
                framebuffer: None,
                response: DmaBuffer::new_uninit_slice(&*dma, L3_PAGE_SIZE)?,
                owner: DisplayOwner::None,
            }),

            dma,

            num_scanouts: num_scanouts as usize,
        })
    }

    fn begin_init(&self) -> Result<DeviceStatus, Error> {
        let mut transport = self.transport.lock();
        let mut status = DeviceStatus::RESET_VALUE;

        log::debug!("Reset device");
        transport.write_device_status(status);
        status |= DeviceStatus::ACKNOWLEDGE;
        transport.write_device_status(status);
        status |= DeviceStatus::DRIVER;
        transport.write_device_status(status);

        let _device_features = transport.read_device_features();

        // TODO select features

        transport.write_driver_features(0);

        status |= DeviceStatus::FEATURES_OK;
        transport.write_device_status(status);

        if !transport
            .read_device_status()
            .contains(DeviceStatus::FEATURES_OK)
        {
            return Err(Error::InvalidOperation);
        }

        Ok(status)
    }

    fn finish_init(&self, status: DeviceStatus) {
        let mut transport = self.transport.lock();

        transport.write_device_status(status | DeviceStatus::DRIVER_OK);
    }

    fn setup_queues(self: &Arc<Self>) -> Result<(), Error> {
        // TODO cursorq
        let mut transport = self.transport.lock();

        let control = VirtQueue::with_capacity_manual(&*self.dma, 0, 128)?;
        transport.set_queue(0, &control, None);

        self.queues.init(Queues { control });

        Ok(())
    }

    fn begin_command(&self) -> CommandExecution<'_, T> {
        CommandExecution {
            transport: &self.transport,
            control: &self.queues.get().control,
        }
    }

    fn setup_display(&self) -> Result<(), Error> {
        let mut config = self.config.write();
        let mut command = self.begin_command();

        let scanouts =
            command.query_scanouts(&*self.dma, self.num_scanouts, &mut config.response)?;
        for (i, scanout) in scanouts.iter().enumerate() {
            log::info!(
                "virtio-gpu: [{i}] {}x{} + {},{}",
                scanout.r.width,
                scanout.r.height,
                scanout.r.x,
                scanout.r.y
            );
        }
        config.scanouts = scanouts.into();

        Ok(())
    }

    fn setup_mode(&self, index: usize) -> Result<(), Error> {
        let mut config = self.config.write();

        if config.framebuffer.is_some() {
            // TODO destroy old resource 2d and free the old framebuffer
            todo!()
        }

        /*
           5.7.6.3 Device Operation: Using pageflip
           It is possible to create multiple framebuffers, flip between them using
           VIRTIO_GPU_CMD_SET_SCANOUT and VIRTIO_GPU_CMD_RESOURCE_FLUSH, and update
           the invisible framebuffer using VIRTIO_GPU_CMD_TRANSFER_TO_HOST_2D.
        */

        let scanout = config.scanouts.get(index).ok_or(Error::DoesNotExist)?;
        let w = scanout.r.width;
        let h = scanout.r.height;
        let stride = w as usize * size_of::<u32>();
        let size = stride * h as usize;

        let dma_buffer = DmaBuffer::new_uninit_slice(&*self.dma, size)?;

        let mut command = self.begin_command();

        let resource_id = command.create_resource_2d(
            &*self.dma,
            &mut config.response,
            w,
            h,
            PixelFormat::R8G8B8A8,
        )?;
        command.attach_backing(
            &*self.dma,
            &mut config.response,
            resource_id,
            dma_buffer.bus_address(),
            size.try_into().unwrap(),
        )?;
        command.set_scanout(
            &*self.dma,
            &mut config.response,
            index as u32,
            resource_id,
            w,
            h,
        )?;

        config.framebuffer = Some(Framebuffer {
            scanout_index: index,
            double: false,

            resource_id,
            size,
            stride,

            dma_buffer,
        });

        Ok(())
    }

    fn flush(&self) -> Result<(), Error> {
        let mut config = self.config.write();

        let framebuffer = config.framebuffer.as_ref().ok_or(Error::DoesNotExist)?;
        let r = config.scanouts[framebuffer.scanout_index].r;

        let mut command = self.begin_command();

        if framebuffer.double {
            // Flip the buffer
            todo!()
        } else {
            let resource_id = framebuffer.resource_id;

            command.transfer_to_host_2d(&*self.dma, &mut config.response, resource_id, r)?;
            command.resource_flush(&*self.dma, &mut config.response, resource_id, r)?;

            Ok(())
        }
    }
}

impl<T: Transport + 'static> Device for VirtioGpu<T> {
    unsafe fn init(self: Arc<Self>, _cx: DeviceInitContext) -> Result<(), Error> {
        let status = self.begin_init()?;
        self.setup_queues()?;
        self.finish_init(status);

        // Set up some initial mode
        self.setup_display()?;
        self.setup_mode(0)?;

        DEVICE_REGISTRY.display.register(self.clone(), false)?;

        Ok(())
    }

    fn display_name(&self) -> &str {
        "VirtIO GPU Device"
    }
}

impl<T: Transport + 'static> PageProvider for VirtioGpu<T> {
    fn ondemand_fetch(&self, _opaque: u64) -> Result<OnDemandPage, Error> {
        unreachable!()
    }

    fn get_page(&self, offset: u64) -> Result<VirtualPage, Error> {
        // TODO check that the page is mapped by framebuffer owner
        let config = self.config.read();
        let framebuffer = config.framebuffer.as_ref().ok_or(Error::DoesNotExist)?;
        if offset as usize + L3_PAGE_SIZE > framebuffer.dma_buffer.page_count() * L3_PAGE_SIZE {
            log::warn!(
                "virtio-gpu: offset {:#x} outside of framebuffer bounds {:#x}",
                offset,
                framebuffer.size
            );
            return Err(Error::InvalidMemoryOperation);
        }
        let phys = unsafe { framebuffer.dma_buffer.as_physical_address() }.add(offset as usize);

        Ok(VirtualPage::Immediate(phys))
    }

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

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

impl<T: Transport + 'static> DisplayDevice for VirtioGpu<T> {
    fn driver_flags(&self) -> DriverFlags {
        DriverFlags::REPLACES_BOOT_FRAMEBUFFER
    }

    fn lock(&self, owner: DisplayOwner) -> Result<(), Error> {
        let mut config = self.config.write();
        if config.owner == owner {
            return Ok(());
        }
        match (config.owner, owner) {
            (DisplayOwner::None | DisplayOwner::Kernel, DisplayOwner::Process(_)) => {
                config.owner = owner;
                Ok(())
            }
            (DisplayOwner::None, DisplayOwner::Kernel) => {
                config.owner = owner;
                Ok(())
            }
            (_, _) => Err(Error::WouldBlock),
        }
    }

    fn unlock(&self, owner: DisplayOwner) -> Result<(), Error> {
        let mut config = self.config.write();
        if config.owner == owner {
            config.owner = DisplayOwner::None;
            Ok(())
        } else {
            Err(Error::InvalidArgument)
        }
    }

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

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

    fn active_framebuffers(
        &self,
        output: &mut [MaybeUninit<FramebufferInfo>],
    ) -> Result<usize, usize> {
        let config = self.config.read();

        if let Some(framebuffer) = config.framebuffer.as_ref() {
            if output.is_empty() {
                return Err(1);
            }

            output[0].write(FramebufferInfo {
                base: unsafe { framebuffer.dma_buffer.as_physical_address() },
                kernel_base: Some(framebuffer.dma_buffer.as_ptr().addr()),
                stride: framebuffer.stride,
                size: framebuffer.size,
            });

            Ok(1)
        } else {
            Ok(0)
        }
    }

    fn active_display_mode(&self) -> Option<DisplayMode> {
        let config = self.config.read();
        let framebuffer = config.framebuffer.as_ref()?;
        let scanout = &config.scanouts[framebuffer.scanout_index];

        Some(DisplayMode {
            index: framebuffer.scanout_index as _,
            width: scanout.r.width,
            height: scanout.r.height,
            pixel_format: PixelFormat::R8G8B8A8,
            frames_per_second: 0,
        })
    }

    fn query_display_modes(&self, _modes: &mut [MaybeUninit<DisplayMode>]) -> Result<usize, Error> {
        todo!()
    }
}

pci_driver! {
    matches: [device (0x1AF4:0x1050)],
    driver: {
        fn driver_name(&self) -> &str {
            "virtio-gpu"
        }

        fn probe(&self, info: &PciDeviceInfo, dma: &Arc<dyn DmaAllocator>) -> Result<Arc<dyn Device>, Error> {
            let space = &info.config_space;

            let transport = PciTransport::from_config_space(space)
                .inspect_err(|error| {
                    log::error!("Couldn't set up PCI virtio transport: {error:?}");
                })
                .map_err(|_| Error::InvalidArgument)?;
            let device = VirtioGpu::new(dma.clone(), transport, info.clone())?;
            let device = Arc::new(device);

            Ok(device)
        }
    }
}
