use core::{mem::MaybeUninit, ops::Deref};

use alloc::boxed::Box;
use async_trait::async_trait;
use libk::{
    dma::{DmaBuffer, DmaSliceMut},
    error::Error,
};

use crate::command::ScsiCommand;

#[async_trait]
pub trait ScsiTransport: Send + Sync {
    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error>;

    /// Perform a no-data request
    async fn perform_request_raw(
        &mut self,
        lun: u8,
        request_data: &[u8],
        response_buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<usize, Error>;

    fn max_bytes_per_request(&self) -> usize;
}

pub struct ScsiTransportWrapper {
    inner: Box<dyn ScsiTransport>,
}

impl ScsiTransportWrapper {
    pub fn new(inner: Box<dyn ScsiTransport>) -> Self {
        Self { inner }
    }

    pub async fn read(
        &mut self,
        lun: u8,
        lba: u64,
        lba_count: u16,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<usize, Error> {
        if lba >= u32::MAX as u64 {
            return Err(Error::InvalidArgument);
        }
        let lba_bytes = (lba as u32).to_be_bytes();
        let lba_count = lba_count.to_be_bytes();
        // Issue a READ (10) command
        let request_buffer = [
            0x28,
            0x00,
            lba_bytes[0],
            lba_bytes[1],
            lba_bytes[2],
            lba_bytes[3],
            0x00,
            lba_count[0],
            lba_count[1],
            0x00,
        ];

        self.inner
            .perform_request_raw(lun, &request_buffer, buffer)
            .await
    }

    pub async fn perform_command<R: ScsiCommand>(
        &mut self,
        lun: u8,
        request: R,
    ) -> Result<R::Response, Error>
    where
        [u8; R::RESPONSE_LEN]: Sized,
        [u8; R::REQUEST_LEN]: Sized,
    {
        let mut response_buffer = self.allocate_buffer(R::RESPONSE_LEN)?;
        let request_buffer = request.into_bytes();

        let response_len = self
            .inner
            .perform_request_raw(
                lun,
                &request_buffer,
                response_buffer.slice_mut(0..R::RESPONSE_LEN),
            )
            .await?;
        let response_bytes = unsafe { response_buffer[..response_len].assume_init_ref() };

        R::parse_response(response_bytes)
    }
}

impl Deref for ScsiTransportWrapper {
    type Target = dyn ScsiTransport;

    fn deref(&self) -> &Self::Target {
        self.inner.as_ref()
    }
}
