use std::collections::BTreeMap;
use std::sync::atomic::{fence, Ordering};
use std::thread;
use std::time::{Duration, Instant};

use redox_driver_sys::dma::DmaBuffer;

use crate::driver::{DriverError, Result};

use super::transport::VirtioModernPciTransport;

const VIRTQ_DESC_F_NEXT: u16 = 1;
const VIRTQ_DESC_F_WRITE: u16 = 2;

#[repr(C)]
#[derive(Clone, Copy, Debug, Default)]
struct VirtqDesc {
    addr: u64,
    len: u32,
    flags: u16,
    next: u16,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Default)]
struct VirtqUsedElem {
    id: u32,
    len: u32,
}

pub struct Virtqueue {
    index: u16,
    size: u16,
    notify_off: u16,
    desc: DmaBuffer,
    avail: DmaBuffer,
    used: DmaBuffer,
    free_list: Vec<u16>,
    pending: BTreeMap<u16, Vec<u16>>,
    last_used_idx: u16,
    failed: bool,
}

impl Virtqueue {
    pub fn new(index: u16, size: u16, notify_off: u16) -> Result<Self> {
        let desc_bytes = usize::from(size) * core::mem::size_of::<VirtqDesc>();
        let avail_bytes = 6 + usize::from(size) * 2;
        let used_bytes = 6 + usize::from(size) * core::mem::size_of::<VirtqUsedElem>();

        let desc = DmaBuffer::allocate(desc_bytes, 16)
            .map_err(|e| DriverError::Buffer(format!("VirtIO desc allocation failed: {e}")))?;
        let avail = DmaBuffer::allocate(avail_bytes, 2)
            .map_err(|e| DriverError::Buffer(format!("VirtIO avail allocation failed: {e}")))?;
        let used = DmaBuffer::allocate(used_bytes, 4)
            .map_err(|e| DriverError::Buffer(format!("VirtIO used allocation failed: {e}")))?;

        Ok(Self {
            index,
            size,
            notify_off,
            desc,
            avail,
            used,
            free_list: (0..size).rev().collect(),
            pending: BTreeMap::new(),
            last_used_idx: 0,
            failed: false,
        })
    }

    pub fn index(&self) -> u16 {
        self.index
    }

    pub fn size(&self) -> u16 {
        self.size
    }

    pub fn notify_off(&self) -> u16 {
        self.notify_off
    }

    pub fn desc_addr(&self) -> u64 {
        self.desc.physical_address() as u64
    }

    pub fn avail_addr(&self) -> u64 {
        self.avail.physical_address() as u64
    }

    pub fn used_addr(&self) -> u64 {
        self.used.physical_address() as u64
    }

    pub fn submit_request(
        &mut self,
        transport: &VirtioModernPciTransport,
        request_addr: u64,
        request_len: u32,
        response_addr: u64,
        response_len: u32,
        timeout: Duration,
    ) -> Result<()> {
        if self.failed {
            return Err(DriverError::Io(format!(
                "VirtIO queue {} is failed and cannot accept more requests",
                self.index
            )));
        }

        let head = self.alloc_desc()?;
        let tail = self.alloc_desc()?;
        self.write_desc(
            head,
            VirtqDesc {
                addr: request_addr,
                len: request_len,
                flags: VIRTQ_DESC_F_NEXT,
                next: tail,
            },
        );
        self.write_desc(
            tail,
            VirtqDesc {
                addr: response_addr,
                len: response_len,
                flags: VIRTQ_DESC_F_WRITE,
                next: 0,
            },
        );

        self.pending.insert(head, vec![head, tail]);

        // VirtIO spec 2.6.8.1: The driver MUST perform a memory barrier before
        // updating the available index to ensure the descriptor table writes and
        // the available ring entry are visible to the device.
        fence(Ordering::SeqCst);

        self.push_avail(head);
        if let Err(error) = transport.notify_queue(self.index, self.notify_off) {
            self.reclaim_pending_chain(head);
            return Err(error);
        }

        if let Err(error) = self.wait_used(head, timeout) {
            self.failed = true;
            return Err(error);
        }

        Ok(())
    }

    fn alloc_desc(&mut self) -> Result<u16> {
        self.free_list.pop().ok_or_else(|| {
            DriverError::Buffer(format!(
                "VirtIO queue {} ran out of descriptors",
                self.index
            ))
        })
    }

    fn push_avail(&mut self, head: u16) {
        let avail_idx = self.read_avail_idx();
        let slot = usize::from(avail_idx % self.size);
        self.write_avail_ring(slot, head);
        fence(Ordering::Release);
        self.write_avail_idx(avail_idx.wrapping_add(1));
    }

    fn wait_used(&mut self, expected_head: u16, timeout: Duration) -> Result<()> {
        let deadline = Instant::now() + timeout;
        loop {
            let used_idx = self.read_used_idx();
            if used_idx != self.last_used_idx {
                let slot = usize::from(self.last_used_idx % self.size);
                let elem = self.read_used_elem(slot);
                self.last_used_idx = self.last_used_idx.wrapping_add(1);
                self.free_chain(elem.id as u16)?;
                if elem.id as u16 != expected_head {
                    self.failed = true;
                    return Err(DriverError::Io(format!(
                        "VirtIO queue {} completed descriptor head {} while waiting for {}",
                        self.index, elem.id, expected_head
                    )));
                }
                return Ok(());
            }

            if Instant::now() >= deadline {
                return Err(DriverError::Io(format!(
                    "VirtIO queue {} timed out waiting for descriptor head {}",
                    self.index, expected_head
                )));
            }

            thread::yield_now();
            thread::sleep(Duration::from_millis(1));
        }
    }

    fn free_chain(&mut self, head: u16) -> Result<()> {
        let chain = self.pending.remove(&head).ok_or_else(|| {
            DriverError::Io(format!(
                "VirtIO queue {} completed unknown descriptor head {}",
                self.index, head
            ))
        })?;
        for index in chain {
            self.free_list.push(index);
        }
        Ok(())
    }

    fn reclaim_pending_chain(&mut self, head: u16) {
        if let Some(chain) = self.pending.remove(&head) {
            for index in chain {
                self.free_list.push(index);
            }
        }
    }

    fn write_desc(&mut self, index: u16, desc: VirtqDesc) {
        let ptr = self.desc.as_mut_ptr() as *mut VirtqDesc;
        // SAFETY: The DMA buffer is sized for `size` descriptors, and callers only pass indices < size.
        unsafe { ptr.add(index as usize).write(desc) };
    }

    fn read_used_elem(&self, slot: usize) -> VirtqUsedElem {
        let offset = 4 + slot * core::mem::size_of::<VirtqUsedElem>();
        let ptr = self.used.as_ptr().wrapping_add(offset) as *const VirtqUsedElem;
        // SAFETY: The used ring allocation covers all ring entries; read_unaligned handles packed access.
        unsafe { ptr.read_unaligned() }
    }

    fn read_avail_idx(&self) -> u16 {
        let ptr = self.avail.as_ptr().wrapping_add(2) as *const u16;
        // SAFETY: Offset 2 is the avail.idx field in the DMA allocation.
        unsafe { ptr.read_unaligned() }
    }

    fn write_avail_idx(&mut self, value: u16) {
        let ptr = self.avail.as_mut_ptr().wrapping_add(2) as *mut u16;
        // SAFETY: Offset 2 is the avail.idx field in the DMA allocation.
        unsafe { ptr.write_unaligned(value) };
    }

    fn write_avail_ring(&mut self, slot: usize, value: u16) {
        let offset = 4 + slot * core::mem::size_of::<u16>();
        let ptr = self.avail.as_mut_ptr().wrapping_add(offset) as *mut u16;
        // SAFETY: The avail ring allocation is sized for `size` u16 entries.
        unsafe { ptr.write_unaligned(value) };
    }

    fn read_used_idx(&self) -> u16 {
        let ptr = self.used.as_ptr().wrapping_add(2) as *const u16;
        // SAFETY: Offset 2 is the used.idx field in the DMA allocation.
        unsafe { ptr.read_unaligned() }
    }
}