#![no_std]

use core::mem::MaybeUninit;

use alloc::{boxed::Box, format, sync::Arc, vec::Vec};
use async_trait::async_trait;
use bytemuck::{Pod, Zeroable};
use device_api::{
    device::{Device, DeviceInitContext},
    dma::DmaAllocator,
    interrupt::{InterruptAffinity, InterruptHandler, IrqVector},
};
use libk::{
    device::{block::BlockDevice, manager::probe_partitions},
    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
    error::Error,
    fs::devfs,
    task::runtime,
};
use libk_mm::{
    address::PhysicalAddress, table::MapAttributes, OnDemandPage, PageProvider, VirtualPage,
};
use libk_util::sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock};
use ygg_driver_pci::{
    device::{PciDeviceInfo, PreferredInterruptMode},
    macros::pci_driver,
};
use ygg_driver_virtio_core::{
    queue::VirtQueue,
    transport::{pci::PciTransport, Transport},
    DeviceStatus,
};
use yggdrasil_abi::{bitflags, io::FileMode};

extern crate alloc;

bitflags! {
    pub struct Features: u64 {
        const F_SIZE_MAX: bit 1;
        const F_SEG_MAX: bit 2;
        const F_RO: bit 5;
    }
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct DeviceConfig {
    capacity: u64,
    size_max: u32,
    seg_max: u32,
    // virtio_blk_geometry {
    cylinders: u16,
    heads: u8,
    sectors: u8,
    // }
    blk_size: u32,
    // virtio_blk_topology {
    physical_block_exp: u8,
    alignment_offset: u8,
    min_io_size: u16,
    opt_io_size: u32,
    // }
    writeback: u8,
    _0: u8,
    num_queues: u16,
    max_discard_sectors: u32,
    max_discard_seg: u32,
    discard_sector_alignment: u32,
    max_write_zeroes_sectors: u32,
    max_write_zeroes_seg: u32,
    write_zeroes_may_unmap: u8,
    _1: [u8; 3],
    max_secure_erase_sectors: u32,
    max_secure_erase_seg: u32,
    secure_erase_sector_alignment: u32,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct CommandHeader {
    ty: u32,
    _0: u32,
    sector: u64,
}

pub struct VirtioBlk<T: Transport + 'static> {
    transport: IrqSafeSpinlock<T>,
    pci_device_info: PciDeviceInfo,
    dma: Arc<dyn DmaAllocator>,

    segment_size: usize,
    read_only: bool,
    capacity: u64,
    request_queue: VirtQueue,
}

impl CommandHeader {
    const TYPE_READ: u32 = 0;
    const TYPE_WRITE: u32 = 1;

    pub fn for_read(lba: u64) -> Self {
        Self {
            ty: Self::TYPE_READ,
            _0: 0,
            sector: lba,
        }
    }

    pub fn for_write(lba: u64) -> Self {
        Self {
            ty: Self::TYPE_WRITE,
            _0: 0,
            sector: lba,
        }
    }
}

impl<T: Transport + 'static> VirtioBlk<T> {
    // Only one VQ
    const VQ_REQUEST_0: u16 = 0;

    fn new(
        dma: Arc<dyn DmaAllocator>,
        mut transport: T,
        pci_device_info: PciDeviceInfo,
    ) -> Result<Self, Error> {
        let features = Features::from(transport.read_device_features());
        let device_cfg = transport
            .device_cfg()
            .ok_or(Error::InvalidArgument)
            .inspect_err(|_| log::error!("virtio-blk does not expose device configuration"))?;
        let device_cfg: &DeviceConfig =
            bytemuck::from_bytes(&device_cfg[..size_of::<DeviceConfig>()]);

        let read_only = features.contains(Features::F_RO);
        let segment_size = if features.contains(Features::F_SIZE_MAX) {
            device_cfg.size_max as usize
        } else {
            // I guess no limit then?
            262144
        };
        let segment_limit = if features.contains(Features::F_SEG_MAX) {
            device_cfg.seg_max as usize
        } else {
            // I guess no limit then?
            8
        };
        let capacity = device_cfg.capacity;

        if segment_limit < 3 {
            // Won't be able to send header + data + status
            log::error!("virtio-blk: allowed segment count too small");
            return Err(Error::InvalidArgument);
        }

        let request_queue = VirtQueue::with_capacity(&*dma, Self::VQ_REQUEST_0, 256)?;

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

            segment_size,
            read_only,
            capacity,
            request_queue,
        })
    }

    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 blah blah blah

        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> {
        self.pci_device_info
            .init_interrupts(PreferredInterruptMode::Msi(true))?;
        let msi_info = self
            .pci_device_info
            .map_interrupt(InterruptAffinity::Any, self.clone())?;
        let vector = msi_info.map(|msi| msi.vector as u16);

        let mut transport = self.transport.lock();

        transport.set_queue(Self::VQ_REQUEST_0, &self.request_queue, vector);

        Ok(())
    }
}

impl<T: Transport + 'static> InterruptHandler for VirtioBlk<T> {
    fn handle_irq(self: Arc<Self>, _vector: IrqVector) -> bool {
        // Only one queue
        self.request_queue.handle_notify();
        true
    }
}

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

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

#[async_trait]
impl<T: Transport + 'static> BlockDevice for VirtioBlk<T> {
    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        DmaBuffer::new_uninit_slice(&*self.dma, size)
    }

    async fn read_aligned(
        &self,
        position: u64,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<(), Error> {
        if position % 512 != 0 || buffer.len() % 512 != 0 {
            return Err(Error::InvalidArgument);
        }
        let lba = position / 512;
        let lba_count = buffer.len() / 512;
        if lba + lba_count as u64 >= self.capacity {
            return Err(Error::InvalidArgument);
        }

        let mut header = DmaBuffer::new_slice(&*self.dma, 0, size_of::<CommandHeader>())?;
        *bytemuck::from_bytes_mut(&mut header[..]) = CommandHeader::for_read(lba);
        let mut status = DmaBuffer::new_uninit_slice(&*self.dma, 1)?;

        self.request_queue
            .enqueue_wait(
                &[header.slice(0..size_of::<CommandHeader>())],
                &[buffer, status.slice_mut(0..1)],
                || {
                    self.transport.lock().notify(Self::VQ_REQUEST_0);
                },
            )
            .await?;

        let status = unsafe { DmaBuffer::assume_init_slice(status) }[0];
        if status == 0 {
            Ok(())
        } else {
            Err(Error::InvalidOperation)
        }
    }
    async fn write_aligned(&self, position: u64, buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
        if self.read_only {
            return Err(Error::ReadOnly);
        }

        if position % 512 != 0 || buffer.len() % 512 != 0 {
            return Err(Error::InvalidArgument);
        }
        let lba = position / 512;
        let lba_count = buffer.len() / 512;
        if lba + lba_count as u64 >= self.capacity {
            return Err(Error::InvalidArgument);
        }

        let mut header = DmaBuffer::new_slice(&*self.dma, 0, size_of::<CommandHeader>())?;
        *bytemuck::from_bytes_mut(&mut header[..]) = CommandHeader::for_write(lba);
        let mut status = DmaBuffer::new_uninit_slice(&*self.dma, 1)?;

        self.request_queue
            .enqueue_wait(
                &[header.slice(0..size_of::<CommandHeader>()), buffer],
                &[status.slice_mut(0..1)],
                || {
                    self.transport.lock().notify(Self::VQ_REQUEST_0);
                },
            )
            .await?;

        let status = unsafe { DmaBuffer::assume_init_slice(status) }[0];
        if status == 0 {
            Ok(())
        } else {
            Err(Error::InvalidOperation)
        }
    }

    fn block_size(&self) -> usize {
        512
    }
    fn block_count(&self) -> u64 {
        self.capacity
    }
    fn max_blocks_per_request(&self) -> usize {
        // TODO this limit can be bumped or scatter-gather operations
        self.segment_size / 512
    }
}

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

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

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

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

static DEVICES: IrqSafeRwLock<Vec<Arc<dyn BlockDevice>>> = IrqSafeRwLock::new(Vec::new());

fn register_virtio_block_device(device: Arc<dyn BlockDevice>) {
    let index = {
        let mut devices = DEVICES.write();
        let index = devices.len();
        devices.push(device.clone());
        index
    };

    let name = format!("vb{index}");
    devfs::add_named_block_device(device.clone(), name.clone(), FileMode::new(0o600)).ok();

    runtime::spawn(async move {
        let name = name;
        log::info!("Probing partitions for {name}");
        probe_partitions(device, |index, partition| {
            let partition_name = format!("{name}p{}", index + 1);
            devfs::add_named_block_device(
                Arc::new(partition),
                partition_name,
                FileMode::new(0o600),
            )
            .ok();
        })
        .await
        .ok();
    })
    .ok();
}

pci_driver! {
    matches: [device (0x1AF4:0x1001)],
    driver: {
        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).unwrap();
            let device = VirtioBlk::new(dma.clone(), transport, info.clone())?;

            let device = Arc::new(device);

            Ok(device)
        }

        fn driver_name(&self) -> &str {
            "virtio-blk"
        }
    }
}
