use core::{future::poll_fn, mem::size_of, ptr::null_mut, task::Poll};

use alloc::collections::{BTreeMap, BTreeSet};
use bytemuck::{Pod, Zeroable};
use device_api::dma::DmaAllocator;
use kernel_arch::{Architecture, ArchitectureImpl};
use libk::dma::{BusAddress, DmaBuffer};
use libk_mm::address::AsPhysicalAddress;
use libk_util::{sync::IrqSafeSpinlock, waker::QueueWaker};
use static_assertions::const_assert;
use yggdrasil_abi::error::Error;

use super::{
    command::{Command, Request},
    error::NvmeError,
};

#[derive(Zeroable, Pod, Clone, Copy, Debug)]
#[repr(C)]
pub struct PhysicalRegionPage(u64);

// Bits:
//
//  16..32 - CID. Command identifier
//  14..16 - PSDT. PRP or SGL for data transfer.
//          0b00 - PRP used
//          0b01 - SGL used. Not implemented
//          0b10 - SGL used. Not implemented
//          0b11 - Reserved
//  10..14 - Reserved
//  08..10 - FUSE. Fused Operation
//  00..08 - OPC. Opcode
#[derive(Zeroable, Pod, Clone, Copy, Debug)]
#[repr(C)]
pub struct CommandDword0(u32);

#[derive(Zeroable, Pod, Clone, Copy, Debug)]
#[repr(C)]
pub struct SubmissionQueueEntry {
    pub command: CommandDword0,                // 0
    pub nsid: u32,                             // 1
    pub io_data: [u32; 2],                     // 2, 3
    pub metadata_pointer: u64,                 // 4, 5
    pub data_pointer: [PhysicalRegionPage; 2], // 6, 7, 8, 9
    pub command_specific: [u32; 6],            // 10, 11, 12, 13, 14, 15
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub struct CommandError {
    sct: u8,
    sc: u8,
}

#[derive(Zeroable, Pod, Clone, Copy, Debug)]
#[repr(C)]
pub struct CompletionQueueEntry {
    dw: [u32; 4],
}

pub struct Queue<T> {
    data: DmaBuffer<[T]>,
    mask: usize,
    head: usize,
    tail: usize,
    phase: bool,

    head_doorbell: *mut u32,
    tail_doorbell: *mut u32,
}

struct Inner {
    sq: Queue<SubmissionQueueEntry>,
    cq: Queue<CompletionQueueEntry>,

    completed: BTreeMap<u32, CompletionQueueEntry>,
    pending: BTreeSet<u32>,
}

pub struct QueuePair {
    id: u32,

    #[allow(unused)]
    vector: usize,

    sq_base: BusAddress,
    cq_base: BusAddress,

    pub completion_notify: QueueWaker,

    inner: IrqSafeSpinlock<Inner>,
}

pub struct PrpList {
    prp1: PhysicalRegionPage,
    prp2: PhysicalRegionPage,
    #[allow(unused)]
    list: Option<DmaBuffer<[BusAddress]>>,
}

impl PrpList {
    pub const fn empty() -> Self {
        Self {
            prp1: PhysicalRegionPage::null(),
            prp2: PhysicalRegionPage::null(),
            list: None,
        }
    }

    pub fn from_buffer(
        dma: &dyn DmaAllocator,
        base: BusAddress,
        size: usize,
    ) -> Result<Self, NvmeError> {
        // TODO hardcoded page size
        if base.into_u64() % 0x1000 != 0 {
            todo!();
        }

        match size {
            0 => Ok(Self::empty()),
            _ if size <= 0x1000 => Ok(Self {
                prp1: PhysicalRegionPage::with_addr(base),
                prp2: PhysicalRegionPage::null(),
                list: None,
            }),
            _ if size <= 0x2000 => Ok(Self {
                prp1: PhysicalRegionPage::with_addr(base),
                prp2: PhysicalRegionPage::with_addr(base.add(0x1000)),
                list: None,
            }),
            _ => {
                let count = size.div_ceil(0x1000);
                let list =
                    DmaBuffer::new_slice_with(dma, |i| base.add((i + 1) * 0x1000), count - 1)
                        .map_err(NvmeError::MemoryError)?;

                Ok(Self {
                    prp1: PhysicalRegionPage::with_addr(base),
                    prp2: PhysicalRegionPage::with_addr(list.bus_address()),
                    list: Some(list),
                })
            }
        }
    }
}

const_assert!(size_of::<CompletionQueueEntry>().is_power_of_two());

impl PhysicalRegionPage {
    pub const fn null() -> Self {
        Self(0)
    }

    pub const fn with_addr(address: BusAddress) -> Self {
        Self(address.into_u64())
    }
}

impl CommandDword0 {
    pub fn set_command_id(&mut self, id: u32) {
        debug_assert!(id & 0xFFFF0000 == 0);
        self.0 &= !(0xFFFF << 16);
        self.0 |= id << 16;
    }

    pub fn set_opcode(&mut self, opcode: u8) {
        self.0 &= !0xFF;
        self.0 |= opcode as u32;
    }
}

impl CompletionQueueEntry {
    pub fn phase(&self) -> bool {
        self.dw[3] & (1 << 16) != 0
    }

    pub fn sub_queue_id(&self) -> u32 {
        self.dw[2] >> 16
    }

    pub fn sub_queue_head(&self) -> usize {
        (self.dw[2] & 0xFFFF) as _
    }

    pub fn command_id(&self) -> u32 {
        self.dw[3] & 0xFFFF
    }

    pub fn error(&self) -> Option<CommandError> {
        let status = (self.dw[3] >> 17) as u16;

        if status != 0 {
            Some(CommandError {
                sct: ((status >> 8) & 0x7) as u8,
                sc: status as u8,
            })
        } else {
            None
        }
    }
}

impl<T> Queue<T> {
    pub fn new(
        data: DmaBuffer<[T]>,
        head_doorbell: *mut u32,
        tail_doorbell: *mut u32,
        phase: bool,
    ) -> Self {
        assert!(
            (head_doorbell.is_null() && !tail_doorbell.is_null())
                || (!head_doorbell.is_null() && tail_doorbell.is_null())
        );

        Self {
            mask: data.len() - 1,
            head: 0,
            tail: 0,
            data,
            head_doorbell,
            tail_doorbell,
            phase,
        }
    }

    pub fn enqueue(&mut self, item: T) -> usize {
        let index = self.tail;
        self.data[self.tail] = item;
        self.phase ^= self.set_tail(self.next_index(self.tail));
        index
    }

    pub fn at_head(&self, offset: usize) -> (&T, bool) {
        let index = (self.head + offset) & self.mask;
        let expected_phase = self.phase ^ (index < self.head);
        (&self.data[index], expected_phase)
    }

    pub fn take(&mut self, count: usize) {
        let index = (self.head + count) & self.mask;
        self.phase ^= self.set_head(index);
    }

    pub fn take_until(&mut self, new_head: usize) {
        self.phase ^= self.set_head(new_head);
    }

    fn next_index(&self, index: usize) -> usize {
        (index + 1) & self.mask
    }

    fn set_tail(&mut self, new_tail: usize) -> bool {
        let wrapped = new_tail < self.tail;

        self.tail = new_tail;

        if !self.tail_doorbell.is_null() {
            self.data.cache_flush_element(self.tail, true);
            unsafe {
                self.tail_doorbell
                    .write_volatile(self.tail.try_into().unwrap());
            }
            ArchitectureImpl::memory_barrier();
        }

        wrapped
    }

    fn set_head(&mut self, new_head: usize) -> bool {
        let wrapped = new_head < self.head;

        self.head = new_head;

        if !self.head_doorbell.is_null() {
            unsafe {
                self.head_doorbell
                    .write_volatile(self.head.try_into().unwrap());
            }
        }

        wrapped
    }
}

impl QueuePair {
    pub fn new(
        dma: &dyn DmaAllocator,
        id: u32,
        vector: usize,
        capacity: usize,
        sq_doorbell: *mut u32,
        cq_doorbell: *mut u32,
    ) -> Result<Self, Error> {
        let sq_data = DmaBuffer::new_slice(dma, SubmissionQueueEntry::zeroed(), capacity)?;
        let cq_data = DmaBuffer::new_slice(dma, CompletionQueueEntry::zeroed(), capacity)?;

        let sq_base = sq_data.bus_address();
        let cq_base = cq_data.bus_address();

        log::debug!("Allocated queue pair: sq={:p}, cq={:p}", sq_data, cq_data);

        let sq = Queue::new(sq_data, null_mut(), sq_doorbell, true);
        let cq = Queue::new(cq_data, cq_doorbell, null_mut(), true);

        let inner = IrqSafeSpinlock::new(Inner {
            sq,
            cq,
            pending: BTreeSet::new(),
            completed: BTreeMap::new(),
        });

        Ok(Self {
            completion_notify: QueueWaker::new(),

            id,
            vector,
            sq_base,
            cq_base,
            inner,
        })
    }

    #[inline]
    pub fn sq_bus_pointer(&self) -> BusAddress {
        self.sq_base
    }

    #[inline]
    pub fn cq_bus_pointer(&self) -> BusAddress {
        self.cq_base
    }

    pub async fn wait_for_completion<T: Unpin>(
        &self,
        command_id: u32,
        result: T,
    ) -> Result<T, CommandError> {
        let mut response = Some(result);
        poll_fn(|cx| {
            let mut inner = self.inner.lock();

            if let Some(entry) = inner.completed.remove(&command_id) {
                self.completion_notify.remove(cx.waker());

                let result = if let Some(error) = entry.error() {
                    Err(error)
                } else {
                    Ok(response.take().unwrap())
                };

                Poll::Ready(result)
            } else {
                self.completion_notify.register(cx.waker());
                Poll::Pending
            }
        })
        .await
    }

    pub fn submit<C: Command>(
        &self,
        cmd: C,
        ranges: &PrpList,
        set_pending: bool,
    ) -> Result<u32, NvmeError> {
        let mut inner = self.inner.lock();
        let mut sqe = SubmissionQueueEntry::zeroed();

        sqe.data_pointer[0] = ranges.prp1;
        sqe.data_pointer[1] = ranges.prp2;

        cmd.fill_sqe(&mut sqe);

        let command_id = inner.sq.tail.try_into().unwrap();
        sqe.command.set_command_id(command_id);

        if set_pending {
            inner.pending.insert(command_id);
        }

        inner.sq.enqueue(sqe);

        Ok(command_id)
    }

    pub async fn request_no_data<C: Command>(&self, req: C) -> Result<(), NvmeError> {
        let list = PrpList::empty();
        let command_id = self.submit(req, &list, true)?;
        self.wait_for_completion(command_id, ())
            .await
            .map_err(NvmeError::CommandError)
    }

    pub async fn request<'r, R: Request>(
        &'r self,
        dma: &dyn DmaAllocator,
        req: R,
    ) -> Result<DmaBuffer<R::Response>, NvmeError>
    where
        R::Response: 'r,
    {
        let response = DmaBuffer::new_uninit(dma).map_err(NvmeError::MemoryError)?;
        let list = PrpList::from_buffer(dma, response.bus_address(), size_of::<R>())?;
        let command_id = self.submit(req, &list, true)?;
        let result = self.wait_for_completion(command_id, response).await?;
        Ok(unsafe { DmaBuffer::assume_init(result) })
    }

    pub fn process_completions(&self) -> usize {
        let mut inner = self.inner.lock();
        let mut n = 0;

        loop {
            let (cmp, expected_phase) = inner.cq.at_head(n);
            let cmp_phase = cmp.phase();

            if cmp_phase != expected_phase {
                break;
            }

            n += 1;

            let command_id = cmp.command_id();
            let sub_queue_id = cmp.sub_queue_id();
            // TODO allow several sqs receive completions through one cq?
            debug_assert_eq!(sub_queue_id, self.id);

            let sub_queue_head = cmp.sub_queue_head();
            let cmp = *cmp;
            inner.sq.take_until(sub_queue_head);

            if inner.pending.remove(&command_id) {
                inner.completed.insert(command_id, cmp);
            }
        }

        if n != 0 {
            inner.cq.take(n);
        }

        if n != 0 {
            self.completion_notify.wake_all();
        }

        n
    }
}

// TODO
unsafe impl Sync for QueuePair {}
unsafe impl Send for QueuePair {}
