#![no_std]
#![forbid(unsafe_code)]

mod keymap;

use core::sync::atomic::{Ordering, fence};
use hw_virtio::mmio::{MmioAccess, Registers};
use libkos::Window;
use port_input_device::{InputDevice, InputEvent};
use types::address::{ByteOffset, PhysicalAddress, VirtualAddress};
use types::bounded::PageCount;

/// Drivers map their devices in the device window (0x01D0_0000 up, 64 KiB
/// each), clear of program images and heaps.
const INPUT_MMIO_VA: VirtualAddress = VirtualAddress::from_raw(0x01D1_0000);
const VQUEUE_VA: VirtualAddress = VirtualAddress::from_raw(0x01D1_1000);
const EVBUF_VA: VirtualAddress = VirtualAddress::from_raw(0x01D1_3000);

const QUEUE_SIZE: usize = 32;
const EVENT_SIZE: usize = 8;
const DESC_SIZE: usize = 16;
/// Where the legacy layout puts the available and used rings.
const AVAIL: ByteOffset = ByteOffset::new(QUEUE_SIZE * DESC_SIZE);
const USED: ByteOffset = ByteOffset::new(4096);

const EV_KEY: u16 = 1;
const KEY_PRESS: u32 = 1;

pub struct VirtioInputDevice {
    regs: Registers,
    /// The queue's rings, and the buffers the device writes events into.
    rings: Window,
    events: Window,
    events_pa: PhysicalAddress,
    last_used: u16,
    next_avail: u16,
    shift: bool,
    ctrl: bool,
}

impl VirtioInputDevice {
    /// Takes over the virtio input device behind `mmio`. `None` if it is
    /// not one or any step fails.
    pub fn try_new(mmio: types::bounded::CapSlot) -> Option<Self> {
        let regs = Registers::new(libkos::device::map_mmio(mmio, INPUT_MMIO_VA).ok()?);
        if regs.read(hw_virtio::mmio::REG_MAGIC) != hw_virtio::mmio::MAGIC
            || regs.read(hw_virtio::mmio::REG_DEVICE_ID) != hw_virtio::mmio::DEVICE_ID_INPUT
        {
            return None;
        }
        let rings = libkos::device::dma(PageCount::new(2)?, VQUEUE_VA).ok()?;
        let events = libkos::device::dma(PageCount::new(1)?, EVBUF_VA).ok()?;
        hw_virtio::mmio::init_legacy_device(&regs, rings.physical, QUEUE_SIZE as u32);

        let mut dev = Self {
            regs,
            rings: rings.window,
            events: events.window,
            events_pa: events.physical,
            last_used: 0,
            next_avail: 0,
            shift: false,
            ctrl: false,
        };
        dev.fill_queue();
        Some(dev)
    }

    /// Every buffer goes to the device, one event each.
    fn fill_queue(&mut self) {
        for i in 0..QUEUE_SIZE {
            let desc = ByteOffset::new(i * DESC_SIZE);
            self.rings.write(desc, (self.events_pa.raw() + i * EVENT_SIZE) as u64);
            self.rings.write(desc + 8, EVENT_SIZE as u32);
            self.rings.write(desc + 12, 2u16); // WRITE flag
            self.rings.write(desc + 14, 0u16); // no next
            self.rings.write(AVAIL + 4 + i * 2, i as u16);
        }
        fence(Ordering::SeqCst);
        self.rings.write(AVAIL + 2, QUEUE_SIZE as u16);
        fence(Ordering::SeqCst);
        self.regs.write(hw_virtio::mmio::REG_QUEUE_NOTIFY, 0);
        self.next_avail = QUEUE_SIZE as u16;
    }

    fn resubmit(&mut self, desc_idx: u16) {
        let slot = self.next_avail as usize % QUEUE_SIZE;
        self.rings.write(AVAIL + 4 + slot * 2, desc_idx);
        fence(Ordering::SeqCst);
        self.next_avail = self.next_avail.wrapping_add(1);
        self.rings.write(AVAIL + 2, self.next_avail);
        fence(Ordering::SeqCst);
        self.regs.write(hw_virtio::mmio::REG_QUEUE_NOTIFY, 0);
    }

    fn read_event(&self, desc_idx: u16) -> (u16, u16, u32) {
        let at = ByteOffset::new(desc_idx as usize % QUEUE_SIZE * EVENT_SIZE);
        (self.events.read(at), self.events.read(at + 2), self.events.read(at + 4))
    }

    fn used_idx(&self) -> u16 {
        self.rings.read(USED + 2)
    }

    fn used_entry_id(&self, idx: u16) -> u16 {
        self.rings.read::<u32>(USED + 4 + (idx as usize % QUEUE_SIZE) * 8) as u16
    }
}

impl InputDevice for VirtioInputDevice {
    fn poll(&mut self) -> Option<InputEvent> {
        loop {
            let current_used = self.used_idx();
            if current_used == self.last_used {
                return None;
            }

            let desc_idx = self.used_entry_id(self.last_used);
            self.last_used = self.last_used.wrapping_add(1);

            let (ev_type, ev_code, ev_value) = self.read_event(desc_idx);
            self.resubmit(desc_idx);

            if ev_type == EV_KEY && matches!(ev_code, keymap::LEFT_SHIFT | keymap::RIGHT_SHIFT) {
                self.shift = ev_value != 0;
                continue;
            }
            if ev_type == EV_KEY && matches!(ev_code, keymap::LEFT_CTRL | keymap::RIGHT_CTRL) {
                self.ctrl = ev_value != 0;
                continue;
            }
            if ev_type == EV_KEY {
                let key = keymap::key(ev_code, self.shift).map(|k| if self.ctrl { keymap::with_ctrl(k) } else { k });
                if let Some(key) = key {
                    if ev_value == KEY_PRESS {
                        return Some(InputEvent::KeyDown(key));
                    } else if ev_value == 0 {
                        return Some(InputEvent::KeyUp(key));
                    }
                }
            }
            // Skip SYN events, repeats, and unknown keys — keep draining
        }
    }

    fn has_event(&self) -> bool {
        self.used_idx() != self.last_used
    }

    fn acknowledge(&mut self) {
        let status = self.regs.read(hw_virtio::mmio::REG_INTERRUPT_STATUS);
        self.regs.write(hw_virtio::mmio::REG_INTERRUPT_ACK, status);
    }
}