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

use core::{
    mem::MaybeUninit,
    sync::atomic::{AtomicBool, Ordering},
    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::{
    block,
    device::{block::BlockDevice, manager::probe_partitions},
    dma::{DmaBuffer, DmaSlice, DmaSliceMut},
    error::Error,
    fs::devfs,
    task::{runtime, sync::AsyncMutex},
};
use libk_mm::{
    address::PhysicalAddress, table::MapAttributes, OnDemandPage, PageProvider, VirtualPage,
};
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;

pub struct ScsiEnclosure {
    transport: AsyncMutex<ScsiTransportWrapper>,
    units: Vec<IrqSafeRwLock<Option<Arc<ScsiUnit>>>>,
    index: OneTimeInit<u32>,
    shutdown: AtomicBool,
}

pub struct ScsiUnit {
    enclosure: Arc<ScsiEnclosure>,
    lun: u8,
    lba_count: u64,
    lba_size: usize,
    max_lba_per_request: usize,
    names: IrqSafeRwLock<Vec<String>>,
}

impl ScsiEnclosure {
    pub async fn setup(
        transport: Box<dyn ScsiTransport>,
        lun_count: usize,
    ) -> Result<Arc<Self>, Error> {
        let transport = AsyncMutex::new(ScsiTransportWrapper::new(transport));
        let units = (0..lun_count).map(|_| IrqSafeRwLock::new(None)).collect();
        let this = Arc::new(Self {
            transport,
            units,
            index: OneTimeInit::new(),
            shutdown: AtomicBool::new(false),
        });
        register_enclosure(this.clone())?;

        // Probe LUNs
        for i in 0..lun_count {
            if this.probe_lun(i as u8).await
                && let Ok(unit) = ScsiUnit::setup(this.clone(), i as u8).await
            {
                *this.units[i].write() = Some(unit);
            }
        }

        // Start enclosure poll task
        let enclosure = this.clone();
        runtime::spawn(async move {
            while !enclosure.shutdown.load(Ordering::Acquire) {
                enclosure.poll().await;
                runtime::sleep(Duration::from_millis(100)).await;
            }
        })
        .ok();

        Ok(this)
    }

    async fn probe_lun(self: &Arc<Self>, lun: u8) -> bool {
        let mut attempts = 3;
        let mut timeout = 10;
        // TODO get statuses to better see if there's a real error or the LUN is not present
        while attempts > 0 {
            let mut transport = self.transport.lock().await;

            // TEST UNIT READY (6)
            if transport
                .perform_command(lun, ScsiTestUnitReady)
                .await
                .is_ok()
            {
                break;
            }

            // If not, send a REQUEST SENSE (6)
            transport.perform_command(lun, ScsiRequestSense).await.ok();

            drop(transport);

            runtime::sleep(Duration::from_millis(timeout)).await;
            timeout *= 2;
            attempts -= 1;
        }

        attempts != 0
    }

    async fn poll(self: &Arc<Self>) {
        let index = *self.index.get();
        for lun in 0..self.units.len() {
            let mut slot = self.units[lun].write();
            let present = self.probe_lun(lun as u8).await;

            if let Some(unit) = slot.as_ref() {
                if !present {
                    log::warn!("scsi{index}u{lun} lost");
                    unit.detach();
                    *slot = None;
                }
            } else if present {
                if let Ok(unit) = ScsiUnit::setup(self.clone(), lun as u8).await {
                    log::info!("scsi{index}u{lun} attached");
                    *slot = Some(unit);
                } else {
                    log::warn!("scsi{index}u{lun} attached, but could not setup");
                }
            }
        }
    }

    pub fn detach(&self) {
        self.shutdown.store(true, Ordering::Release);
        let index = self.index.try_get().copied();

        for unit in self.units.iter() {
            if let Some(unit) = unit.write().take() {
                unit.detach();
            }
        }

        // Deregister the enclosure
        if let Some(index) = index {
            remove_enclosure(index);
        }
    }
}

impl ScsiUnit {
    pub async fn setup(enclosure: Arc<ScsiEnclosure>, lun: u8) -> Result<Arc<Self>, Error> {
        let enclosure_index = *enclosure.index.get();
        let mut transport = enclosure.transport.lock().await;

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

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

        drop(transport);

        let unit = Arc::new(Self {
            enclosure,
            lun,
            lba_count: capacity_info.block_count.into(),
            lba_size: capacity_info.block_size as usize,
            max_lba_per_request,
            names: IrqSafeRwLock::new(Vec::new()),
        });

        register_unit(enclosure_index, lun, unit.clone());

        Ok(unit)
    }

    fn detach(&self) {
        let id = *self.enclosure.index.get();
        log::info!("scsi{id}u{} detached", self.lun);
        for name in self.names.read().iter() {
            devfs::remove_node(name).ok();
        }
    }
}

#[async_trait]
impl BlockDevice for ScsiUnit {
    fn allocate_buffer(&self, size: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error> {
        block!(self.enclosure.transport.lock().await.allocate_buffer(size))?
    }

    async fn read_aligned(
        &self,
        position: u64,
        buffer: DmaSliceMut<'_, MaybeUninit<u8>>,
    ) -> Result<(), Error> {
        if position % self.lba_size as u64 != 0 {
            log::warn!("scsi: misaligned read");
            return Err(Error::InvalidArgument);
        }
        if buffer.len() % self.lba_size != 0 {
            log::warn!("scsi: misaligned buffer size");
            return Err(Error::InvalidArgument);
        }
        let lba_start = position / self.lba_size as u64;
        let lba_count = buffer.len() / self.lba_size;

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

        let lba_end = lba_start + lba_count as u64;

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

        // TODO DmaSliceMut subslicing
        let (buffer, range) = buffer.into_parts();
        let mut offset = range.start;

        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 dst_slice = buffer.slice_mut(offset..offset + amount);
            let len = transport
                .read(self.lun, lba, count as u16, dst_slice)
                .await?;
            if len != amount {
                return Err(Error::InvalidArgument);
            }

            offset += amount;
        }

        Ok(())
    }

    async fn write_aligned(&self, _position: u64, _buffer: DmaSlice<'_, u8>) -> Result<(), Error> {
        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 ScsiUnit {
    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!()
    }
}

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

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

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

fn register_enclosure(enclosure: Arc<ScsiEnclosure>) -> 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 enclosure: too many of them"))?
            as u32;
        let mut devices = SCSI_ENCLOSURES.lock();
        *bitmap |= 1 << index;
        assert!(!devices.contains_key(&index));
        devices.insert(index, enclosure.clone());
        index
    };

    enclosure.index.init(index);

    Ok(())
}

fn register_unit(enclosure_index: u32, lun: u8, unit: Arc<ScsiUnit>) {
    let name = format!("scsi{enclosure_index}u{lun}");
    unit.names.write().push(name.clone());
    devfs::add_named_block_device(unit.clone(), name.clone(), FileMode::new(0o600)).ok();

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

fn remove_enclosure(index: u32) {
    let mut devices = SCSI_ENCLOSURES.lock();
    let mut bitmap = SCSI_BITMAP.lock();

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