#![feature(generic_const_exprs, maybe_uninit_slice)]
#![allow(incomplete_features)]
#![no_std]

use core::{mem::MaybeUninit, time::Duration};

use alloc::{
    boxed::Box, collections::btree_map::BTreeMap, format, string::String, sync::Arc, vec::Vec,
};
use async_trait::async_trait;
use command::{ScsiReadCapacity, ScsiRequestSense, ScsiTestUnitReady};
use device_api::device::Device;
use libk::{
    device::{block::BlockDevice, manager::probe_partitions},
    error::Error,
    fs::devfs,
    task::runtime,
};
use libk_mm::{address::PhysicalAddress, table::MapAttributes, PageProvider, PageSlice};
use libk_util::{
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
    OneTimeInit,
};
use transport::{ScsiTransport, ScsiTransportWrapper};
use yggdrasil_abi::io::FileMode;

extern crate alloc;

pub mod command;
pub mod device;
pub mod transport;

// TODO SCSI detach
pub struct ScsiDevice {
    transport: IrqSafeSpinlock<ScsiTransportWrapper>,
    lba_count: u64,
    lba_size: usize,
    max_lba_per_request: usize,
    index: OneTimeInit<u32>,
    names: IrqSafeRwLock<Vec<String>>,
}

impl ScsiDevice {
    // TODO support LUNs other than 0
    pub async fn setup<T: ScsiTransport + 'static>(transport: T) -> Result<Arc<Self>, Error> {
        let mut transport = ScsiTransportWrapper::new(transport);

        let mut attempts = 5;
        let mut timeout = 100;
        while attempts > 0 {
            // TEST UNIT READY (6)
            if transport
                .perform_command(0, ScsiTestUnitReady)
                .await
                .is_ok()
            {
                break;
            }

            // If not, send a REQUEST SENSE (6)
            if transport.perform_command(0, ScsiRequestSense).await.is_ok() {
                break;
            }

            log::warn!("scsi: unit not ready [{attempts}/5]");

            runtime::sleep(Duration::from_millis(timeout)).await;
            timeout *= 2;
            attempts -= 1;
        }
        if attempts == 0 {
            log::error!("scsi: unit not ready");
            return Err(Error::DoesNotExist);
        }

        // TODO INQUIRY fails for real USB flash drives
        // transport.perform_command(0, ScsiInquiry).await?;

        let capacity_info = transport.perform_command(0, ScsiReadCapacity).await?;
        let max_lba_per_request =
            transport.max_bytes_per_request() / capacity_info.block_size as usize;
        log::info!(
            "scsi: lba_size={}, lba_count={}, max_lba_per_request={}",
            capacity_info.block_size,
            capacity_info.block_count,
            max_lba_per_request
        );

        Ok(Arc::new(Self {
            transport: IrqSafeSpinlock::new(transport),
            lba_count: capacity_info.block_count.into(),
            lba_size: capacity_info.block_size as usize,
            max_lba_per_request,
            index: OneTimeInit::new(),
            names: IrqSafeRwLock::new(Vec::new()),
        }))
    }

    pub fn detach(&self) {
        if let Some(&index) = self.index.try_get() {
            detach(index);
        }
    }
}

#[async_trait]
impl BlockDevice for ScsiDevice {
    // TODO avoid copies by reading directly into the cache?
    async fn read_aligned(
        &self,
        position: u64,
        buffer: &mut PageSlice<MaybeUninit<u8>>,
    ) -> Result<(), Error> {
        if buffer.len() % self.lba_size != 0 {
            log::warn!("scsi: buffer is not multiple of LBA size");
            return Err(Error::InvalidArgument);
        }

        let lba_start = position / self.lba_size as u64;
        let lba_count = buffer.len() / self.lba_size;
        let lba_end = lba_start + lba_count as u64;

        if lba_start.saturating_add(lba_count as u64) >= self.lba_count {
            log::warn!("scsi: read beyond medium end");
            return Err(Error::InvalidArgument);
        }

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

        let mut offset = 0;
        for i in (0..lba_count).step_by(self.max_lba_per_request) {
            let lba = lba_start + i as u64;
            let end = (lba + self.max_lba_per_request as u64).min(lba_end);
            let count = (end - lba) as usize;
            let amount = count * self.lba_size;

            let slice =
                unsafe { MaybeUninit::slice_assume_init_mut(&mut buffer[offset..offset + amount]) };
            let len = transport.read(0, lba, count as _, slice).await?;
            if len != amount {
                return Err(Error::InvalidArgument);
            }
            offset += amount;
        }

        Ok(())
    }

    async fn write_aligned(&self, _position: u64, _buffer: &PageSlice<u8>) -> Result<(), Error> {
        // TODO AtaWriteDmaEx
        Err(Error::NotImplemented)
    }

    fn block_size(&self) -> usize {
        self.lba_size
    }

    fn block_count(&self) -> u64 {
        self.lba_count
    }

    fn max_blocks_per_request(&self) -> usize {
        self.max_lba_per_request
    }
}

impl PageProvider for ScsiDevice {
    fn get_page(&self, _offset: u64) -> Result<PhysicalAddress, Error> {
        Err(Error::NotImplemented)
    }

    fn release_page(&self, _offset: u64, _phys: PhysicalAddress) -> Result<(), Error> {
        Err(Error::NotImplemented)
    }

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

impl Device for ScsiDevice {
    fn display_name(&self) -> &str {
        "SCSI Storage Device"
    }
}

impl Drop for ScsiDevice {
    fn drop(&mut self) {
        if let Some(index) = self.index.try_get() {
            log::info!("scsi{index} dropped");
        }
    }
}

// TODO this is crap
static SCSI_DEVICES: IrqSafeSpinlock<BTreeMap<u32, Arc<ScsiDevice>>> =
    IrqSafeSpinlock::new(BTreeMap::new());
static SCSI_BITMAP: IrqSafeSpinlock<u32> = IrqSafeSpinlock::new(0);

pub fn attach(device: Arc<ScsiDevice>) -> Result<(), Error> {
    let index = {
        let mut bitmap = SCSI_BITMAP.lock();
        let index = (0..8)
            .position(|p| *bitmap & (1 << p) == 0)
            .ok_or(Error::InvalidOperation)
            .inspect_err(|_| log::warn!("Cannot attach SCSI device: too many of them"))?
            as u32;
        let mut devices = SCSI_DEVICES.lock();
        *bitmap |= 1 << index;
        assert!(!devices.contains_key(&index));
        devices.insert(index, device.clone());
        index
    };

    let name = format!("scsi{index}");
    device.index.init(index);
    device.names.write().push(name.clone());
    devfs::add_named_block_device(device.clone(), name.clone(), FileMode::new(0o600)).ok();
    log::info!("{name} attached");

    // TODO this code is repeated everywhere
    runtime::spawn(async move {
        let name = name;
        probe_partitions(device.clone(), |index, partition| {
            let partition_name = format!("{name}p{}", index + 1);
            log::info!("{name}: partition {partition_name}");
            device.names.write().push(partition_name.clone());
            devfs::add_named_block_device(
                Arc::new(partition),
                partition_name,
                FileMode::new(0o600),
            )
            .ok();
        })
        .await
        .ok();
    })
    .ok();

    Ok(())
}

pub fn detach(index: u32) {
    let mut devices = SCSI_DEVICES.lock();
    let mut bitmap = SCSI_BITMAP.lock();

    if let Some(device) = devices.remove(&index) {
        {
            let names = device.names.read();
            for name in names.iter() {
                devfs::remove_node(name).ok();
            }
        }

        *bitmap &= !(1 << index);
        log::info!("scsi{index} detached");
    }
}