use core::array;
use core::ops::Deref;

use alloc::{boxed::Box, sync::Arc};
use mammoth::{mem, sync::Mutex, zion::ZError};

use crate::ahci::port::AhciPortInterruptStatus;

use super::{
    ahci_command::{CommandList, CommandTable, FisType, HostToDeviceRegisterFis, ReceivedFis},
    command::{CommandFuture, CommandStatus},
    port::AhciPortHba,
    Command,
};

struct CommandStructures {
    command_list: &'static mut CommandList,
    received_fis: &'static mut ReceivedFis,
    command_tables: &'static mut [CommandTable; 32],
}

pub struct PortController {
    ahci_port_hba: Arc<Mutex<&'static mut AhciPortHba>>,
    command_structures: Mutex<CommandStructures>,

    command_slots: [Arc<Mutex<CommandStatus>>; 32],
}

impl PortController {
    pub fn new(ahci_port_hba: &'static mut AhciPortHba) -> Self {
        let (command_vaddr, command_paddr) = mem::map_physical_and_leak(0x2500);
        ahci_port_hba.init(command_paddr, command_paddr + 0x400);

        let command_list = unsafe { (command_vaddr as *mut CommandList).as_mut().unwrap() };
        let received_fis = unsafe {
            ((command_vaddr + 0x400) as *mut ReceivedFis)
                .as_mut()
                .unwrap()
        };
        let command_tables = unsafe {
            ((command_vaddr + 0x500) as *mut [CommandTable; 32])
                .as_mut()
                .unwrap()
        };

        // This leaves space for 8 prdt entries.
        for (i, header) in command_list.iter_mut().enumerate() {
            header.command_table_base_addr = (command_paddr + 0x500) + (0x100 * (i as u64));
        }

        let command_slots = array::from_fn(|_| Arc::new(Mutex::new(CommandStatus::Empty)));

        Self {
            ahci_port_hba: Arc::new(Mutex::new(ahci_port_hba)),
            command_structures: Mutex::new(CommandStructures {
                command_list,
                received_fis,
                command_tables,
            }),
            command_slots,
        }
    }

    pub fn get_signature(&self) -> u32 {
        self.ahci_port_hba.lock().signature.read()
    }

    pub fn issue_command(&self, command: &Command) -> Result<CommandFuture, ZError> {
        let slot = self.select_slot()?;
        let command_slot = self.command_slots[slot].clone();

        let ahci_port_hba_clone = self.ahci_port_hba.clone();
        let slot_clone = slot;
        let future = CommandFuture::new(
            command_slot,
            Box::new(move || {
                ahci_port_hba_clone.lock().issue_command(slot_clone);
            }),
        );

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

        command_structures.command_tables[slot]
            .command_fis
            .host_to_device = command.into();

        command_structures.command_tables[slot].prdt[0].region_address = command.paddr;
        command_structures.command_tables[slot].prdt[0].byte_count =
            512 * (command.sector_cnt as u32);

        command_structures.command_list[slot].prd_table_length = 1;

        command_structures.command_list[slot].command =
            (size_of::<HostToDeviceRegisterFis>() as u16 / 4) & 0x1F;
        command_structures.command_list[slot].command |= 1 << 7;

        Ok(future)
    }

    fn select_slot(&self) -> Result<usize, ZError> {
        // TODO: We have a race condition here.
        for i in 0..self.command_slots.len() {
            let mut slot = self.command_slots[i].lock();
            if matches!(*slot, CommandStatus::Empty) {
                *slot = CommandStatus::NotSent;
                return Ok(i);
            }
        }
        Err(ZError::EXHAUSTED)
    }

    pub fn handle_interrupt(&self) {
        let int_status = self.ahci_port_hba.lock().interrupt_status.read();
        if int_status.device_to_host_register_fis_interrupt() {
            let received_fis = self
                .command_structures
                .lock()
                .received_fis
                .device_to_host_register_fis;
            assert_eq!(
                received_fis.fis_type as u8,
                FisType::RegisterDeviceToHost as u8
            );
            if received_fis.error != 0 {
                mammoth::debug!("D2H err: {:#0x}", received_fis.error);
                mammoth::debug!("Status: {:#0x}", received_fis.status);
            }

            self.ahci_port_hba.lock().interrupt_status.write(
                AhciPortInterruptStatus::new().with_device_to_host_register_fis_interrupt(true),
            );
        }
        if int_status.pio_setup_fis_interrupt() {
            self.ahci_port_hba
                .lock()
                .interrupt_status
                .write(AhciPortInterruptStatus::new().with_pio_setup_fis_interrupt(true));
        }

        for i in 0..32 {
            let int_offset = 1 << i;

            // If there is no longer a command issued on a slot and we have something in
            // the command list that is pending we know that this is the command that finished.
            if (self.ahci_port_hba.lock().command_issue.read() & int_offset) != int_offset {
                let mut command_status = self.command_slots[i].lock();
                let mut did_complete = false;
                if let CommandStatus::Pending(ref waker) = command_status.deref() {
                    waker.wake_by_ref();
                    did_complete = true;
                }

                if did_complete {
                    *command_status = CommandStatus::Complete;
                }
            }
        }
    }
}