// Copyright 2024 The ChromiumOS Authors

// Use of this source code is governed by a BSD-style license that can be

// found in the LICENSE file.

//! Virtio console device worker thread.

use std::collections::BTreeMap;

use std::collections::VecDeque;

use std::sync::mpsc;

use std::sync::Arc;

use anyhow::anyhow;

use anyhow::Context;

use base::error;

use base::Event;

use base::EventToken;

use base::WaitContext;

use base::WorkerThread;

use sync::Mutex;

use crate::virtio::console::control::process_control_receive_queue;

use crate::virtio::console::control::process_control_transmit_queue;

use crate::virtio::console::control::ControlMsgBytes;

use crate::virtio::console::input::process_receive_queue;

use crate::virtio::console::output::process_transmit_queue;

use crate::virtio::console::port::ConsolePort;

use crate::virtio::console::port::ConsolePortInfo;

use crate::virtio::Queue;

const PORT0_RECEIVEQ_IDX: usize = 0;

const PORT0_TRANSMITQ_IDX: usize = 1;

const CONTROL_RECEIVEQ_IDX: usize = 2;

const CONTROL_TRANSMITQ_IDX: usize = 3;

const PORT1_RECEIVEQ_IDX: usize = 4;

const PORT1_TRANSMITQ_IDX: usize = 5;

pub struct WorkerPort {

    info: Option<ConsolePortInfo>,

    in_avail_evt: Event,

    input_buffer: Arc<Mutex<VecDeque<u8>>>,

    output: Box<dyn std::io::Write + Send>,

}

impl WorkerPort {

    pub fn from_console_port(port: &mut ConsolePort) -> WorkerPort {

        let in_avail_evt = port.clone_in_avail_evt().unwrap();

        let input_buffer = port.clone_input_buffer();

        let output = port

            .take_output()

            .unwrap_or_else(|| Box::new(std::io::sink()));

        let info = port.port_info().cloned();

        WorkerPort {

            info,

            in_avail_evt,

            input_buffer,

            output,

        }

    }

    /// Restore the state retrieved from `ConsolePort` by `WorkerPort::from_console_port()`.

    pub fn into_console_port(self, console_port: &mut ConsolePort) {

        console_port.restore_output(self.output);

    }

    pub fn is_console(&self) -> bool {

        self.info

            .as_ref()

            .map(|info| info.console)

            .unwrap_or_default()

    }

    pub fn name(&self) -> Option<&str> {

        self.info.as_ref().and_then(ConsolePortInfo::name)

    }

}

#[derive(EventToken)]

enum Token {

    ReceiveQueueAvailable(u32),

    TransmitQueueAvailable(u32),

    InputAvailable(u32),

    ControlReceiveQueueAvailable,

    ControlTransmitQueueAvailable,

    WorkerRequest,

    Kill,

}

pub enum WorkerRequest {

    StartQueue {

        idx: usize,

        queue: Queue,

        response_sender: mpsc::SyncSender<anyhow::Result<()>>,

    },

    StopQueue {

        idx: usize,

        response_sender: mpsc::SyncSender<Option<Queue>>,

    },

}

pub struct Worker {

    wait_ctx: WaitContext<Token>,

    // Currently running queues.

    queues: BTreeMap<usize, Queue>,

    // Console ports indexed by port ID. At least port 0 will exist, and other ports may be

    // available if `VIRTIO_CONSOLE_F_MULTIPORT` is enabled.

    ports: Vec<WorkerPort>,

    // Device-to-driver messages to be received by the driver via the control receiveq.

    pending_receive_control_msgs: VecDeque<ControlMsgBytes>,

    worker_receiver: mpsc::Receiver<WorkerRequest>,

    worker_event: Event,

}

impl Worker {

    pub fn new(

        ports: Vec<WorkerPort>,

        worker_receiver: mpsc::Receiver<WorkerRequest>,

        worker_event: Event,

    ) -> anyhow::Result<Self> {

        let wait_ctx = WaitContext::new().context("WaitContext::new() failed")?;

        wait_ctx.add(&worker_event, Token::WorkerRequest)?;

        for (index, port) in ports.iter().enumerate() {

            let port_id = index as u32;

            wait_ctx.add(&port.in_avail_evt, Token::InputAvailable(port_id))?;

        }

        Ok(Worker {

            wait_ctx,

            queues: BTreeMap::new(),

            ports,

            pending_receive_control_msgs: VecDeque::new(),

            worker_receiver,

            worker_event,

        })

    }

    pub fn run(&mut self, kill_evt: &Event) -> anyhow::Result<()> {

        self.wait_ctx.add(kill_evt, Token::Kill)?;

        let res = self.run_loop();

        self.wait_ctx.delete(kill_evt)?;

        res

    }

    fn run_loop(&mut self) -> anyhow::Result<()> {

        let mut running = true;

        while running {

            let events = self.wait_ctx.wait()?;

            for event in events.iter().filter(|e| e.is_readable) {

                match event.token {

                    Token::TransmitQueueAvailable(port_id) => {

                        if let (Some(port), Some(transmitq)) = (

                            self.ports.get_mut(port_id as usize),

                            transmitq_idx(port_id).and_then(|idx| self.queues.get_mut(&idx)),

                        ) {

                            transmitq

                                .event()

                                .wait()

                                .context("failed reading transmit queue Event")?;

                            process_transmit_queue(transmitq, &mut port.output);

                        }

                    }

                    Token::ReceiveQueueAvailable(port_id) | Token::InputAvailable(port_id) => {

                        let port = self.ports.get_mut(port_id as usize);

                        let receiveq =

                            receiveq_idx(port_id).and_then(|idx| self.queues.get_mut(&idx));

                        let event = if matches!(event.token, Token::ReceiveQueueAvailable(..)) {

                            receiveq.as_ref().map(|q| q.event())

                        } else {

                            port.as_ref().map(|p| &p.in_avail_evt)

                        };

                        if let Some(event) = event {

                            event.wait().context("failed to clear receive event")?;

                        }

                        if let (Some(port), Some(receiveq)) = (port, receiveq) {

                            let mut input_buffer = port.input_buffer.lock();

                            process_receive_queue(&mut input_buffer, receiveq);

                        }

                    }

                    Token::ControlReceiveQueueAvailable => {

                        if let Some(ctrl_receiveq) = self.queues.get_mut(&CONTROL_RECEIVEQ_IDX) {

                            ctrl_receiveq

                                .event()

                                .wait()

                                .context("failed waiting on control event")?;

                            process_control_receive_queue(

                                ctrl_receiveq,

                                &mut self.pending_receive_control_msgs,

                            );

                        }

                    }

                    Token::ControlTransmitQueueAvailable => {

                        if let Some(ctrl_transmitq) = self.queues.get_mut(&CONTROL_TRANSMITQ_IDX) {

                            ctrl_transmitq

                                .event()

                                .wait()

                                .context("failed waiting on control event")?;

                            process_control_transmit_queue(

                                ctrl_transmitq,

                                &self.ports,

                                &mut self.pending_receive_control_msgs,

                            );

                        }

                        // Attempt to send any new replies if there is space in the receiveq.

                        if let Some(ctrl_receiveq) = self.queues.get_mut(&CONTROL_RECEIVEQ_IDX) {

                            process_control_receive_queue(

                                ctrl_receiveq,

                                &mut self.pending_receive_control_msgs,

                            )

                        }

                    }

                    Token::WorkerRequest => {

                        self.worker_event.wait()?;

                        self.process_worker_requests();

                    }

                    Token::Kill => running = false,

                }

            }

        }

        Ok(())

    }

    fn process_worker_requests(&mut self) {

        while let Ok(request) = self.worker_receiver.try_recv() {

            match request {

                WorkerRequest::StartQueue {

                    idx,

                    queue,

                    response_sender,

                } => {

                    let res = self.start_queue(idx, queue);

                    let _ = response_sender.send(res);

                }

                WorkerRequest::StopQueue {

                    idx,

                    response_sender,

                } => {

                    let res = self.stop_queue(idx);

                    let _ = response_sender.send(res);

                }

            }

        }

    }

    fn start_queue(&mut self, idx: usize, queue: Queue) -> anyhow::Result<()> {

        if let Some(port_id) = receiveq_port_id(idx) {

            self.wait_ctx

                .add(queue.event(), Token::ReceiveQueueAvailable(port_id))?;

        } else if let Some(port_id) = transmitq_port_id(idx) {

            self.wait_ctx

                .add(queue.event(), Token::TransmitQueueAvailable(port_id))?;

        } else if idx == CONTROL_RECEIVEQ_IDX {

            self.wait_ctx

                .add(queue.event(), Token::ControlReceiveQueueAvailable)?;

        } else if idx == CONTROL_TRANSMITQ_IDX {

            self.wait_ctx

                .add(queue.event(), Token::ControlTransmitQueueAvailable)?;

        } else {

            return Err(anyhow!("unhandled queue idx {idx}"));

        }

        let prev = self.queues.insert(idx, queue);

        assert!(prev.is_none());

        Ok(())

    }

    fn stop_queue(&mut self, idx: usize) -> Option<Queue> {

        if let Some(queue) = self.queues.remove(&idx) {

            let _ = self.wait_ctx.delete(queue.event());

            Some(queue)

        } else {

            None

        }

    }

}

pub struct WorkerHandle {

    worker_thread: WorkerThread<Vec<WorkerPort>>,

    worker_sender: mpsc::Sender<WorkerRequest>,

    worker_event: Event,

}

impl WorkerHandle {

    pub fn new(ports: Vec<WorkerPort>) -> anyhow::Result<Self> {

        let worker_event = Event::new().context("Event::new")?;

        let worker_event_clone = worker_event.try_clone().context("Event::try_clone")?;

        let (worker_sender, worker_receiver) = mpsc::channel();

        let worker_thread = WorkerThread::start("v_console", move |kill_evt| {

            let mut worker = Worker::new(ports, worker_receiver, worker_event_clone)

                .expect("console Worker::new() failed");

            if let Err(e) = worker.run(&kill_evt) {

                error!("console worker failed: {:#}", e);

            }

            worker.ports

        });

        Ok(WorkerHandle {

            worker_thread,

            worker_sender,

            worker_event,

        })

    }

    pub fn start_queue(&mut self, idx: usize, queue: Queue) -> anyhow::Result<()> {

        let (response_sender, response_receiver) = mpsc::sync_channel(0);

        self.worker_sender

            .send(WorkerRequest::StartQueue {

                idx,

                queue,

                response_sender,

            })

            .context("mpsc::Sender::send")?;

        self.worker_event.signal().context("Event::signal")?;

        response_receiver.recv().context("mpsc::Receiver::recv")?

    }

    pub fn stop_queue(&mut self, idx: usize) -> anyhow::Result<Option<Queue>> {

        let (response_sender, response_receiver) = mpsc::sync_channel(0);

        self.worker_sender

            .send(WorkerRequest::StopQueue {

                idx,

                response_sender,

            })

            .context("mpsc::Sender::send")?;

        self.worker_event.signal().context("Event::signal")?;

        response_receiver.recv().context("mpsc::Receiver::recv")

    }

    pub fn stop(self) -> Vec<WorkerPort> {

        self.worker_thread.stop()

    }

}

fn receiveq_idx(port_id: u32) -> Option<usize> {

    if port_id == 0 {

        Some(PORT0_RECEIVEQ_IDX)

    } else {

        PORT1_RECEIVEQ_IDX.checked_add((port_id - 1).checked_mul(2)?.try_into().ok()?)

    }

}

fn transmitq_idx(port_id: u32) -> Option<usize> {

    if port_id == 0 {

        Some(PORT0_TRANSMITQ_IDX)

    } else {

        PORT1_TRANSMITQ_IDX.checked_add((port_id - 1).checked_mul(2)?.try_into().ok()?)

    }

}

fn receiveq_port_id(queue_idx: usize) -> Option<u32> {

    if queue_idx == PORT0_RECEIVEQ_IDX {

        Some(0)

    } else if queue_idx >= PORT1_RECEIVEQ_IDX && (queue_idx & 1) == 0 {

        ((queue_idx - PORT1_RECEIVEQ_IDX) / 2)

            .checked_add(1)?

            .try_into()

            .ok()

    } else {

        None

    }

}

fn transmitq_port_id(queue_idx: usize) -> Option<u32> {

    if queue_idx == PORT0_TRANSMITQ_IDX {

        Some(0)

    } else if queue_idx >= PORT1_TRANSMITQ_IDX && (queue_idx & 1) == 1 {

        ((queue_idx - PORT1_TRANSMITQ_IDX) / 2)

            .checked_add(1)?

            .try_into()

            .ok()

    } else {

        None

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_receiveq_idx() {

        assert_eq!(receiveq_idx(0), Some(0));

        assert_eq!(receiveq_idx(1), Some(4));

        assert_eq!(receiveq_idx(2), Some(6));

        assert_eq!(receiveq_idx(3), Some(8));

    }

    #[test]

    fn test_transmitq_idx() {

        assert_eq!(transmitq_idx(0), Some(1));

        assert_eq!(transmitq_idx(1), Some(5));

        assert_eq!(transmitq_idx(2), Some(7));

        assert_eq!(transmitq_idx(3), Some(9));

    }

    #[test]

    fn test_receiveq_port_id() {

        assert_eq!(receiveq_port_id(0), Some(0));

        assert_eq!(receiveq_port_id(1), None); // port0 transmitq

        assert_eq!(receiveq_port_id(2), None); // ctrl receiveq

        assert_eq!(receiveq_port_id(3), None); // ctrl transmitq

        assert_eq!(receiveq_port_id(4), Some(1));

        assert_eq!(receiveq_port_id(5), None);

        assert_eq!(receiveq_port_id(6), Some(2));

        assert_eq!(receiveq_port_id(7), None);

        assert_eq!(receiveq_port_id(8), Some(3));

        assert_eq!(receiveq_port_id(9), None);

    }

    #[test]

    fn test_transmitq_port_id() {

        assert_eq!(transmitq_port_id(0), None); // port0 receiveq

        assert_eq!(transmitq_port_id(1), Some(0));

        assert_eq!(transmitq_port_id(2), None); // ctrl receiveq

        assert_eq!(transmitq_port_id(3), None); // ctrl transmitq

        assert_eq!(transmitq_port_id(4), None); // port1 receiveq

        assert_eq!(transmitq_port_id(5), Some(1));

        assert_eq!(transmitq_port_id(6), None);

        assert_eq!(transmitq_port_id(7), Some(2));

        assert_eq!(transmitq_port_id(8), None);

        assert_eq!(transmitq_port_id(9), Some(3));

    }

}