// SPDX-License-Identifier: MPL-2.0

use alloc::{boxed::Box, fmt::Debug, string::ToString, sync::Arc, vec::Vec};
use core::hint::spin_loop;

use aster_console::{AnyConsoleDevice, ConsoleCallback};
use aster_util::mem_obj_slice::Slice;
use log::debug;
use ostd::{
    arch::trap::TrapFrame,
    mm::{DmaDirection, DmaStream, FrameAllocOptions, VmReader, io_util::HasVmReaderWriter},
    sync::{Rcu, SpinLock},
};

use super::{DEVICE_NAME, config::VirtioConsoleConfig};
use crate::{
    device::{VirtioDeviceError, console::config::ConsoleFeatures},
    queue::VirtQueue,
    transport::{ConfigManager, VirtioTransport},
};

pub struct ConsoleDevice {
    config_manager: ConfigManager<VirtioConsoleConfig>,
    transport: SpinLock<Box<dyn VirtioTransport>>,
    receive_queue: SpinLock<VirtQueue>,
    transmit_queue: SpinLock<VirtQueue>,
    send_buffer: Arc<DmaStream>,
    receive_buffer: Arc<DmaStream>,
    #[expect(clippy::box_collection)]
    callbacks: Rcu<Box<Vec<&'static ConsoleCallback>>>,
}

impl AnyConsoleDevice for ConsoleDevice {
    fn send(&self, value: &[u8]) {
        let mut transmit_queue = self.transmit_queue.disable_irq().lock();
        let mut reader = VmReader::from(value);

        while reader.remain() > 0 {
            let mut writer = self.send_buffer.writer().unwrap();
            let len = writer.write(&mut reader);
            self.send_buffer.sync(0..len).unwrap();

            let slice = Slice::new(&self.send_buffer, 0..len);
            transmit_queue.add_dma_buf(&[&slice], &[]).unwrap();

            if transmit_queue.should_notify() {
                transmit_queue.notify();
            }
            while !transmit_queue.can_pop() {
                spin_loop();
            }
            transmit_queue.pop_used().unwrap();
        }
    }

    fn register_callback(&self, callback: &'static ConsoleCallback) {
        loop {
            let callbacks = self.callbacks.read();
            let mut callbacks_cloned = callbacks.get().clone();
            callbacks_cloned.push(callback);
            if callbacks.compare_exchange(callbacks_cloned).is_ok() {
                break;
            }
            // Contention on pushing, retry.
            core::hint::spin_loop();
        }
    }
}

impl Debug for ConsoleDevice {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("ConsoleDevice")
            .field("config", &self.config_manager.read_config())
            .field("transport", &self.transport)
            .field("receive_queue", &self.receive_queue)
            .field("transmit_queue", &self.transmit_queue)
            .finish()
    }
}

impl ConsoleDevice {
    pub fn negotiate_features(features: u64) -> u64 {
        let mut features = ConsoleFeatures::from_bits_truncate(features);
        // A virtio console device may have multiple ports, but we only use one port to communicate now.
        features.remove(ConsoleFeatures::VIRTIO_CONSOLE_F_MULTIPORT);
        features.bits()
    }

    pub fn init(mut transport: Box<dyn VirtioTransport>) -> Result<(), VirtioDeviceError> {
        let config_manager = VirtioConsoleConfig::new_manager(transport.as_ref());
        debug!("virtio_console_config = {:?}", config_manager.read_config());

        const RECV0_QUEUE_INDEX: u16 = 0;
        const TRANSMIT0_QUEUE_INDEX: u16 = 1;
        let receive_queue =
            SpinLock::new(VirtQueue::new(RECV0_QUEUE_INDEX, 2, transport.as_mut()).unwrap());
        let transmit_queue =
            SpinLock::new(VirtQueue::new(TRANSMIT0_QUEUE_INDEX, 2, transport.as_mut()).unwrap());

        let send_buffer = {
            let segment = FrameAllocOptions::new().alloc_segment(1).unwrap();
            Arc::new(DmaStream::map(segment.into(), DmaDirection::ToDevice, false).unwrap())
        };

        let receive_buffer = {
            let segment = FrameAllocOptions::new().alloc_segment(1).unwrap();
            Arc::new(DmaStream::map(segment.into(), DmaDirection::FromDevice, false).unwrap())
        };

        let device = Arc::new(Self {
            config_manager,
            transport: SpinLock::new(transport),
            receive_queue,
            transmit_queue,
            send_buffer,
            receive_buffer,
            callbacks: Rcu::new(Box::new(Vec::new())),
        });

        device.activate_receive_buffer(&mut device.receive_queue.disable_irq().lock());

        // Register irq callbacks
        let mut transport = device.transport.disable_irq().lock();
        let handle_console_input = {
            let device = device.clone();
            move |_: &TrapFrame| device.handle_recv_irq()
        };
        transport
            .register_queue_callback(RECV0_QUEUE_INDEX, Box::new(handle_console_input), false)
            .unwrap();
        transport
            .register_cfg_callback(Box::new(config_space_change))
            .unwrap();
        transport.finish_init();
        drop(transport);

        aster_console::register_device(DEVICE_NAME.to_string(), device);

        Ok(())
    }

    fn handle_recv_irq(&self) {
        let mut receive_queue = self.receive_queue.disable_irq().lock();

        let Ok((_, len)) = receive_queue.pop_used() else {
            return;
        };
        self.receive_buffer.sync(0..len as usize).unwrap();

        let callbacks = self.callbacks.read();
        for callback in callbacks.get().iter() {
            let mut reader = self.receive_buffer.reader().unwrap();
            reader.limit(len as usize);
            callback(reader);
        }
        drop(callbacks);

        self.activate_receive_buffer(&mut receive_queue);
    }

    fn activate_receive_buffer(&self, receive_queue: &mut VirtQueue) {
        receive_queue
            // We limit the buffer length to one to work around a QEMU bug that causes incorrect
            // results when pasting more than 32 bytes into the virtio console. This has no
            // performance penalty, since QEMU always gets one byte at a time, regardless of
            // whether we have this limit or not.
            //
            // For the QEMU bug, see details at
            // <https://lore.kernel.org/qemu-devel/20240707111940.232549-3-lrh2000@pku.edu.cn/T/#u>.
            .add_dma_buf(&[], &[&Slice::new(&self.receive_buffer, 0..1)])
            .unwrap();

        if receive_queue.should_notify() {
            receive_queue.notify();
        }
    }
}

fn config_space_change(_: &TrapFrame) {
    debug!("Virtio-Console device configuration space change");
}
