use core::mem::size_of;

use alloc::string::String;
use bytemuck::{Pod, Zeroable};
use libk::dma::BusAddress;
use libk_util::{ConstAssert, IsTrue};
use static_assertions::const_assert_eq;

use crate::{
    command::{AtaCommand, AtaIdentify, AtaIdentifyResponse},
    error::AhciError,
    MAX_PRD_SIZE,
};

pub const COMMAND_LIST_LENGTH: usize = 32;

const AHCI_FIS_REG_H2D_COMMAND: u8 = 1 << 7;
const AHCI_FIS_REG_H2D: u8 = 0x27;

#[repr(C)]
pub struct AtaString<const N: usize>
where
    ConstAssert<{ N % 2 == 0 }>: IsTrue,
{
    data: [u8; N],
}

#[derive(Debug, Clone, Copy, Zeroable)]
#[repr(C)]
pub struct PhysicalRegionDescriptor {
    buffer_address: u64,
    _0: u32,
    dbc: u32,
}

#[derive(Debug, Clone, Copy, Zeroable)]
#[repr(C)]
pub struct CommandListEntry {
    attr: u16,
    prdtl: u16,
    prdbc: u32,
    ctba: u64,
    _0: [u32; 4],
}

#[derive(Clone, Copy, Zeroable)]
#[repr(C)]
pub union SentFis {
    reg_h2d: RegisterHostToDeviceFis,
    raw: RawFis,
}

#[derive(Clone, Copy, Zeroable)]
#[repr(C)]
pub struct RegisterHostToDeviceFis {
    pub ty: u8,
    pub cmd_port: u8,
    pub cmd: u8,
    pub feature_low: u8,

    pub lba0: u8,
    pub lba1: u8,
    pub lba2: u8,
    pub device: u8,

    pub lba3: u8,
    pub lba4: u8,
    pub lba5: u8,
    pub feature_high: u8,

    pub count: u16,
    pub icc: u8,
    pub control: u8,

    _0: u32,
}

#[derive(Clone, Copy, Zeroable, Pod)]
#[repr(C)]
pub struct RawFis {
    pub bytes: [u8; 64],
}

#[derive(Clone, Copy)]
#[repr(C)]
pub struct CommandTable {
    fis: SentFis,                          // 0x00..0x40
    _0: [u8; 16],                          // 0x40..0x50
    _1: [u8; 48],                          // 0x50..0x80
    prdt: [PhysicalRegionDescriptor; 248], // 0x80..0x1000
}

#[derive(Clone, Copy, Zeroable, Pod)]
#[repr(C)]
pub struct ReceivedFis {
    _dsfis: [u8; 0x1C],  // 0x00..0x1C
    _0: [u8; 0x04],      // 0x1C..0x20
    _psfis: [u8; 0x14],  // 0x20..0x34
    _1: [u8; 0x0C],      // 0x34..0x40
    _rfis: [u8; 0x14],   // 0x40..0x54
    _2: [u8; 0x04],      // 0x54..0x58
    _sdbfis: [u8; 0x08], // 0x58..0x60
    _ufis: [u8; 0x40],   // 0x60..0xA0
    _3: [u8; 0x60],      // 0xA0..0x100
}

const_assert_eq!(size_of::<SentFis>(), 0x40);
const_assert_eq!(size_of::<CommandTable>(), 0x1000);
const_assert_eq!(size_of::<CommandListEntry>(), 32);
const_assert_eq!(size_of::<ReceivedFis>(), 0x100);

impl CommandTable {
    pub fn setup_command<C: AtaCommand>(&mut self, command: &C) -> Result<(), AhciError> {
        let lba = command.lba();
        assert_eq!(lba & !0xFFFFFFFFFF, 0);
        let count = command.sector_count().try_into().unwrap();

        if C::COMMAND_ID == AtaIdentify::COMMAND_ID {
            self.fis = SentFis {
                reg_h2d: RegisterHostToDeviceFis {
                    ty: AHCI_FIS_REG_H2D,
                    cmd_port: AHCI_FIS_REG_H2D_COMMAND,
                    cmd: C::COMMAND_ID as _,

                    ..RegisterHostToDeviceFis::zeroed()
                },
            };
        } else {
            self.fis = SentFis {
                reg_h2d: RegisterHostToDeviceFis {
                    ty: AHCI_FIS_REG_H2D,
                    cmd_port: AHCI_FIS_REG_H2D_COMMAND,
                    cmd: C::COMMAND_ID as _,
                    device: 1 << 6, // LBA mode
                    lba0: lba as u8,
                    lba1: (lba >> 8) as u8,
                    lba2: (lba >> 16) as u8,
                    lba3: (lba >> 24) as u8,
                    lba4: (lba >> 32) as u8,
                    lba5: (lba >> 40) as u8,
                    count,

                    ..RegisterHostToDeviceFis::zeroed()
                },
            };
        }

        // Setup PRDs
        if let Some((base, size)) = command.buffer() {
            let mut remaining = size;
            let mut prd = 0;
            while remaining != 0 {
                let rem = remaining.min(MAX_PRD_SIZE);
                let last = rem <= MAX_PRD_SIZE;
                log::trace!(
                    target: "io",
                    "prd[{prd}] base={:#x}, size={rem}",
                    base.add(prd * MAX_PRD_SIZE)
                );

                self.prdt[prd] =
                    PhysicalRegionDescriptor::new(base.add(prd * MAX_PRD_SIZE), rem, last)?;

                prd += 1;
                remaining -= rem;
            }
            assert_eq!(prd, command.prd_count());

            self.prdt[prd..].fill_with(PhysicalRegionDescriptor::zeroed);
        }

        Ok(())
    }
}

impl CommandListEntry {
    pub fn new(command_table_entry: BusAddress, prd_count: usize) -> Result<Self, AhciError> {
        if prd_count > 0xFFFF {
            todo!()
        }
        Ok(Self {
            // attr = FIS size in dwords
            attr: (size_of::<RegisterHostToDeviceFis>() / size_of::<u32>()) as _,
            prdtl: prd_count as _,
            prdbc: 0,
            ctba: command_table_entry.into_u64(),
            _0: [0; 4],
        })
    }
}

unsafe impl Zeroable for CommandTable {
    fn zeroed() -> Self {
        Self {
            fis: SentFis::zeroed(),
            _0: [0; 16],
            _1: [0; 48],
            prdt: [PhysicalRegionDescriptor::zeroed(); 248],
        }
    }
}

impl PhysicalRegionDescriptor {
    pub fn new(address: BusAddress, byte_count: usize, is_last: bool) -> Result<Self, AhciError> {
        if byte_count > MAX_PRD_SIZE {
            return Err(AhciError::RegionTooLarge);
        }

        let dbc_mask = (is_last as u32) << 31;
        Ok(Self {
            buffer_address: address.into_u64(),
            _0: 0,
            dbc: ((byte_count as u32 - 1) << 1) | 1 | dbc_mask,
        })
    }
}

impl AtaIdentifyResponse {
    pub fn logical_sector_count(&self) -> u64 {
        // If logical_sector_count_28 == 0x0FFFFFFF, and logical_sector_count_qw >= 0x0FFFFFFF,
        //      then ACCEESSIBLE CAPACITY (?) field contains the total number of user addressable
        //      LBAs (see 4.1)
        // bit 3 in additional_features -> logical_sector_count_qw:
        //      0 -> max value = 0xFFFFFFFFFFFF (48)
        //      1 -> max value = 0xFFFFFFFF (32)
        //  If bit 3 in additional_features is set, ext_logical_sector_count_qw contains
        //      the maximum addressable LBA count. Max value = 0xFFFFFFFFFFFF (48)

        if self.command_sets[1] & (1 << 10) != 0 {
            // 48-bit supported
            if self.additional_features & (1 << 3) != 0 {
                // Use ext_logical_sector_count_qw
                todo!()
            } else {
                // Use logical_sector_count_qw
                self.logical_sector_count_qw
            }
        } else {
            todo!()
        }
    }
}

impl<const N: usize> AtaString<N>
where
    ConstAssert<{ N % 2 == 0 }>: IsTrue,
{
    #[allow(clippy::inherent_to_string)]
    pub fn to_string(&self) -> String {
        let mut buf = [0; N];

        for i in (0..N).step_by(2) {
            buf[i] = self.data[i + 1];
            buf[i + 1] = self.data[i];
        }
        let mut len = 0;
        for i in (0..N).rev() {
            if buf[i] != b' ' {
                len = i + 1;
                break;
            }
        }

        String::from(core::str::from_utf8(&buf[..len]).unwrap())
    }
}
