use core::mem::{size_of, MaybeUninit};

use device_api::dma::DmaAllocator;
use libk::dma::{BusAddress, DmaBuffer, DmaSliceMut};
use tock_registers::register_structs;

use crate::{data::AtaString, error::AhciError, MAX_PRD_SIZE};

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[repr(u8)]
pub enum AtaCommandId {
    Identify = 0xEC,
    ReadDmaEx = 0x25,
}

pub trait AtaCommand {
    type Response;

    const COMMAND_ID: AtaCommandId;

    fn lba(&self) -> u64;
    fn sector_count(&self) -> usize;
    fn buffer(&self) -> Option<(BusAddress, usize)>;
    unsafe fn into_response(self) -> Self::Response;

    fn prd_count(&self) -> usize {
        if let Some((_, size)) = self.buffer() {
            size.div_ceil(MAX_PRD_SIZE)
        } else {
            0
        }
    }
}

register_structs! {
    // Offsets in the ATA8-ACS spec are in words, so each value take from there is
    // multiplied by two
    pub AtaIdentifyResponse {
        (0  => _0),
        (20 => pub serial_number: AtaString<20>),
        (40 => _1),
        (54 => pub model_number: AtaString<40>),
        (94 => _2),
        (98 => pub capabilities: [u16; 2]),
        (102 => _3),
        (120 => pub logical_sector_count_28: [u8; 4]),
        (124 => _4),
        (138 => pub additional_features: u16),
        (140 => _5),
        (164 => pub command_sets: [u16; 6]),
        (176 => _6),
        (200 => pub logical_sector_count_qw: u64),
        (208 => _7),
        (212 => pub phys_logical_sector_size: u16),
        (214 => _8),
        (234 => pub logical_sector_size: [u16; 2]),
        (238 => _9),
        (460 => pub ext_logical_sector_count_qw: [u8; 8]),
        (468 => _10),
        (512 => @END),
    }
}

pub struct AtaIdentify {
    buffer: DmaBuffer<MaybeUninit<AtaIdentifyResponse>>,
}

pub struct AtaReadDmaEx {
    lba: u64,
    sector_count: usize,
    buffer_base: BusAddress,
    buffer_size: usize,
}

impl AtaIdentify {
    pub fn create(dma: &dyn DmaAllocator) -> Result<Self, AhciError> {
        DmaBuffer::new_uninit(dma)
            .map(Self::with_data)
            .map_err(AhciError::MemoryError)
    }

    pub fn with_data(buffer: DmaBuffer<MaybeUninit<AtaIdentifyResponse>>) -> Self {
        Self { buffer }
    }
}

impl AtaReadDmaEx {
    pub fn new(lba: u64, sector_count: usize, buffer: DmaSliceMut<MaybeUninit<u8>>) -> Self {
        Self {
            lba,
            sector_count,
            buffer_base: buffer.bus_address(),
            buffer_size: buffer.len(),
        }
    }
}

impl AtaCommand for AtaIdentify {
    type Response = DmaBuffer<AtaIdentifyResponse>;

    const COMMAND_ID: AtaCommandId = AtaCommandId::Identify;

    fn lba(&self) -> u64 {
        0
    }

    fn sector_count(&self) -> usize {
        0
    }

    fn buffer(&self) -> Option<(BusAddress, usize)> {
        let base = self.buffer.bus_address();
        let size = size_of::<AtaIdentifyResponse>();
        Some((base, size))
    }

    unsafe fn into_response(self) -> Self::Response {
        DmaBuffer::assume_init(self.buffer)
    }
}

impl AtaCommand for AtaReadDmaEx {
    type Response = ();

    const COMMAND_ID: AtaCommandId = AtaCommandId::ReadDmaEx;

    fn lba(&self) -> u64 {
        self.lba
    }

    fn sector_count(&self) -> usize {
        self.sector_count
    }

    fn buffer(&self) -> Option<(BusAddress, usize)> {
        Some((self.buffer_base, self.buffer_size))
    }

    unsafe fn into_response(self) -> Self::Response {}
}
