// Copyright (c) 2021 Intel Corporation
//
// SPDX-License-Identifier: Apache-2.0
//

use anyhow::{anyhow, Result};
use futures::StreamExt;
use std::io;
use std::io::ErrorKind;
use std::os::unix::io::{FromRawFd, RawFd};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::watch::Receiver;
use tokio_vsock::{Incoming, VsockListener, VsockStream};
use tracing::instrument;

// Size of I/O read buffer
const BUF_SIZE: usize = 8192;

// Interruptable I/O copy using readers and writers
// (an interruptable version of "io::copy()").
pub async fn interruptable_io_copier<R, W>(
    mut reader: R,
    mut writer: W,
    mut shutdown: Receiver<bool>,
) -> io::Result<u64>
where
    R: tokio::io::AsyncRead + Unpin + Sized,
    W: tokio::io::AsyncWrite + Unpin + Sized,
{
    let mut total_bytes: u64 = 0;

    let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];

    loop {
        tokio::select! {
            _ = shutdown.changed() => {
                eprintln!("INFO: interruptable_io_copier: got shutdown request");
                break;
            },

            result = reader.read(&mut buf) => {
                let bytes = match result {
                    Ok(0) => return Ok(total_bytes),
                    Ok(len) => len,
                    Err(ref e) if e.kind() == ErrorKind::Interrupted => continue,
                    Err(e) => return Err(e),
                };

                total_bytes += bytes as u64;

                // Actually copy the data ;)
                writer.write_all(&buf[..bytes]).await?;
            },
        };
    }

    Ok(total_bytes)
}

#[instrument]
pub fn get_vsock_incoming(fd: RawFd) -> Incoming {
    unsafe { VsockListener::from_raw_fd(fd).incoming() }
}

#[instrument]
pub async fn get_vsock_stream(fd: RawFd) -> Result<VsockStream> {
    let stream = get_vsock_incoming(fd)
        .next()
        .await
        .ok_or_else(|| anyhow!("cannot handle incoming vsock connection"))?;

    Ok(stream?)
}

pub fn merge<T>(v1: &mut Option<Vec<T>>, v2: &Option<Vec<T>>) -> Option<Vec<T>>
where
    T: Clone,
{
    let mut result = v1.clone().map(|mut vec| {
        if let Some(ref other) = v2 {
            vec.extend(other.iter().cloned());
        }
        vec
    });

    if result.is_none() {
        result.clone_from(v2);
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use rstest::rstest;
    use std::io;
    use std::io::Cursor;
    use std::io::Write;
    use std::pin::Pin;
    use std::sync::{Arc, Mutex};
    use std::task::{Context, Poll, Poll::Ready};
    use tokio::pin;
    use tokio::select;
    use tokio::sync::watch::channel;
    use tokio::task::JoinError;
    use tokio::time::Duration;

    #[derive(Debug, Default, Clone)]
    struct BufWriter {
        data: Arc<Mutex<Vec<u8>>>,
        write_delay: Duration,
    }

    impl BufWriter {
        fn new() -> Self {
            BufWriter {
                data: Arc::new(Mutex::new(Vec::<u8>::new())),
                write_delay: Duration::new(0, 0),
            }
        }

        fn write_vec(&mut self, buf: &[u8]) -> io::Result<usize> {
            let vec_ref = self.data.clone();

            let mut vec_locked = vec_ref.lock();

            let mut v = vec_locked.as_deref_mut().unwrap();

            if self.write_delay.as_nanos() > 0 {
                std::thread::sleep(self.write_delay);
            }

            std::io::Write::write(&mut v, buf)
        }
    }

    impl Write for BufWriter {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.write_vec(buf)
        }

        fn flush(&mut self) -> io::Result<()> {
            let vec_ref = self.data.clone();

            let mut vec_locked = vec_ref.lock();

            let v = vec_locked
                .as_deref_mut()
                .map_err(|e| std::io::Error::other(e.to_string()))?;

            std::io::Write::flush(v)
        }
    }

    impl tokio::io::AsyncWrite for BufWriter {
        fn poll_write(
            mut self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<Result<usize, io::Error>> {
            let result = self.write_vec(buf);

            Ready(result)
        }

        fn poll_flush(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), std::io::Error>> {
            // NOP
            Ready(Ok(()))
        }

        fn poll_shutdown(
            self: Pin<&mut Self>,
            _cx: &mut Context<'_>,
        ) -> Poll<Result<(), std::io::Error>> {
            // NOP
            Ready(Ok(()))
        }
    }

    impl std::fmt::Display for BufWriter {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            let data_ref = self.data.clone();
            let output = data_ref.lock().unwrap();
            let s = (*output).clone();

            write!(f, "{}", String::from_utf8(s).unwrap())
        }
    }

    #[rstest]
    #[case("".into())]
    #[case("a".into())]
    #[case("foo".into())]
    #[case("b".repeat(BUF_SIZE - 1))]
    #[case("c".repeat(BUF_SIZE))]
    #[case("d".repeat(BUF_SIZE + 1))]
    #[case("e".repeat((2 * BUF_SIZE) - 1))]
    #[case("f".repeat(2 * BUF_SIZE))]
    #[case("g".repeat((2 * BUF_SIZE) + 1))]
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_interruptable_io_copier_reader(#[case] reader_value: String) {
        let (tx, rx) = channel(true);
        let reader = Cursor::new(reader_value.clone());
        let writer = BufWriter::new();

        // XXX: Pass a copy of the writer to the copier to allow the
        // result of the write operation to be checked below.
        let handle = tokio::spawn(interruptable_io_copier(reader, writer.clone(), rx));

        // Allow time for the thread to be spawned.
        tokio::time::sleep(Duration::from_secs(1)).await;

        let timeout = tokio::time::sleep(Duration::from_secs(1));
        pin!(timeout);

        // Since the readers only specify a small number of bytes, the
        // copier will quickly read zero and kill the task, closing the
        // Receiver.
        assert!(tx.is_closed());

        let spawn_result: std::result::Result<std::result::Result<u64, std::io::Error>, JoinError>;

        select! {
            res = handle => spawn_result = res,
            _ = &mut timeout => panic!("timed out"),
        }

        assert!(spawn_result.is_ok());

        let result: std::result::Result<u64, std::io::Error> = spawn_result.unwrap();

        assert!(result.is_ok());

        let byte_count = result.unwrap() as usize;
        assert_eq!(byte_count, reader_value.len());

        let value = writer.to_string();
        assert_eq!(value, reader_value);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_interruptable_io_copier_eof() {
        // Create an async reader that always returns EOF
        let reader = tokio::io::empty();

        let (tx, rx) = channel(true);
        let writer = BufWriter::new();

        let handle = tokio::spawn(interruptable_io_copier(reader, writer.clone(), rx));

        // Allow time for the thread to be spawned.
        tokio::time::sleep(Duration::from_secs(1)).await;

        let timeout = tokio::time::sleep(Duration::from_secs(1));
        pin!(timeout);

        assert!(tx.is_closed());

        let spawn_result: std::result::Result<std::result::Result<u64, std::io::Error>, JoinError>;

        select! {
            res = handle => spawn_result = res,
            _ = &mut timeout => panic!("timed out"),
        }

        assert!(spawn_result.is_ok());

        let result: std::result::Result<u64, std::io::Error> = spawn_result.unwrap();

        assert!(result.is_ok());

        let byte_count = result.unwrap();
        assert_eq!(byte_count, 0);
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_interruptable_io_copier_shutdown() {
        // Create an async reader that creates an infinite stream of bytes
        // (which allows us to interrupt it, since we know it is always busy ;)
        const REPEAT_CHAR: u8 = b'r';

        let reader = tokio::io::repeat(REPEAT_CHAR);

        let (tx, rx) = channel(true);
        let writer = BufWriter::new();

        let handle = tokio::spawn(interruptable_io_copier(reader, writer.clone(), rx));

        // Allow time for the thread to be spawned.
        tokio::time::sleep(Duration::from_secs(1)).await;

        let timeout = tokio::time::sleep(Duration::from_secs(1));
        pin!(timeout);

        assert!(!tx.is_closed());

        tx.send(true).expect("failed to request shutdown");

        let spawn_result: std::result::Result<std::result::Result<u64, std::io::Error>, JoinError>;

        select! {
            res = handle => spawn_result = res,
            _ = &mut timeout => panic!("timed out"),
        }

        assert!(spawn_result.is_ok());

        let result: std::result::Result<u64, std::io::Error> = spawn_result.unwrap();

        assert!(result.is_ok());

        let byte_count = result.unwrap();

        let value = writer.to_string();

        let writer_byte_count = value.len() as u64;

        assert_eq!(byte_count, writer_byte_count);

        // Remove the char used as a payload. If anything else remins,
        // something went wrong.
        let mut remainder = value;

        remainder.retain(|c| c != REPEAT_CHAR as char);

        assert_eq!(remainder.len(), 0);
    }
}
