use crate::hw::uart::Uart;
use crate::{board, mmu};
use types::id::ProcessId;

const UART_VA: usize = 0x20000;
const INPUT_MMIO_VA: usize = 0x25000;
const INPUT_VQUEUE_VA: usize = 0x26000;
const INPUT_EVBUF_VA: usize = 0x28000;
const INPUT_PHYSINFO_VA: usize = 0x29000;

fn kern() -> &'static mut kernel::Kernel<crate::platform::RiscVPlatform> {
    unsafe { &mut *crate::kernel_ptr() }
}

pub fn configure(uart: &Uart, input_pid: ProcessId) {
    let pt_addr = kern().platform_mut().page_table_for(input_pid);
    let pt = unsafe { &mut *(pt_addr as *mut mmu::PageTable) };

    // Always map UART for fallback
    map_page(pt, UART_VA, board::UART_BASE);

    // Try to find VirtIO input device
    if let Some(mmio) = hw_virtio::mmio::scan_for_device(hw_virtio::mmio::DEVICE_ID_INPUT) {
        let base = mmio.base();
        map_page(pt, INPUT_MMIO_VA, base);

        let vq_phys = alloc_contiguous(2);
        map_range(pt, INPUT_VQUEUE_VA, vq_phys, 2);

        let evbuf_phys = mmu::allocator::alloc_zeroed_page();
        map_page(pt, INPUT_EVBUF_VA, evbuf_phys);

        let info_phys = mmu::allocator::alloc_zeroed_page();
        map_page(pt, INPUT_PHYSINFO_VA, info_phys);
        unsafe {
            core::ptr::write_volatile(info_phys as *mut u64, evbuf_phys as u64);
            core::ptr::write_volatile((info_phys as *mut u64).add(1), vq_phys as u64);
        }

        hw_virtio::mmio::init_legacy_device(&mmio, vq_phys, 32);
        uart.write_str("k-os: VirtIO keyboard configured\n");
    } else {
        uart.write_str("k-os: UART mapped for input\n");
    }
}

fn map_page(pt: &mut mmu::PageTable, va: usize, pa: usize) {
    mmu::map(pt, mmu::Mapping {
        va: types::address::VirtualAddress::from_raw(va),
        pa: types::address::PhysicalAddress::from_raw(pa),
        pages: 1,
        flags: mmu::Sv39Flags::user_rw(),
    });
}

fn map_range(pt: &mut mmu::PageTable, va_base: usize, pa_base: usize, count: usize) {
    for i in 0..count {
        map_page(pt, va_base + i * 4096, pa_base + i * 4096);
    }
}

fn alloc_contiguous(pages: usize) -> usize {
    let first = mmu::allocator::alloc_zeroed_page();
    for _ in 1..pages {
        mmu::allocator::alloc_zeroed_page();
    }
    first
}