use core::mem::MaybeUninit;

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use bytemuck::{Pod, Zeroable};
use libk::{
    dma::{DmaBuffer, DmaSliceMut},
    error::Error,
};
use ygg_driver_scsi::{transport::ScsiTransport, ScsiEnclosure};

use crate::{
    communication::UsbDirection,
    device::{UsbDeviceAccess, UsbDeviceDetachHandler},
    error::UsbError,
    info::{UsbDeviceClass, UsbEndpointType},
    pipe::{
        control::{ControlTransferSetup, UsbClassSpecificRequest},
        normal::{UsbBulkInPipeAccess, UsbBulkOutPipeAccess},
    },
};

use super::{UsbClassInfo, UsbDriver};

pub struct UsbMassStorageDriverBulkOnly;

#[derive(Debug, Clone, Copy, Zeroable, Pod)]
#[repr(C)]
struct Cbw {
    signature: u32,       // 0x00
    tag: u32,             // 0x04
    transfer_length: u32, // 0x08
    flags: u8,            // 0x0C
    lun: u8,              // 0x0D
    cb_length: u8,        // 0x0E
    cb_data: [u8; 16],    // 0x0F
    // Not sent
    _0: u8,
}

#[derive(Debug, Clone, Copy, Zeroable, Pod)]
#[repr(C)]
struct Csw {
    signature: u32,
    tag: u32,
    data_residue: u32,
    status: u8,
    _0: [u8; 3],
}

struct Bbb {
    #[allow(unused)]
    device: Arc<UsbDeviceAccess>,
    in_pipe: UsbBulkInPipeAccess,
    out_pipe: UsbBulkOutPipeAccess,
    last_tag: u32,
}

struct DetachHandler(Arc<ScsiEnclosure>);

impl Bbb {
    pub fn new(
        device: Arc<UsbDeviceAccess>,
        in_pipe: UsbBulkInPipeAccess,
        out_pipe: UsbBulkOutPipeAccess,
    ) -> Result<Self, UsbError> {
        Ok(Self {
            device,
            in_pipe,
            out_pipe,
            last_tag: 0,
        })
    }
}

impl Bbb {
    async fn send_cbw(
        &mut self,
        lun: u8,
        host_to_dev: bool,
        command: &[u8],
        response_len: usize,
    ) -> Result<u32, Error> {
        self.last_tag = self.last_tag.wrapping_add(1);

        let flags = if !host_to_dev { 1 << 7 } else { 0 };
        let tag = self.last_tag;
        let mut cbw_bytes = [0; 32];
        let cbw = bytemuck::from_bytes_mut::<Cbw>(&mut cbw_bytes);

        cbw.signature = 0x43425355;
        cbw.transfer_length = response_len as u32;
        cbw.flags = flags;
        cbw.tag = tag;
        cbw.lun = lun;
        cbw.cb_length = command.len() as u8;
        cbw.cb_data[..command.len()].copy_from_slice(command);

        self.out_pipe
            .write(&cbw_bytes[..31])
            .await
            .inspect_err(|error| log::error!("msc: CBW send error: {error:?}"))?;

        Ok(tag)
    }

    async fn read_csw(&mut self, tag: u32) -> Result<(), Error> {
        let mut csw_bytes = [0; 16];
        self.in_pipe
            .read_exact(&mut csw_bytes[..13])
            .await
            .inspect_err(|error| log::error!("msc: CSW receive error: {error:?}"))?;
        let csw = bytemuck::from_bytes::<Csw>(&csw_bytes);

        if csw.signature != 0x53425355 {
            log::warn!("msc: invalid csw signature");
            return Err(Error::InvalidArgument);
        }
        if csw.tag != tag {
            let csw_tag = csw.tag;
            log::warn!("msc: invalid csw tag (got {}, expected {tag})", csw_tag);
            return Err(Error::InvalidArgument);
        }
        if csw.status != 0x00 {
            return Err(Error::InvalidArgument);
        }
        Ok(())
    }

    async fn read_response_data(
        &mut self,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<usize, Error> {
        if buffer.is_empty() {
            return Ok(0);
        }
        let len = self
            .in_pipe
            .read_dma(buffer)
            .await
            .inspect_err(|error| log::error!("msc: DMA read error: {error:?}"))?;
        Ok(len)
    }
}

#[async_trait]
impl ScsiTransport for Bbb {
    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        Ok(self.in_pipe.allocate_dma_buffer(size)?)
    }

    async fn perform_request_raw(
        &mut self,
        lun: u8,
        request_data: &[u8],
        response_buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<usize, Error> {
        if request_data.len() > 16 || response_buffer.len() > self.max_bytes_per_request() {
            return Err(Error::InvalidArgument);
        }

        let tag = self
            .send_cbw(lun, false, request_data, response_buffer.len())
            .await?;
        let response_len = self.read_response_data(response_buffer).await?;
        self.read_csw(tag).await?;
        Ok(response_len)
    }

    fn max_bytes_per_request(&self) -> usize {
        32768
    }
}

impl UsbDeviceDetachHandler for DetachHandler {
    fn handle_device_detach(&self) {
        log::info!("Mass storage detached");
        self.0.detach();
    }
}

#[derive(Debug, Pod, Zeroable, Clone, Copy)]
#[repr(C)]
pub struct BulkOnlyMassStorageReset;

#[derive(Debug, Pod, Zeroable, Clone, Copy)]
#[repr(C)]
pub struct GetMaxLun;

impl UsbClassSpecificRequest for BulkOnlyMassStorageReset {
    const BM_REQUEST_TYPE: u8 = 0b00100001;
    const B_REQUEST: u8 = 0b11111111;
}

impl UsbClassSpecificRequest for GetMaxLun {
    const BM_REQUEST_TYPE: u8 = 0b10100001;
    const B_REQUEST: u8 = 0b11111110;
}

#[async_trait]
impl UsbDriver for UsbMassStorageDriverBulkOnly {
    async fn run(self: Arc<Self>, device: Arc<UsbDeviceAccess>) -> Result<(), UsbError> {
        // TODO filter to only accept BBB config
        let config = device.select_configuration(|_| true).await?.unwrap();
        // Bulk-in, bulk-out
        assert_eq!(config.endpoints.len(), 2);
        let control_pipe = device.control_pipe();
        let (in_index, in_info) = config
            .find_endpoint(|ep| ep.is(UsbEndpointType::Bulk, UsbDirection::In))
            .ok_or(UsbError::InvalidConfiguration)?;
        let (out_index, out_info) = config
            .find_endpoint(|ep| ep.is(UsbEndpointType::Bulk, UsbDirection::Out))
            .ok_or(UsbError::InvalidConfiguration)?;
        let in_pipe = device
            .open_bulk_in_pipe(in_index, in_info.max_packet_size as u16)
            .await?;
        let out_pipe = device
            .open_bulk_out_pipe(out_index, out_info.max_packet_size as u16)
            .await?;

        // Perform a Bulk-Only Mass Storage Reset
        // TODO interface id?
        control_pipe
            .control_transfer(ControlTransferSetup {
                bm_request_type: BulkOnlyMassStorageReset::BM_REQUEST_TYPE,
                b_request: BulkOnlyMassStorageReset::B_REQUEST,
                w_value: 0,
                w_index: 0,
                w_length: 0,
            })
            .await?;

        // Get max LUN
        // TODO on devices which do not support multiple LUNs, this command may STALL
        let mut buffer = [MaybeUninit::uninit()];
        let len = control_pipe
            .control_transfer_in(
                ControlTransferSetup {
                    bm_request_type: GetMaxLun::BM_REQUEST_TYPE,
                    b_request: GetMaxLun::B_REQUEST,
                    w_value: 0,
                    w_index: 0,
                    w_length: 1,
                },
                &mut buffer,
            )
            .await?;
        let max_lun = if len < 1 {
            0
        } else {
            unsafe { buffer[0].assume_init() }
        };

        let bbb = Bbb::new(device.clone(), in_pipe, out_pipe)?;
        let scsi = ScsiEnclosure::setup(Box::new(bbb), max_lun as usize + 1)
            .await
            .inspect_err(|error| log::error!("msc: scsi error {error:?}"))
            .map_err(|_| UsbError::DriverError)?;
        let detach = DetachHandler(scsi.clone());
        device.set_detach_handler(Arc::new(detach));

        Ok(())
    }

    fn name(&self) -> &'static str {
        "USB Mass Storage"
    }

    fn probe(&self, class: &UsbClassInfo, _device: &UsbDeviceAccess) -> bool {
        // TODO support other protocols
        class.class == UsbDeviceClass::MassStorage && class.interface_protocol_number == 0x50
    }
}
