// Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::env::args_os;
use std::fs::remove_file;
use std::io::{stdout, Write};
use std::os::unix::prelude::*;
use std::process;

use devices::virtio::{remote, virtualized};
use msg_socket::{MsgReceiver, MsgSender, MsgSocket};
use sys_util::net::{UnixSeqpacket, UnixSeqpacketListener};
use sys_util::{PollContext, PollToken};
use vm_control::{VmMemoryRequest, VmMemoryResponse};

fn main() {
    // Create and display the incoming socket.
    let mut path = std::env::var_os("XDG_RUNTIME_DIR").expect("XDG_RUNTIME_DIR missing");
    path.push("/crosvm_wl-");
    path.push(process::id().to_string());
    path.push(".sock");
    let _ = remove_file(&path);
    let server = UnixSeqpacketListener::bind(&path).expect("failed to create socket");
    match stdout().write_all(path.as_bytes()) {
        Ok(()) => println!(),
        Err(e) => eprintln!("{}", e),
    }

    // Receive connection from crosvm.
    let conn = server.accept().expect("accept failed");
    let client_virtio_device_msg_socket: msg_socket2::Socket<remote::Response, remote::Request> =
        msg_socket2::Socket::new(conn);

    let (memory_params, client_vm_control_msg_socket) = match client_virtio_device_msg_socket.recv()
    {
        Ok(remote::Request::Create {
            memory_params,
            vm_control_socket,
        }) => (memory_params, vm_control_socket.owned()),

        Ok(msg) => panic!("unexpected message: {:?}", msg),
        Err(e) => panic!("recv error: {}", e),
    };

    let server_virtio_device_socket_path = args_os().nth(1).expect("missing server socket path");
    let server_virtio_device_socket =
        UnixSeqpacket::connect(server_virtio_device_socket_path).expect("connect failed");
    let server_virtio_device_msg_socket: msg_socket2::Socket<
        virtualized::Request,
        virtualized::Response,
    > = msg_socket2::Socket::new(server_virtio_device_socket);

    let (server_vm_control_socket, ext_vm_control_socket) =
        UnixSeqpacket::pair().expect("pair failed");
    let server_vm_control_msg_socket = MsgSocket::new(server_vm_control_socket);

    server_virtio_device_msg_socket
        .send(virtualized::Request::Create {
            memory_params,
            ext_socket: ext_vm_control_socket,
        })
        .expect("send failed");

    #[derive(Debug, PollToken)]
    enum Token {
        ClientVirtioDevice,
        ClientVmControl,
        ServerVirtioDevice,
        ServerVmControl,
    }

    let poll_ctx = match PollContext::build_with(&[
        (&client_virtio_device_msg_socket, Token::ClientVirtioDevice),
        (&client_vm_control_msg_socket, Token::ClientVmControl),
        (&server_virtio_device_msg_socket, Token::ServerVirtioDevice),
        (&server_vm_control_msg_socket, Token::ServerVmControl),
    ]) {
        Ok(pc) => pc,
        Err(e) => panic!("failed creating PollContext: {}", e),
    };

    loop {
        let events = match poll_ctx.wait() {
            Ok(v) => v,
            Err(e) => panic!("failed polling for events: {}", e),
        };

        for event in &events {
            match event.token() {
                Token::ClientVirtioDevice => {
                    let msg = match dbg!(client_virtio_device_msg_socket.recv()) {
                        Ok(remote::Request::DebugLabel) => virtualized::Request::DebugLabel,
                        Ok(remote::Request::DeviceType) => virtualized::Request::DeviceType,
                        Ok(remote::Request::QueueMaxSizes) => virtualized::Request::QueueMaxSizes,
                        Ok(remote::Request::Features) => virtualized::Request::Features,
                        Ok(remote::Request::AckFeatures(value)) => {
                            virtualized::Request::AckFeatures(value)
                        }
                        Ok(remote::Request::ReadConfig { offset, len }) => {
                            virtualized::Request::ReadConfig { offset, len }
                        }
                        Ok(remote::Request::WriteConfig { offset, data }) => {
                            virtualized::Request::WriteConfig { offset, data }
                        }
                        Ok(remote::Request::Activate {
                            shm,
                            interrupt,
                            interrupt_resample_evt,
                            queues,
                            queue_evts,
                        }) => virtualized::Request::Activate {
                            shm,
                            interrupt,
                            interrupt_resample_evt,
                            queues,
                            queue_evts,
                        },
                        Ok(remote::Request::Reset) => virtualized::Request::Reset,
                        Ok(remote::Request::GetDeviceBars(address)) => {
                            virtualized::Request::GetDeviceBars(address)
                        }
                        Ok(remote::Request::GetDeviceCaps) => virtualized::Request::GetDeviceCaps,
                        Ok(remote::Request::Kill) => virtualized::Request::Kill,

                        Ok(req @ remote::Request::Create { .. }) => {
                            panic!("unexpected message: {:?}", req)
                        }

                        Err(e) => panic!("recv failed: {}", e),
                    };

                    if let Err(e) = server_virtio_device_msg_socket.send(msg) {
                        panic!("send failed: {}", e);
                    }
                }

                Token::ClientVmControl => {
                    use VmMemoryResponse::*;
                    let msg = match client_vm_control_msg_socket.recv().expect("recv failed") {
                        RegisterMemory { pfn, slot } => RegisterMemory { pfn, slot },
                        AllocateAndRegisterGpuMemory {
                            fd,
                            pfn,
                            slot,
                            desc,
                        } => AllocateAndRegisterGpuMemory {
                            fd,
                            pfn,
                            slot,
                            desc,
                        },
                        Ok => Ok,
                        Err(error) => Err(error),
                    };

                    server_vm_control_msg_socket
                        .send(dbg!(&msg))
                        .expect("send failed");
                }

                Token::ServerVirtioDevice => {
                    let msg = match server_virtio_device_msg_socket.recv() {
                        Ok(virtualized::Response::DebugLabel(label)) => {
                            remote::Response::DebugLabel(label)
                        }
                        Ok(virtualized::Response::DeviceType(typ)) => {
                            remote::Response::DeviceType(typ)
                        }
                        Ok(virtualized::Response::Features(value)) => {
                            remote::Response::Features(value)
                        }
                        Ok(virtualized::Response::GetDeviceBars(bars)) => {
                            remote::Response::GetDeviceBars(bars)
                        }
                        Ok(virtualized::Response::GetDeviceCaps(caps)) => {
                            remote::Response::GetDeviceCaps(caps)
                        }
                        Ok(virtualized::Response::Kill) => remote::Response::Kill,
                        Ok(virtualized::Response::QueueMaxSizes(sizes)) => {
                            remote::Response::QueueMaxSizes(sizes)
                        }
                        Ok(virtualized::Response::ReadConfig(config)) => {
                            remote::Response::ReadConfig(config)
                        }
                        Ok(virtualized::Response::Reset(result)) => remote::Response::Reset(result),
                        Err(e) => panic!("recv failed: {}", e),
                    };

                    if let Err(e) = client_virtio_device_msg_socket.send(dbg!(msg)) {
                        panic!("send failed: {}", e);
                    }
                }

                Token::ServerVmControl => {
                    use VmMemoryRequest::*;
                    let msg = match server_vm_control_msg_socket.recv().expect("recv failed") {
                        RegisterMemory(fd, size) => RegisterMemory(fd, size),
                        RegisterFdAtPciBarOffset(alloc, fd, size, offset) => {
                            RegisterFdAtPciBarOffset(alloc, fd, size, offset)
                        }
                        UnregisterMemory(slot) => UnregisterMemory(slot),
                        AllocateAndRegisterGpuMemory {
                            width,
                            height,
                            format,
                        } => AllocateAndRegisterGpuMemory {
                            width,
                            height,
                            format,
                        },
                        RegisterMmapMemory {
                            fd,
                            size,
                            offset,
                            gpa,
                        } => RegisterMmapMemory {
                            fd,
                            size,
                            offset,
                            gpa,
                        },
                    };

                    client_vm_control_msg_socket
                        .send(dbg!(&msg))
                        .expect("send failed");
                }
            }
        }
    }
}