#![no_std]
#![no_main]
#![feature(asm_const)]
#![feature(abi_x86_interrupt)]

#[macro_use]
extern crate log;
extern crate alloc;

mod boot;
mod hal;
mod heap;
mod logger;
mod trap;

#[cfg(feature = "tcp")]
mod tcp;

use self::hal::HalImpl;
use virtio_drivers::{
    device::{blk::VirtIOBlk, gpu::VirtIOGpu, net::VirtIONet},
    transport::{
        pci::{
            bus::{BarInfo, Cam, Command, DeviceFunction, PciRoot},
            virtio_device_type, PciTransport,
        },
        DeviceType, Transport,
    },
};

/// Memory mapped address space to access PCI configuration.
///
/// Currently it is hardcoded (from qemu/roms/seabios/src/hw/dev-q35.h)
///
/// TODO: get it from ACPI MCFG table.
const MMCONFIG_BASE: usize = 0xB000_0000;

const NET_BUFFER_LEN: usize = 2048;
const NET_QUEUE_SIZE: usize = 16;

fn system_off() -> ! {
    use x86_64::instructions::{hlt, port::PortWriteOnly};
    unsafe {
        PortWriteOnly::new(0x604).write(0x2000u16);
        loop {
            hlt();
        }
    }
}

#[no_mangle]
extern "C" fn main(_mbi: *const u8) -> ! {
    logger::init(log::LevelFilter::Info).unwrap();
    info!("virtio-drivers example started.");

    trap::init();
    heap::init_heap();

    enumerate_pci(MMCONFIG_BASE as _);

    info!("test end");
    system_off();
}

fn virtio_device(transport: impl Transport) {
    match transport.device_type() {
        DeviceType::Block => virtio_blk(transport),
        DeviceType::GPU => virtio_gpu(transport),
        DeviceType::Network => virtio_net(transport),
        t => warn!("Unrecognized virtio device: {:?}", t),
    }
}

fn virtio_blk<T: Transport>(transport: T) {
    let mut blk = VirtIOBlk::<HalImpl, T>::new(transport).expect("failed to create blk driver");
    assert!(!blk.readonly());
    let mut input = [0xffu8; 512];
    let mut output = [0; 512];
    for i in 0..32 {
        for x in input.iter_mut() {
            *x = i as u8;
        }
        blk.write_block(i, &input).expect("failed to write");
        blk.read_block(i, &mut output).expect("failed to read");
        assert_eq!(input, output);
    }
    info!("virtio-blk test finished");
}

fn virtio_gpu<T: Transport>(transport: T) {
    let mut gpu = VirtIOGpu::<HalImpl, T>::new(transport).expect("failed to create gpu driver");
    let (width, height) = gpu.resolution().expect("failed to get resolution");
    let width = width as usize;
    let height = height as usize;
    info!("GPU resolution is {}x{}", width, height);
    let fb = gpu.setup_framebuffer().expect("failed to get fb");
    for y in 0..height {
        for x in 0..width {
            let idx = (y * width + x) * 4;
            fb[idx] = x as u8;
            fb[idx + 1] = y as u8;
            fb[idx + 2] = (x + y) as u8;
        }
    }
    gpu.flush().expect("failed to flush");
    //delay some time
    info!("virtio-gpu show graphics....");
    for _ in 0..10000 {
        for _ in 0..100000 {
            unsafe {
                core::arch::asm!("nop");
            }
        }
    }

    info!("virtio-gpu test finished");
}

fn virtio_net<T: Transport>(transport: T) {
    let net = VirtIONet::<HalImpl, T, NET_QUEUE_SIZE>::new(transport, NET_BUFFER_LEN)
        .expect("failed to create net driver");
    info!("MAC address: {:02x?}", net.mac_address());

    #[cfg(not(feature = "tcp"))]
    {
        let mut net = net;
        loop {
            match net.receive() {
                Ok(buf) => {
                    info!("RECV {} bytes: {:02x?}", buf.packet_len(), buf.packet());
                    let tx_buf = virtio_drivers::device::net::TxBuffer::from(buf.packet());
                    net.send(tx_buf).expect("failed to send");
                    net.recycle_rx_buffer(buf).unwrap();
                    break;
                }
                Err(virtio_drivers::Error::NotReady) => continue,
                Err(err) => panic!("failed to recv: {:?}", err),
            }
        }
        info!("virtio-net test finished");
    }

    #[cfg(feature = "tcp")]
    tcp::test_echo_server(net);
}

fn enumerate_pci(mmconfig_base: *mut u8) {
    info!("mmconfig_base = {:#x}", mmconfig_base as usize);

    let mut pci_root = unsafe { PciRoot::new(mmconfig_base, Cam::Ecam) };
    for (device_function, info) in pci_root.enumerate_bus(0) {
        let (status, command) = pci_root.get_status_command(device_function);
        info!(
            "Found {} at {}, status {:?} command {:?}",
            info, device_function, status, command
        );
        if let Some(virtio_type) = virtio_device_type(&info) {
            info!("  VirtIO {:?}", virtio_type);

            // Enable the device to use its BARs.
            pci_root.set_command(
                device_function,
                Command::IO_SPACE | Command::MEMORY_SPACE | Command::BUS_MASTER,
            );
            dump_bar_contents(&mut pci_root, device_function, 4);

            let mut transport =
                PciTransport::new::<HalImpl>(&mut pci_root, device_function).unwrap();
            info!(
                "Detected virtio PCI device with device type {:?}, features {:#018x}",
                transport.device_type(),
                transport.read_device_features(),
            );
            virtio_device(transport);
        }
    }
}

fn dump_bar_contents(root: &mut PciRoot, device_function: DeviceFunction, bar_index: u8) {
    let bar_info = root.bar_info(device_function, bar_index).unwrap();
    trace!("Dumping bar {}: {:#x?}", bar_index, bar_info);
    if let BarInfo::Memory { address, size, .. } = bar_info {
        let start = address as *const u8;
        unsafe {
            let mut buf = [0u8; 32];
            for i in 0..size / 32 {
                let ptr = start.add(i as usize * 32);
                core::ptr::copy(ptr, buf.as_mut_ptr(), 32);
                if buf.iter().any(|b| *b != 0xff) {
                    trace!("  {:?}: {:x?}", ptr, buf);
                }
            }
        }
    }
    trace!("End of dump");
}

#[panic_handler]
fn panic(info: &core::panic::PanicInfo) -> ! {
    error!("{}", info);
    system_off()
}