use libk::{dma::BusAddress, error::Error};
use tock_registers::{
    interfaces::{ReadWriteable, Readable, Writeable},
    register_bitfields, register_structs,
    registers::{ReadOnly, ReadWrite, WriteOnly},
};

use crate::queue::{CompletionQueueEntry, SubmissionQueueEntry};

register_bitfields! {
    u32,
    pub CC [
        IOCQES OFFSET(20) NUMBITS(4) [],
        IOSQES OFFSET(16) NUMBITS(4) [],
        AMS OFFSET(11) NUMBITS(3) [],
        MPS OFFSET(7) NUMBITS(4) [],
        CSS OFFSET(4) NUMBITS(3) [
            NvmCommandSet = 0
        ],
        ENABLE OFFSET(0) NUMBITS(1) [],
    ],
    pub CSTS [
        CFS OFFSET(1) NUMBITS(1) [],
        RDY OFFSET(0) NUMBITS(1) [],
    ],
    pub AQA [
        /// Admin Completion Queue Size in entries - 1
        ACQS OFFSET(16) NUMBITS(12) [],
        /// Admin Submission Queue Size in entries - 1
        ASQS OFFSET(0) NUMBITS(12) [],
    ]
}

register_bitfields! {
    u64,
    pub CAP [
        /// Maximum Queue Entries Supported - 1. i.e., 0 means maximum queue len of 1, 1 = 2 etc.
        MQES OFFSET(0) NUMBITS(16) [],
        /// Timeout. Represents the worst-case time the host software should wait for CSTS.RDY to
        /// change its state.
        TO OFFSET(24) NUMBITS(8) [],
        /// Doorbell stride. Stride in bytes = pow(2, 2 + DSTRD).
        DSTRD OFFSET(32) NUMBITS(4) [],
        /// NVM Subsystem Reset Supported (see NVMe BS Section 3.7.1)
        NSSRS OFFSET(36) NUMBITS(1) [],
        /// Controller supports one or more I/O command sets
        CSS_IO_COMMANDS OFFSET(43) NUMBITS(1) [],
        /// Controller only supports admin commands and no I/O commands
        CSS_ADMIN_ONLY OFFSET(44) NUMBITS(1) [],
        /// Memory page size minimum (bytes = pow(2, 12 + MPSMIN))
        MPSMIN OFFSET(48) NUMBITS(4) [],
        /// Memory page size maximum -|-
        MPSMAX OFFSET(52) NUMBITS(4) [],
    ]
}

register_structs! {
    #[allow(non_snake_case)]
    pub Regs {
        (0x00 => pub CAP: ReadOnly<u64, CAP::Register>),
        (0x08 => pub VS: ReadOnly<u32>),
        (0x0C => pub INTMS: WriteOnly<u32>),
        (0x10 => pub INTMC: WriteOnly<u32>),
        (0x14 => pub CC: ReadWrite<u32, CC::Register>),
        (0x18 => _0),
        (0x1C => pub CSTS: ReadOnly<u32, CSTS::Register>),
        (0x20 => _1),
        (0x24 => AQA: ReadWrite<u32, AQA::Register>),
        (0x28 => ASQ: ReadWrite<u64>),
        (0x30 => ACQ: ReadWrite<u64>),
        (0x38 => _2),
        (0x2000 => @END),
    }
}

impl Regs {
    pub fn configure_admin_queue(
        &self,
        submission_queue_pointer: BusAddress,
        completion_queue_pointer: BusAddress,
        submission_queue_size: usize,
        completion_queue_size: usize,
    ) -> Result<(), Error> {
        let max_queue_size = self.CAP.read(CAP::MQES) + 1;
        if submission_queue_size as u64 > max_queue_size {
            log::error!("admin submission queue too large");
            return Err(Error::InvalidArgument);
        }
        if completion_queue_size as u64 > max_queue_size {
            log::error!("admin completion queue too large");
            return Err(Error::InvalidArgument);
        }

        self.AQA.write(
            AQA::ASQS.val(submission_queue_size as u32 - 1)
                + AQA::ACQS.val(completion_queue_size as u32 - 1),
        );
        self.ASQ.set(submission_queue_pointer.into_u64());
        self.ACQ.set(completion_queue_pointer.into_u64());

        Ok(())
    }

    pub fn configure_controller(&self) {
        const IOSQES: u32 = size_of::<SubmissionQueueEntry>().ilog2();
        const IOCQES: u32 = size_of::<CompletionQueueEntry>().ilog2();

        self.CC.modify(
            CC::IOCQES.val(IOCQES)
                + CC::IOSQES.val(IOSQES)
                + CC::MPS.val(0)
                + CC::CSS::NvmCommandSet,
        );
    }

    pub fn enable_controller(&self, mut timeout_cycles: u64) -> Result<(), Error> {
        self.CC.modify(CC::ENABLE::SET);

        while timeout_cycles > 0 && !self.CSTS.matches_any(&[CSTS::RDY::SET, CSTS::CFS::SET]) {
            timeout_cycles -= 1;
            core::hint::spin_loop();
        }

        if timeout_cycles == 0 {
            return Err(Error::TimedOut);
        }

        if self.CSTS.matches_all(CSTS::CFS::SET) {
            log::error!("nvme: controller fatal status after enable");
            return Err(Error::InvalidArgument);
        }

        Ok(())
    }

    pub fn disable_controller(&self, mut timeout_cycles: u64) -> Result<(), Error> {
        self.CC.modify(CC::ENABLE::CLEAR);

        while timeout_cycles > 0 && self.CSTS.matches_all(CSTS::RDY::SET) {
            timeout_cycles -= 1;
            core::hint::spin_loop();
        }

        if timeout_cycles > 0 {
            Ok(())
        } else {
            Err(Error::TimedOut)
        }
    }
}
