// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

use crate::queue::QueueCoreCompleteWork;
use crate::queue::QueueCoreGetWork;
use crate::queue::QueueError;
use crate::queue::QueueParams;
use crate::queue::QueueWork;
use crate::queue::VirtioQueuePayload;
use crate::queue::new_queue;
use crate::spec::VirtioDeviceFeatures;
use futures::FutureExt;
use futures::Stream;
use guestmem::DoorbellRegistration;
use guestmem::GuestMemory;
use guestmem::GuestMemoryError;
use guestmem::MappedMemoryRegion;
use inspect::Inspect;
use pal_async::wait::PolledWait;
use pal_event::Event;
use parking_lot::Mutex;
use std::io::Error;
use std::pin::Pin;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use std::task::ready;
use thiserror::Error;
use vmcore::interrupt::Interrupt;

#[derive(Debug)]
pub(crate) struct VirtioQueueUsedHandler {
    core: QueueCoreCompleteWork,
    notify_guest: Interrupt,
}

impl VirtioQueueUsedHandler {
    pub(crate) fn new(core: QueueCoreCompleteWork, notify_guest: Interrupt) -> Self {
        Self { core, notify_guest }
    }

    pub(crate) fn complete_descriptor(&mut self, work: &QueueWork, bytes_written: u32) {
        match self.core.complete_descriptor(work, bytes_written) {
            Ok(true) => {
                self.notify_guest.deliver();
            }
            Ok(false) => {}
            Err(err) => {
                tracelimit::error_ratelimited!(
                    error = &err as &dyn std::error::Error,
                    "failed to complete descriptor"
                );
            }
        }
    }
}

pub struct VirtioQueueCallbackWork {
    used_queue_handler: Arc<Mutex<VirtioQueueUsedHandler>>,
    work: QueueWork,
    pub payload: Vec<VirtioQueuePayload>,
    completed: bool,
}

impl VirtioQueueCallbackWork {
    pub(crate) fn new(
        mut work: QueueWork,
        used_queue_handler: &Arc<Mutex<VirtioQueueUsedHandler>>,
    ) -> Self {
        let used_queue_handler = used_queue_handler.clone();
        let payload = std::mem::take(&mut work.payload);
        Self {
            work,
            payload,
            used_queue_handler,
            completed: false,
        }
    }

    pub fn complete(&mut self, bytes_written: u32) {
        assert!(!self.completed);
        self.used_queue_handler
            .lock()
            .complete_descriptor(&self.work, bytes_written);
        self.completed = true;
    }

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

    // Determine the total size of all readable or all writeable payload buffers.
    pub fn get_payload_length(&self, writeable: bool) -> u64 {
        self.payload
            .iter()
            .filter(|x| x.writeable == writeable)
            .fold(0, |acc, x| acc + x.length as u64)
    }

    // Read all payload into a buffer.
    pub fn read(&self, mem: &GuestMemory, target: &mut [u8]) -> Result<usize, GuestMemoryError> {
        let mut remaining = target;
        let mut read_bytes: usize = 0;
        for payload in &self.payload {
            if payload.writeable {
                continue;
            }

            let size = std::cmp::min(payload.length as usize, remaining.len());
            let (current, next) = remaining.split_at_mut(size);
            mem.read_at(payload.address, current)?;
            read_bytes += size;
            if next.is_empty() {
                break;
            }

            remaining = next;
        }

        Ok(read_bytes)
    }

    // Write the specified buffer to the payload buffers.
    pub fn write_at_offset(
        &self,
        offset: u64,
        mem: &GuestMemory,
        source: &[u8],
    ) -> Result<(), VirtioWriteError> {
        let mut skip_bytes = offset;
        let mut remaining = source;
        for payload in &self.payload {
            if !payload.writeable {
                continue;
            }

            let payload_length = payload.length as u64;
            if skip_bytes >= payload_length {
                skip_bytes -= payload_length;
                continue;
            }

            let size = std::cmp::min(
                payload_length as usize - skip_bytes as usize,
                remaining.len(),
            );
            let (current, next) = remaining.split_at(size);
            mem.write_at(payload.address + skip_bytes, current)?;
            remaining = next;
            if remaining.is_empty() {
                break;
            }
            skip_bytes = 0;
        }

        if !remaining.is_empty() {
            return Err(VirtioWriteError::NotAllWritten(source.len()));
        }

        Ok(())
    }

    pub fn write(&self, mem: &GuestMemory, source: &[u8]) -> Result<(), VirtioWriteError> {
        self.write_at_offset(0, mem, source)
    }
}

#[derive(Debug, Error)]
pub enum VirtioWriteError {
    #[error(transparent)]
    Memory(#[from] GuestMemoryError),
    #[error("{0:#x} bytes not written")]
    NotAllWritten(usize),
}

impl Drop for VirtioQueueCallbackWork {
    fn drop(&mut self) {
        if !self.completed {
            self.complete(0);
        }
    }
}

#[derive(Debug, Inspect)]
pub struct VirtioQueue {
    #[inspect(flatten)]
    core: QueueCoreGetWork,
    #[inspect(skip)]
    used_handler: Arc<Mutex<VirtioQueueUsedHandler>>,
    #[inspect(skip)]
    queue_event: PolledWait<Event>,
}

impl VirtioQueue {
    pub fn new(
        features: VirtioDeviceFeatures,
        params: QueueParams,
        mem: GuestMemory,
        notify: Interrupt,
        queue_event: PolledWait<Event>,
    ) -> Result<Self, QueueError> {
        let (get_work, complete_work) = new_queue(features, mem, params)?;
        let used_handler = Arc::new(Mutex::new(VirtioQueueUsedHandler::new(
            complete_work,
            notify,
        )));
        Ok(Self {
            core: get_work,
            used_handler,
            queue_event,
        })
    }

    /// Polls until the queue is kicked by the guest, indicating new work may be available.
    pub fn poll_kick(&mut self, cx: &mut Context<'_>) -> Poll<()> {
        ready!(self.queue_event.wait().poll_unpin(cx)).expect("waits on Event cannot fail");
        Poll::Ready(())
    }

    /// Try to get the next work item from the queue. Returns `Ok(None)` if no
    /// work is currently available, or an error if there was an issue accessing
    /// the queue.
    ///
    /// If `None` is returned, then the queue will be armed so that the guest
    /// will kick it when new work is available; the caller can use
    /// [`poll_kick`](Self::poll_kick) to wait for this.
    pub fn try_next(&mut self) -> Result<Option<VirtioQueueCallbackWork>, Error> {
        Ok(self
            .core
            .try_next_work()
            .map_err(Error::other)?
            .map(|work| VirtioQueueCallbackWork::new(work, &self.used_handler)))
    }

    fn poll_next_buffer(
        &mut self,
        cx: &mut Context<'_>,
    ) -> Poll<Result<Option<VirtioQueueCallbackWork>, Error>> {
        loop {
            if let Some(work) = self.try_next()? {
                return Ok(Some(work)).into();
            }
            ready!(self.poll_kick(cx));
        }
    }
}

impl Drop for VirtioQueue {
    fn drop(&mut self) {
        if Arc::get_mut(&mut self.used_handler).is_none() {
            tracing::error!("Virtio queue dropped with outstanding work pending")
        }
    }
}

impl Stream for VirtioQueue {
    type Item = Result<VirtioQueueCallbackWork, Error>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        ready!(self.get_mut().poll_next_buffer(cx))
            .transpose()
            .into()
    }
}

pub(crate) struct VirtioDoorbells {
    registration: Option<Arc<dyn DoorbellRegistration>>,
    doorbells: Vec<Box<dyn Send + Sync>>,
}

impl VirtioDoorbells {
    pub fn new(registration: Option<Arc<dyn DoorbellRegistration>>) -> Self {
        Self {
            registration,
            doorbells: Vec::new(),
        }
    }

    pub fn add(&mut self, address: u64, value: Option<u64>, length: Option<u32>, event: &Event) {
        if let Some(registration) = &mut self.registration {
            let doorbell = registration.register_doorbell(address, value, length, event);
            if let Ok(doorbell) = doorbell {
                self.doorbells.push(doorbell);
            }
        }
    }

    pub fn clear(&mut self) {
        self.doorbells.clear();
    }
}

#[derive(Copy, Clone, Debug, Default)]
pub struct DeviceTraitsSharedMemory {
    pub id: u8,
    pub size: u64,
}

#[derive(Clone, Debug, Default)]
pub struct DeviceTraits {
    pub device_id: u16,
    pub device_features: VirtioDeviceFeatures,
    pub max_queues: u16,
    pub device_register_length: u32,
    pub shared_memory: DeviceTraitsSharedMemory,
}

pub trait VirtioDevice: inspect::InspectMut + Send {
    fn traits(&self) -> DeviceTraits;
    fn read_registers_u32(&mut self, offset: u16) -> u32;
    fn write_registers_u32(&mut self, offset: u16, val: u32);
    /// Enable the device with the given resources.
    ///
    /// Called when the guest sets `DRIVER_OK`. On success, the device should
    /// start processing queues and the transport will reflect `DRIVER_OK` in
    /// the device status. On failure, the transport will log the error and
    /// leave `DRIVER_OK` unset, so the device remains inert and the guest
    /// will observe failures through IO timeouts.
    fn enable(&mut self, resources: Resources) -> anyhow::Result<()>;
    /// Poll the device to complete a disable/reset operation.
    ///
    /// This is called when the guest writes status=0 (device reset). The device
    /// should stop workers and drain any in-flight IO. Returns `Poll::Ready(())`
    /// when the disable is complete, or `Poll::Pending` if more work is needed.
    ///
    /// Devices that don't need async cleanup can return `Poll::Ready(())`
    /// immediately.
    fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()>;
}

pub struct QueueResources {
    pub params: QueueParams,
    pub notify: Interrupt,
    pub event: Event,
}

pub struct Resources {
    pub features: VirtioDeviceFeatures,
    pub queues: Vec<QueueResources>,
    pub shared_memory_region: Option<Arc<dyn MappedMemoryRegion>>,
    pub shared_memory_size: u64,
}
