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::{ScsiEnclosure, transport::ScsiTransport};

use crate::{
    class_driver::UsbInterfaceDriver,
    communication::UsbDirection,
    device::{UsbDeviceDetachHandler, UsbInterfaceAccess},
    error::UsbError,
    info::UsbEndpointType,
    pipe::{
        control::ControlTransferSetup,
        normal::{UsbBulkInPipeAccess, UsbBulkOutPipeAccess},
    },
};

const BM_REQUEST_TYPE_BULK_ONLY_MASS_STORAGE_RESET: u8 = 0b00100001;
const BM_REQUEST_TYPE_GET_MAX_LUN: u8 = 0b10100001;

const B_REQUEST_BULK_ONLY_MASS_STORAGE_RESET: u8 = 0b11111111;
const B_REQUEST_GET_MAX_LUN: u8 = 0b11111110;

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)]
    interface: UsbInterfaceAccess,
    in_pipe: UsbBulkInPipeAccess,
    out_pipe: UsbBulkOutPipeAccess,
    last_tag: u32,
}

struct DetachHandler(Arc<ScsiEnclosure>);

impl Bbb {
    pub fn new(
        interface: UsbInterfaceAccess,
        in_pipe: UsbBulkInPipeAccess,
        out_pipe: UsbBulkOutPipeAccess,
    ) -> Result<Self, UsbError> {
        Ok(Self {
            interface,
            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
    }
}

#[async_trait]
impl UsbDeviceDetachHandler for DetachHandler {
    async fn handle_device_detach(&self) {
        self.0.detach().await;
    }
}

#[async_trait]
impl UsbInterfaceDriver for UsbMassStorageDriverBulkOnly {
    async fn run(self: Arc<Self>, interface: UsbInterfaceAccess) -> Result<(), UsbError> {
        let endpoints = interface.endpoints();
        let bulk_in = endpoints
            .iter()
            .find(|ep| ep.ty == UsbEndpointType::Bulk && ep.direction == UsbDirection::In);
        let bulk_out = endpoints
            .iter()
            .find(|ep| ep.ty == UsbEndpointType::Bulk && ep.direction == UsbDirection::Out);

        let (Some(bulk_in), Some(bulk_out)) = (bulk_in, bulk_out) else {
            log::warn!(
                "{}: BBB mass storage needs at least 2 bulk endpoints",
                interface.address()
            );
            return Err(UsbError::InvalidConfiguration);
        };

        let control_pipe = interface.device().control_pipe();
        let in_pipe = interface
            .device()
            .open_bulk_in_pipe(bulk_in.number, bulk_in.max_packet_size as _)
            .await?;
        let out_pipe = interface
            .device()
            .open_bulk_out_pipe(bulk_out.number, bulk_out.max_packet_size as _)
            .await?;

        // Perform a Bulk-Only Mass Storage Reset
        control_pipe
            .control_transfer(ControlTransferSetup {
                bm_request_type: BM_REQUEST_TYPE_BULK_ONLY_MASS_STORAGE_RESET,
                b_request: B_REQUEST_BULK_ONLY_MASS_STORAGE_RESET,
                w_value: 0,
                w_index: interface.number() as _,
                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: BM_REQUEST_TYPE_GET_MAX_LUN,
                    b_request: B_REQUEST_GET_MAX_LUN,
                    w_value: 0,
                    w_index: interface.number() as _,
                    w_length: 1,
                },
                &mut buffer,
            )
            .await?;
        let max_lun = if len < 1 {
            0
        } else {
            unsafe { buffer[0].assume_init() }
        };

        let bbb = Bbb::new(interface.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());
        interface.device().set_detach_handler(Arc::new(detach));

        Ok(())
    }

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

    fn probe(&self, class: u8, subclass: u8, protocol: u8) -> bool {
        // TODO support other protocols
        let _ = subclass;
        class == 0x08 && protocol == 0x50
    }
}
