/*
 * Copyright (c) 2022, Umut İnan Erdoğan <umutinanerdogan@pm.me>
 *
 * SPDX-License-Identifier: MPL-2.0
 */

use core::num::TryFromIntError;

// FIXME: platform agnostic-ify these
use x86_64::instructions::{interrupts, port::Port};
use x86_64::structures::paging::{mapper::MapToError, Mapper, Page, PhysFrame, Size4KiB};
use x86_64::structures::paging::{FrameAllocator, FrameDeallocator};
use x86_64::VirtAddr;

use crate::arch::memory::BootInfoFrameAllocator;
use crate::devices::pci::check_device;

use super::PciDeviceInfo;

// FIXME: un-hardcode these
const PRDT_START: u64 = 0xffff_ffff_0000_0000;
const BUFFER_START: u64 = 0xffff_ffff_0000_1000;

/// ATA logical sector size, in bytes
const SECTOR_SIZE: u16 = 512;

/// Bus Master IDE Command
const BMIC_OFFSET: u16 = 0;
/// Bus Master IDE Status
const BMIS_OFFSET: u16 = 2;
/// Bus Master IDE Descriptor Table Pointer
const BMIDTP_OFFSET: u16 = 4;

/// Bus Master IDE Secondary Offset
const BMI_SECONDARY: u16 = 8;

/// Primary command block offset
const PRIMARY_COMMAND: u16 = 0x01F0;

/// Secondary command block offset
const SECONDARY_COMMAND: u16 = 0x0170;

/// Data register offset
const DATA_OFFSET: u16 = 0;

/// Sector count register offset
const SECCOUNT_OFFSET: u16 = 2;

/// LBA0 register offset
const LBA0_OFFSET: u16 = 3;

/// LBA1 register offset
const LBA1_OFFSET: u16 = 4;

/// LBA2 register offset
const LBA2_OFFSET: u16 = 5;

/// Drive/Head register offset
const DRIVE_HEAD_OFFSET: u16 = 6;

/// Command/status register offset
const COMMAND_STATUS_OFFSET: u16 = 7;

/// Secondary control block offset
const SECONDARY_CONTROL: u16 = 0x0374;

/// Primary control block offset
const PRIMARY_CONTROL: u16 = 0x03F4;

/// Alternative status offset
const ALT_STATUS_OFFSET: u16 = 2;

/// ATA identification command
const CMD_IDENTIFY: u8 = 0xEC;

/// ATA read using LBA28 DMA command
const CMD_READ_DMA: u8 = 0xC8;

/// ATA write using LBA28 DMA command
const CMD_WRITE_DMA: u8 = 0xCA;

/// ATA read using LBA48 DMA command
const CMD_READ_DMA_EXT: u8 = 0x25;

/// ATA write using LBA48 DMA command
const CMD_WRITE_DMA_EXT: u8 = 0x35;
#[derive(Debug)]
pub struct PciIde {
    device_info: PciDeviceInfo,
    ide_devices: Vec<IdeDevice>,
    prdt_frame: Option<PhysFrame>,
    buffer_frames: Option<Vec<PhysFrame>>,
    bmiba: u16,
}

impl PciIde {
    // FIXME: make this return a Result
    pub fn new(bus: u8, device: u8) -> Option<Self> {
        let device_info = check_device(bus, device)?;
        device_info.enable_bus_mastering();
        let idetim = unsafe { device_info.read(0, 0x40) };
        trace!("idetim: {idetim:b}");
        // FIXME: enable the right bits in idetim (and sidetim) to use fast timings

        let mut ide_devices = Vec::with_capacity(4);
        for ch in 0..2 {
            let channel = if ch == 0 {
                Channel::Primary
            } else {
                Channel::Secondary
            };
            'drive: for dr in 0..2 {
                let drive = if dr == 0 { Drive::Master } else { Drive::Slave };

                unsafe {
                    select_drive(drive, channel);
                    // FIXME: clear sector count and lba0, lba1, lba2 registers
                    let status = ata_send_command(CMD_IDENTIFY, channel);
                    if status == 0 {
                        continue; // If status = 0, no device
                    }

                    loop {
                        let status = {
                            let addr = if channel.secondary() {
                                SECONDARY_COMMAND
                            } else {
                                PRIMARY_COMMAND
                            } + COMMAND_STATUS_OFFSET;
                            Port::<u8>::new(addr).read()
                        };

                        if status & 1 == 1 {
                            // if error (bit 0), device is not ATA
                            // FIXME: ATAPI devices
                            continue 'drive;
                        }
                        if !((status >> 7) & 1 == 1) && (status >> 3) & 1 == 1 {
                            // BSY cleared, DRQ set, everything is right
                            break;
                        }
                    }

                    // Read identification space of device
                    let addr = if channel.secondary() {
                        SECONDARY_COMMAND
                    } else {
                        PRIMARY_COMMAND
                    } + DATA_OFFSET;
                    let mut buffer = [0_u8; 512];
                    read_dword_buffer(addr, buffer.as_mut_ptr() as *mut _, 128);
                    // for (i, byte) in buffer.iter().enumerate() {
                    //     if byte.is_ascii() {
                    //         trace!("byte {i}: {byte:b}, ascii: {}", *byte as char);
                    //     } else {
                    //         trace!("byte {i}: {byte:b}");
                    //     }
                    // }

                    if buffer[99] & 1 != 1 {
                        // FIXME: PIO mode support
                        error!("IDE drive {channel:?}/{drive:?} does not support DMA");
                        continue;
                    }

                    // FIXME: CHS support
                    let lba48 = (buffer[167] >> 2) & 1 == 1;

                    let size = buffer[200] as u64
                        | (buffer[201] as u64) << 8
                        | (buffer[202] as u64) << 16
                        | (buffer[203] as u64) << 24
                        | (buffer[204] as u64) << 32
                        | (buffer[205] as u64) << 40
                        | (buffer[206] as u64) << 48
                        | (buffer[207] as u64) << 54;
                    trace!("IDE drive {channel:?}/{drive:?} has {size} sectors");

                    ide_devices.push(IdeDevice {
                        channel,
                        drive,
                        size,
                        lba48_support: lba48,
                    });
                }
            }
        }

        let bmiba = device_info.bar(4) & 0xFFFFFFFC;

        Some(Self {
            device_info,
            ide_devices,
            prdt_frame: None,
            buffer_frames: None,
            bmiba: bmiba.try_into().ok()?,
        })
    }

    pub fn allocate_dma_frame(
        &mut self,
        mapper: &mut impl Mapper<Size4KiB>,
        frame_allocator: &mut BootInfoFrameAllocator,
    ) -> Result<(), MapToError<Size4KiB>> {
        use x86_64::structures::paging::PageTableFlags as Flags;

        let prdt_frame = frame_allocator
            .allocate_frame()
            .ok_or(MapToError::FrameAllocationFailed)?;

        let buffer_frames = {
            let mut frame = frame_allocator
                .allocate_frame()
                .ok_or(MapToError::FrameAllocationFailed)?;
            while !frame.start_address().is_aligned(0x10000u64) {
                unsafe {
                    frame_allocator.deallocate_frame(frame);
                }

                frame = frame_allocator
                    .allocate_frame()
                    .ok_or(MapToError::FrameAllocationFailed)?;
            }

            let mut frames = Vec::with_capacity(16);
            frames.push(frame);
            for _ in 0..15 {
                let frame = frame_allocator
                    .allocate_frame()
                    .ok_or(MapToError::FrameAllocationFailed)?;
                frames.push(frame);
            }

            frames
        };
        let flags = Flags::NO_CACHE | Flags::PRESENT | Flags::WRITABLE;

        unsafe {
            mapper
                .map_to(
                    Page::containing_address(VirtAddr::new(PRDT_START)),
                    prdt_frame,
                    flags,
                    frame_allocator,
                )?
                .flush();

            for (i, frame) in buffer_frames.iter().enumerate() {
                mapper
                    .map_to(
                        Page::containing_address(VirtAddr::new(BUFFER_START + i as u64 * 0x1000)),
                        *frame,
                        flags,
                        frame_allocator,
                    )?
                    .flush()
            }
        }

        self.prdt_frame = Some(prdt_frame);
        self.buffer_frames = Some(buffer_frames);
        Ok(())
    }

    pub fn read(
        &mut self,
        channel: Channel,
        drive: Drive,
        lba: u64,
        sector_count: u16,
        buffer: &mut Vec<u8>,
    ) -> Result<(), TryFromIntError> {
        let lba48_support = self
            .ide_devices
            .iter()
            .find(|d| d.channel == channel && d.drive == drive)
            .map(|d| d.lba48_support)
            .unwrap(); // FIXME: make this an error
        let lba48 = lba > 0xFFFFFFF && lba48_support;

        // FIXME: make this an error
        assert!((lba48 && lba > 0xFFFFFFF) || (!lba48 && lba <= 0xFFFFFFF));

        let byte_count = sector_count * SECTOR_SIZE;

        // prepare PRD table
        let prd = PRDT_START as *mut PhysRegionDescriptor;
        unsafe {
            (*prd).data_buffer = self.buffer_frames.as_ref().unwrap()[0]
                .start_address()
                .as_u64()
                .try_into()?;
            (*prd).byte_count = byte_count;
            // this is the end of table
            (*prd).eot = 1 << 7;
            // this byte is reserved, we should probably set it to 0
            (*prd)._0 = 0;
        }

        unsafe {
            self.load_prdt(channel);
            self.stop(channel);
            self.set_read(channel);
            self.clear_bmi_status(channel);
            select_drive(drive, channel);
            set_lba(channel, lba, sector_count, lba48);

            if lba48 {
                ata_send_command(CMD_READ_DMA_EXT, channel);
            } else {
                ata_send_command(CMD_READ_DMA, channel);
            }

            self.start(channel);
        }

        loop {
            let status = unsafe { self.bmi_status(channel) };

            // FIXME: error handling

            // Bit 2 (INT) set?
            if (status >> 2) & 1 == 1 {
                break;
            }
        }

        unsafe {
            // Stop DMA
            self.stop(channel);

            // Clear the interrupt bit
            self.clear_bmi_status(channel);
        }

        for i in 0..byte_count as u64 {
            let addr = (BUFFER_START + i) as *mut u8;
            buffer.push(unsafe { *addr });
        }

        Ok(())
    }

    pub fn write(
        &mut self,
        channel: Channel,
        drive: Drive,
        lba: u64,
        data: &[u8],
    ) -> Result<(), TryFromIntError> {
        // FIXME: make this an error
        assert!(data.len() % SECTOR_SIZE as usize == 0);

        let lba48_support = self
            .ide_devices
            .iter()
            .find(|d| d.channel == channel && d.drive == drive)
            .map(|d| d.lba48_support)
            .unwrap(); // FIXME: make this an error
        let lba48 = lba > 0xFFFFFFF && lba48_support;

        // FIXME: make this an error
        assert!((lba48 && lba > 0xFFFFFFF) || (!lba48 && lba <= 0xFFFFFFF));

        let byte_count = data.len() as u16;
        let sector_count = byte_count / SECTOR_SIZE;

        // prepare PRD table
        let prd = PRDT_START as *mut PhysRegionDescriptor;
        unsafe {
            (*prd).data_buffer = self.buffer_frames.as_ref().unwrap()[0]
                .start_address()
                .as_u64()
                .try_into()?;
            (*prd).byte_count = byte_count;
            // this is the end of table
            (*prd).eot = 1 << 7;
            // this byte is reserved, we should probably set it to 0
            (*prd)._0 = 0;
        }

        // copy the data over to the DMA buffer
        for i in 0..byte_count {
            let addr = (BUFFER_START + i as u64) as *mut u8;
            unsafe {
                *addr = *data.get(i as usize).unwrap_or(&0);
            }
        }

        unsafe {
            self.load_prdt(channel);
            self.stop(channel);
            self.set_write(channel);
            self.clear_bmi_status(channel);
            select_drive(drive, channel);
            set_lba(channel, lba, sector_count, lba48);

            if lba48 {
                ata_send_command(CMD_WRITE_DMA_EXT, channel);
            } else {
                ata_send_command(CMD_WRITE_DMA, channel);
            }

            self.start(channel);
        }

        loop {
            let status = unsafe { self.bmi_status(channel) };

            // FIXME: error handling

            // Bit 2 (INT) set?
            if (status >> 2) & 1 == 1 {
                break;
            }
        }

        unsafe {
            // Stop DMA
            self.stop(channel);

            // Clear the interrupt bit
            self.clear_bmi_status(channel);
        }

        Ok(())
    }

    pub fn device_info(&self) -> PciDeviceInfo {
        self.device_info
    }

    unsafe fn load_prdt(&self, channel: Channel) {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIDTP_OFFSET;
        Port::<u32>::new(addr).write(
            self.prdt_frame
                .unwrap()
                .start_address()
                .as_u64()
                .try_into()
                .unwrap(),
        );
    }

    unsafe fn start(&self, channel: Channel) {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIC_OFFSET;
        let mut port: Port<u8> = Port::new(addr);
        let mut bmic = port.read();
        // start transfer
        bmic |= 1;
        // write the new bmic
        port.write(bmic);
    }

    unsafe fn stop(&self, channel: Channel) {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIC_OFFSET;
        let mut port: Port<u8> = Port::new(addr);
        let mut bmic = port.read();
        // stop ongoing transfer
        bmic &= !1;
        // write the new bmic
        port.write(bmic);
    }

    unsafe fn set_read(&self, channel: Channel) {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIC_OFFSET;
        let mut port: Port<u8> = Port::new(addr);
        let mut bmic = port.read();
        // mark bit 3 as 0 (read)
        bmic &= !(1 << 3);
        // write the new bmic
        port.write(bmic);
    }

    unsafe fn set_write(&self, channel: Channel) {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIC_OFFSET;
        let mut port: Port<u8> = Port::new(addr);
        let mut bmic = port.read();
        // mark bit 3 as 1 (write)
        bmic |= 1 << 3;
        // write the new bmic
        port.write(bmic);
    }

    unsafe fn bmi_status(&self, channel: Channel) -> u8 {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIS_OFFSET;
        let mut port = Port::new(addr);
        port.read()
    }

    unsafe fn clear_bmi_status(&self, channel: Channel) {
        let addr = if channel.secondary() {
            BMI_SECONDARY
        } else {
            0
        } + self.bmiba
            + BMIS_OFFSET;
        let mut port: Port<u8> = Port::new(addr);
        let mut bmis = port.read();
        // write 1 to bits 1 (DMA error) and 2 (int status) which clears them
        bmis |= 1 << 1 | 1 << 2;
        // write the new bmis
        port.write(bmis);
    }
}

unsafe fn select_drive(drive: Drive, channel: Channel) {
    let addr = if channel.secondary() {
        SECONDARY_COMMAND
    } else {
        PRIMARY_COMMAND
    } + DRIVE_HEAD_OFFSET;
    let mut port: Port<u8> = Port::new(addr);
    // FIXME: CHS support
    let drive_command = if drive.slave() {
        // slave & LBA
        0b11110000
    } else {
        // master & LBA
        0b11100000
    };

    // write the new drive/head register
    port.write(drive_command);
    ata_delay(channel);
}

/// Send ATA command and read status afterwards
unsafe fn ata_send_command(command: u8, channel: Channel) -> u8 {
    let addr = if channel.secondary() {
        SECONDARY_COMMAND
    } else {
        PRIMARY_COMMAND
    } + COMMAND_STATUS_OFFSET;
    let mut port = Port::new(addr);
    port.write(command);
    ata_delay(channel);
    port.read()
}

/// Read the alternate status register 14 times to create a ~420ns delay
unsafe fn ata_delay(channel: Channel) {
    let addr = if channel.secondary() {
        SECONDARY_CONTROL
    } else {
        PRIMARY_CONTROL
    } + ALT_STATUS_OFFSET;
    let mut port: Port<u8> = Port::new(addr);
    for _ in 0..14 {
        port.read();
    }
}

/// Set LBA and sector count registers. sector_count of 0 means 65536 sectors
unsafe fn set_lba(channel: Channel, lba: u64, sector_count: u16, lba48: bool) {
    let command_block = if channel.secondary() {
        SECONDARY_COMMAND
    } else {
        PRIMARY_COMMAND
    };

    let mut seccount = Port::new(command_block + SECCOUNT_OFFSET);
    let mut lba0 = Port::new(command_block + LBA0_OFFSET);
    let mut lba1 = Port::new(command_block + LBA1_OFFSET);
    let mut lba2 = Port::new(command_block + LBA2_OFFSET);
    let mut head = Port::new(command_block + DRIVE_HEAD_OFFSET);
    let head_value: u8 = head.read();

    let lba_bytes = lba.to_le_bytes();
    let sector_count_bytes = sector_count.to_le_bytes();

    // write the new LBA & sector count registers
    // FIXME: CHS support
    if lba48 {
        seccount.write(sector_count_bytes[1]);
        lba0.write(lba_bytes[3]);
        lba1.write(lba_bytes[4]);
        lba2.write(lba_bytes[5]);
    } else {
        head.write(head_value | (lba_bytes[3] & 0x0F));
    }
    seccount.write(sector_count_bytes[0]);
    lba0.write(lba_bytes[0]);
    lba1.write(lba_bytes[1]);
    lba2.write(lba_bytes[2]);
}

unsafe fn read_dword_buffer(port: u16, buffer: *mut u32, mut count: u32) {
    // FIXME: this assumes x86-64
    interrupts::without_interrupts(|| {
        asm!("
        cld
        repne
        insd",
            in("di") buffer,
            in("dx") port,
            inout("cx") count,
        )
    });
}

#[derive(Debug)]
struct IdeDevice {
    pub channel: Channel,
    pub drive: Drive,
    pub size: u64, // in sectors
    pub lba48_support: bool,
    // FIXME: model
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Channel {
    Primary,
    Secondary,
}

impl Channel {
    fn secondary(&self) -> bool {
        matches!(self, Self::Secondary)
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum Drive {
    Master,
    Slave,
}

impl Drive {
    fn slave(&self) -> bool {
        matches!(self, Self::Slave)
    }
}

#[repr(C, packed)]
struct PhysRegionDescriptor {
    /// Pointer to the data buffer
    pub data_buffer: u32,
    /// Byte count, 64K maximum per PRD transfer. 0 means 64K
    pub byte_count: u16,
    /// Reserved byte
    pub _0: u8,
    /// MSB marks end of transfer
    pub eot: u8,
}