starry-kernel 0.10.1

use alloc::{collections::VecDeque, format, string::ToString, sync::Arc, vec, vec::Vec};
use core::{
    any::Any,
    mem::offset_of,
    sync::atomic::{AtomicU8, AtomicU32, Ordering},
    time::Duration,
};

/// Number of registered `/dev/input/event*` nodes. Populated by
/// [`input_devices`] at boot and read by sysfs so
/// `/sys/class/input/event<N>` matches reality.
static EVENT_DEVICE_COUNT: AtomicU32 = AtomicU32::new(0);

/// Returns the number of `/dev/input/event*` devices currently exposed.
pub fn input_device_count() -> u32 {
    EVENT_DEVICE_COUNT.load(Ordering::Acquire)
}

use ax_input::{ErasedInputDevice, Event, EventType, InputDevice, InputDeviceId, InputError};
use ax_lazyinit::OnceLock;
use ax_runtime::hal::{irq::IrqId, time::wall_time};
use ax_std::os::arceos::{
    task as scheduler,
    task::sync::{
        WaitQueue,
        irq::{IrqWaitCell, IrqWaitRegistration},
    },
};
use axfs_ng_vfs::{DeviceId, NodeFlags, NodeType, VfsError, VfsResult};
use axpoll::{ExclusiveRegistrationSink, IoEvents, Pollable, SharedRegistrationSink};
use axpoll_set::PollSet;
use bitmaps::Bitmap;
use linux_raw_sys::{
    general::{__kernel_old_time_t, __kernel_suseconds_t},
    ioctl::{EVIOCGID, EVIOCGRAB, EVIOCGVERSION},
};
use zerocopy::{FromBytes, Immutable, IntoBytes};

use crate::{
    mm::UserPtr,
    pseudofs::{
        Device, DeviceOps, DirMapping, SimpleFs, dev::irq_service::complete_irq_service_cycle,
    },
    sync::IrqMutex,
};
const KEY_CNT: usize = EventType::Key.bits_count();

/// Bound on the in-kernel evdev buffer. Linux uses a per-client ring of
/// 64 entries by default; we hold a bit more headroom so a 20-key burst
/// (key down + key up + EV_SYN per key = 60 entries) never drops events
/// before userspace drains it. When the queue is full we follow Linux's
/// behavior and drop the oldest entry rather than blocking the driver.
const READ_AHEAD_CAP: usize = 256;

const IRQ_SERVICE_STOPPED: u8 = 0;
const IRQ_SERVICE_STARTING: u8 = 1;
const IRQ_SERVICE_STARTED: u8 = 2;

struct Inner {
    device: ErasedInputDevice,
    read_ahead: VecDeque<(Duration, Event)>,
    key_state: Bitmap<KEY_CNT>,
}
impl Inner {
    /// Drain everything the driver currently has buffered into `read_ahead`,
    /// updating cached key state along the way. Stops at the first
    /// `InputError::Again` (driver queue empty) or after a hard ceiling of
    /// `READ_AHEAD_CAP` pulls per call to bound a single pass.
    ///
    /// Returns `true` if at least one event is now queued for userspace.
    fn drain_into_queue(&mut self) -> bool {
        for _ in 0..READ_AHEAD_CAP {
            match self.device.read_event() {
                Ok(event) => {
                    if event.event_type == EventType::Key as u16 {
                        if event.value == 0 {
                            self.key_state.set(event.code as usize, false);
                        } else if event.value == 1 {
                            self.key_state.set(event.code as usize, true);
                        }
                    }
                    if self.read_ahead.len() >= READ_AHEAD_CAP {
                        // Mirror Linux evdev: drop oldest on overflow so
                        // the most recent input wins. Keeps the driver
                        // ring from stalling under a burst we cannot
                        // forward to a slow reader.
                        self.read_ahead.pop_front();
                    }
                    self.read_ahead.push_back((wall_time(), event));
                }
                Err(InputError::Again) => break,
                Err(err) => {
                    warn!("Failed to read event: {err:?}");
                    break;
                }
            }
        }
        !self.read_ahead.is_empty()
    }

    fn has_event(&mut self) -> bool {
        self.drain_into_queue()
    }
}

/// Linux `INPUT_PROP_CNT` — the property bitmap is 4 bytes (32 properties).
const INPUT_PROP_CNT: usize = 0x20;
/// Linux `INPUT_PROP_POINTER` — emulates a relative pointer or maps absolute
/// coordinates to screen space. libinput hides the cursor on absolute-axis
/// devices that do not advertise this until proximity is reported.
const INPUT_PROP_POINTER: usize = 0x00;
/// Linux `INPUT_PROP_DIRECT` — direct-mapped axes (touchscreens).
const INPUT_PROP_DIRECT: usize = 0x01;

/// Linux uapi `struct input_absinfo` — six `i32`s returned by
/// `EVIOCGABS(axis)`.
#[repr(C)]
#[derive(Default, Clone, Copy, FromBytes, IntoBytes, Immutable)]
struct InputAbsInfo {
    value: i32,
    minimum: i32,
    maximum: i32,
    fuzz: i32,
    flat: i32,
    resolution: i32,
}

/// Maximum number of absolute axes Linux's EVIOCGABS encodes (0..0x3F).
const ABS_MAX: usize = 0x40;

pub struct EventDev {
    inner: IrqMutex<Inner>,
    waiters: PollSet,
    /// IRQ domain id the runtime resolved for the underlying driver.
    irq: Option<IrqId>,
    irq_handle: OnceLock<ax_runtime::hal::irq::IrqHandle>,
    irq_notify: IrqWaitCell,
    irq_service_park: WaitQueue,
    irq_service_state: AtomicU8,
    ev_bits: Bitmap<{ EventType::COUNT as usize }>,
    /// Cached `EVIOCGPROP` bitmap. Computed once at probe from the driver's
    /// raw bits with a synthesized `INPUT_PROP_POINTER` for absolute or
    /// relative pointing devices that aren't touchscreens. QEMU's
    /// virtio-mouse / virtio-tablet do not set the bit themselves, so
    /// libinput would otherwise classify the tablet as a graphics tablet
    /// and suppress the cursor pending a never-firing `BTN_TOOL_PEN`.
    prop_bits: [u8; INPUT_PROP_CNT.div_ceil(8)],
    /// Cached `EV_ABS` bitmap. Used by `EVIOCGABS` to refuse axes the
    /// device doesn't advertise with `EINVAL`, matching Linux's
    /// `evdev_handle_get_val` behavior. virtio-drivers reports the
    /// underlying `Error::IoError` when the AbsInfo selector has size 0,
    /// which would otherwise surface as EIO and confuse libinput.
    abs_bits: [u8; ABS_MAX.div_ceil(8)],
}

impl EventDev {
    pub fn new(mut device: ErasedInputDevice) -> Self {
        let mut ev_bits = Bitmap::new();
        for i in 0..EventType::COUNT {
            let Some(ty) = EventType::from_repr(i) else {
                continue;
            };
            if device
                .get_event_bits(ty, &mut [])
                .is_ok_and(|success| success)
            {
                ev_bits.set(i as usize, true);
            }
        }

        let mut prop_bits = [0u8; INPUT_PROP_CNT.div_ceil(8)];
        let prop_bits_reliable = match device.get_prop_bits(&mut prop_bits) {
            Ok(_) => true,
            Err(err) => {
                warn!("Failed to get input property bits: {err:?}");
                false
            }
        };
        let is_touchscreen = prop_bits[INPUT_PROP_DIRECT / 8] & (1 << (INPUT_PROP_DIRECT % 8)) != 0;
        let has_axes =
            ev_bits.get(EventType::Relative as usize) || ev_bits.get(EventType::Absolute as usize);
        if prop_bits_reliable && has_axes && !is_touchscreen {
            prop_bits[INPUT_PROP_POINTER / 8] |= 1 << (INPUT_PROP_POINTER % 8);
        }

        let mut abs_bits = [0u8; ABS_MAX.div_ceil(8)];
        if ev_bits.get(EventType::Absolute as usize) {
            let _ = device.get_event_bits(EventType::Absolute, &mut abs_bits);
        }

        let irq = device.irq_id();
        Self {
            inner: IrqMutex::new(Inner {
                device,
                read_ahead: VecDeque::with_capacity(READ_AHEAD_CAP),
                key_state: Bitmap::new(),
            }),
            waiters: PollSet::new(),
            irq,
            irq_handle: OnceLock::new(),
            irq_notify: IrqWaitCell::new(),
            irq_service_park: WaitQueue::new(),
            irq_service_state: AtomicU8::new(IRQ_SERVICE_STOPPED),
            ev_bits,
            prop_bits,
            abs_bits,
        }
    }

    fn axis_supported(&self, axis: u8) -> bool {
        let bit = axis as usize;
        if bit >= ABS_MAX {
            return false;
        }
        self.abs_bits[bit / 8] & (1 << (bit % 8)) != 0
    }

    fn get_event_bits(
        &self,
        current: &crate::task::UserTaskRef,
        arg: usize,
        size: usize,
        ty: u8,
    ) -> VfsResult<usize> {
        if ty == 0 {
            write_user_bytes(current, arg, size, self.ev_bits.as_bytes())
        } else {
            let ty = EventType::from_repr(ty).ok_or(VfsError::InvalidInput)?;
            let mut kernel_bits = vec![0; size];
            {
                let mut inner = self.inner.lock();
                match inner.device.get_event_bits(ty, &mut kernel_bits) {
                    Ok(true) => {}
                    Ok(false) => {
                        debug!("No events for {ty:?}");
                    }
                    Err(err) => {
                        warn!("Failed to get event bits: {err:?}");
                    }
                }
            }
            let bytes = size.min(ty.bits_count().div_ceil(8));
            UserPtr::<u8>::from(arg).write_slice(current, &kernel_bits[..bytes])?;
            Ok(bytes)
        }
    }

    fn register_irq(self: &Arc<Self>) {
        let Some(irq) = self.irq else {
            return;
        };
        let event_dev = Arc::clone(self);
        let request = ax_runtime::hal::irq::IrqRequest::new(move |_| event_dev.handle_irq())
            .share_mode(ax_runtime::hal::irq::ShareMode::Shared)
            .auto_enable(ax_runtime::hal::irq::AutoEnable::No);
        match ax_runtime::hal::irq::request_irq(irq, request) {
            Ok(handle) => {
                if !self.start_irq_service() {
                    warn!("failed to start evdev IRQ service for irq {irq:?}");
                    return;
                }
                self.irq_handle.call_once(|| handle);
                self.inner.lock().device.enable_irq();
                if let Some(handle) = self.irq_handle.get().copied()
                    && let Err(err) = ax_runtime::hal::irq::enable_irq(handle)
                {
                    warn!("failed to enable evdev irq handler for irq {irq:?}: {err:?}");
                    self.inner.lock().device.disable_irq();
                }
            }
            Err(err) => {
                warn!("failed to register evdev irq handler for irq {irq:?}: {err:?}");
                self.inner.lock().device.disable_irq();
            }
        }
    }

    fn start_irq_service(self: &Arc<Self>) -> bool {
        if self
            .irq_service_state
            .compare_exchange(
                IRQ_SERVICE_STOPPED,
                IRQ_SERVICE_STARTING,
                Ordering::AcqRel,
                Ordering::Acquire,
            )
            .is_err()
        {
            return self.irq_service_state.load(Ordering::Acquire) == IRQ_SERVICE_STARTED;
        }

        let event_dev = Arc::clone(self);
        match crate::task::kernel_thread_builder("evdev-irq-service".into())
            .spawn(move || event_dev.run_irq_service())
        {
            Ok(_service) => {
                self.irq_service_state
                    .store(IRQ_SERVICE_STARTED, Ordering::Release);
                true
            }
            Err(error) => {
                self.irq_service_state
                    .store(IRQ_SERVICE_STOPPED, Ordering::Release);
                warn!("failed to spawn evdev IRQ service: {error}");
                false
            }
        }
    }

    fn run_irq_service(self: Arc<Self>) {
        let current = scheduler::thread::current::current_thread_handle()
            .unwrap_or_else(|error| panic!("evdev IRQ service has no scheduler thread: {error}"));
        let waiter = EventIrqWaiter::new(&current);

        loop {
            let registration = self.irq_notify.register(&waiter.registration);
            let completed = complete_irq_service_cycle(
                registration,
                |token| self.irq_service_park.wait_until(|| !token.is_attached()),
                || self.drain_irq_events(),
            )
            .unwrap_or_else(|error| panic!("evdev IRQ waiter could not quiesce: {error}"));
            if !completed {
                panic!("evdev IRQ service registration was occupied concurrently");
            }
        }
    }

    fn drain_irq_events(&self) {
        let ready = self.inner.lock().drain_into_queue();
        if ready {
            unsafe { self.waiters.wake(IoEvents::IN) };
        }
    }

    fn handle_irq(&self) -> ax_runtime::hal::irq::IrqReturn {
        // Use `lock()` rather than `try_lock()` so the virtio ISR is always
        // acknowledged. `IrqMutex` guarantees the holder has local IRQs
        // disabled, so this IRQ can only fire on a different CPU. Without the
        // ack, a level-triggered shared IRQ line stays asserted and can starve
        // other devices on the same line.
        let mut inner = self.inner.lock();
        let event = inner.device.handle_irq();
        drop(inner);
        if event.input_ready {
            let _result = self.irq_notify.notify();
            return ax_runtime::hal::irq::IrqReturn::Wake;
        }
        if event.handled {
            ax_runtime::hal::irq::IrqReturn::Handled
        } else {
            ax_runtime::hal::irq::IrqReturn::Unhandled
        }
    }
}

struct EventIrqWaiter {
    registration: IrqWaitRegistration,
}

impl EventIrqWaiter {
    fn new(current: &scheduler::thread::ThreadHandle) -> Self {
        Self {
            registration: IrqWaitRegistration::new(current.wake_handle()),
        }
    }
}

fn write_user_bytes(
    current: &crate::task::UserTaskRef,
    arg: usize,
    capacity: usize,
    source: &[u8],
) -> VfsResult<usize> {
    let len = source.len().min(capacity);
    UserPtr::<u8>::from(arg).write_slice(current, &source[..len])?;
    Ok(len)
}

fn return_str(
    current: &crate::task::UserTaskRef,
    arg: usize,
    size: usize,
    s: &str,
) -> VfsResult<usize> {
    write_user_bytes(current, arg, size, s.as_bytes())
}

fn input_error_to_vfs_error(err: InputError) -> VfsError {
    match err {
        InputError::AlreadyExists => VfsError::AlreadyExists,
        InputError::Again => VfsError::WouldBlock,
        InputError::BadState => VfsError::BadState,
        InputError::InvalidInput | InputError::Unsupported => VfsError::InvalidInput,
        InputError::Io => VfsError::Io,
        InputError::NoMemory => VfsError::NoMemory,
        InputError::ResourceBusy => VfsError::ResourceBusy,
    }
}

fn return_zero_bits(
    current: &crate::task::UserTaskRef,
    arg: usize,
    size: usize,
    bits: usize,
) -> VfsResult<usize> {
    let len = bits.div_ceil(8).min(size);
    UserPtr::<u8>::from(arg).write_slice(current, &vec![0; len])?;
    Ok(len)
}

#[repr(C)]
#[derive(FromBytes, IntoBytes, Immutable)]
pub struct KernelTimeval {
    pub tv_sec: __kernel_old_time_t,
    pub tv_usec: __kernel_suseconds_t,
}

#[repr(C)]
#[derive(FromBytes, IntoBytes, Immutable)]
struct InputEvent {
    time: KernelTimeval,
    event_type: u16,
    code: u16,
    value: i32,
}

#[unsafe(no_mangle)]
#[inline(never)]
pub extern "C" fn ongkey() {
    core::hint::black_box(());
}

impl DeviceOps for EventDev {
    fn read_at(&self, buf: &mut [u8], _offset: u64) -> VfsResult<usize> {
        if buf.is_empty() {
            return Ok(0);
        }
        if buf.len() < size_of::<InputEvent>() {
            return Err(VfsError::InvalidInput);
        }
        let mut read = 0;
        let mut inner = self.inner.lock();
        // Drain the driver queue once up front so a single read() syscall
        // can return as many buffered events as the user buffer holds.
        inner.drain_into_queue();
        for out in buf.as_chunks_mut::<{ size_of::<InputEvent>() }>().0 {
            let Some((time, event)) = inner.read_ahead.pop_front() else {
                break;
            };
            let input_event = InputEvent {
                time: KernelTimeval {
                    tv_sec: time.as_secs() as _,
                    tv_usec: time.subsec_micros() as _,
                },
                event_type: event.event_type,
                code: event.code,
                value: event.value as _,
            };
            out.copy_from_slice(input_event.as_bytes());
            read += out.len();
        }
        if read == 0 {
            Err(VfsError::WouldBlock)
        } else {
            Ok(read)
        }
    }

    fn write_at(&self, _buf: &[u8], _offset: u64) -> VfsResult<usize> {
        Err(VfsError::InvalidInput)
    }

    fn flags(&self) -> NodeFlags {
        NodeFlags::NON_CACHEABLE | NodeFlags::STREAM
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_pollable(&self) -> Option<&dyn Pollable> {
        Some(self)
    }

    fn ioctl(&self, current: &crate::task::UserTaskRef, cmd: u32, arg: usize) -> VfsResult<usize> {
        match cmd {
            EVIOCGVERSION => {
                UserPtr::<u32>::from(arg).write(current, 0x10001)?;
                Ok(0)
            }
            EVIOCGID => {
                let device_id = self.inner.lock().device.device_id();
                let user = UserPtr::<InputDeviceId>::from(arg);
                user.write_field(
                    current,
                    offset_of!(InputDeviceId, bus_type),
                    device_id.bus_type,
                )?;
                user.write_field(current, offset_of!(InputDeviceId, vendor), device_id.vendor)?;
                user.write_field(
                    current,
                    offset_of!(InputDeviceId, product),
                    device_id.product,
                )?;
                user.write_field(
                    current,
                    offset_of!(InputDeviceId, version),
                    device_id.version,
                )?;
                Ok(0)
            }
            EVIOCGRAB => Ok(0),
            other => {
                // variable-length command
                let mut tmp = other;
                let nr = (tmp & 0xff) as u8;
                tmp >>= 8;
                let ty = (tmp & 0xff) as u8;
                tmp >>= 8;
                let size = (tmp & 0x3fff) as usize;
                tmp >>= 14;
                let dir = tmp & 0x3;

                if ty != b'E' {
                    warn!("unknown ioctl for evdev: {cmd} {arg}");
                    return Err(VfsError::InvalidInput);
                }

                match dir {
                    // IOC_WRITE
                    1 => return Err(VfsError::InvalidInput),
                    // IOC_READ
                    2 => {
                        #[allow(clippy::single_match)]
                        match nr {
                            // EVIOCGNAME
                            0x06 => {
                                let name = self.inner.lock().device.name().to_string();
                                return return_str(current, arg, size, &name);
                            }
                            // EVIOCGPHYS
                            0x07 => {
                                let location =
                                    self.inner.lock().device.physical_location().to_string();
                                return return_str(current, arg, size, &location);
                            }
                            // EVIOCGUNIQ
                            0x08 => {
                                let unique_id = self.inner.lock().device.unique_id().to_string();
                                return return_str(current, arg, size, &unique_id);
                            }
                            // EVIOCGPROP — device property bitmap. libinput
                            // uses INPUT_PROP_POINTER to keep the cursor
                            // visible on absolute-axis pointing devices like
                            // virtio-tablet; we synthesize the bit at probe
                            // for any non-touchscreen with REL/ABS axes.
                            0x09 => {
                                return write_user_bytes(current, arg, size, &self.prop_bits);
                            }
                            // EVIOCGKEY
                            0x18 => {
                                let key_state = {
                                    let inner = self.inner.lock();
                                    let bytes = inner.key_state.as_bytes();
                                    let mut key_state = Vec::with_capacity(bytes.len());
                                    key_state.extend_from_slice(bytes);
                                    key_state
                                };
                                return write_user_bytes(current, arg, size, &key_state);
                            }
                            // EVIOCGLED
                            0x19 => {
                                return return_zero_bits(
                                    current,
                                    arg,
                                    size,
                                    EventType::Led.bits_count(),
                                );
                            }
                            // EVIOCGSND
                            0x1a => {
                                return return_zero_bits(
                                    current,
                                    arg,
                                    size,
                                    EventType::Sound.bits_count(),
                                );
                            }
                            // EVIOCGSW
                            0x1b => {
                                return return_zero_bits(
                                    current,
                                    arg,
                                    size,
                                    EventType::Switch.bits_count(),
                                );
                            }
                            _ => {}
                        }
                        if nr & !EventType::MAX == EventType::COUNT {
                            return self.get_event_bits(current, arg, size, nr & EventType::MAX);
                        }
                        const ABS_CNT: u8 = 0x40;
                        if nr & !(ABS_CNT - 1) == ABS_CNT {
                            // EVIOCGABS(axis) — absolute axis info.
                            // libinput needs min/max/res to map the
                            // virtio-tablet's 0..0x7FFF absolute range to
                            // screen pixels; without it motion is treated
                            // as noise.
                            if size < size_of::<InputAbsInfo>() {
                                return Err(VfsError::InvalidInput);
                            }
                            let axis = nr & (ABS_CNT - 1);
                            // Linux's evdev returns EINVAL for any axis the
                            // device does not advertise in its EV_ABS bitmap.
                            // virtio-drivers surfaces the same as Error::IoError
                            // (size==0 selector), so without this pre-check
                            // userspace would see EIO and reject the device.
                            if !self.axis_supported(axis) {
                                return Err(VfsError::InvalidInput);
                            }
                            let info = match self.inner.lock().device.get_abs_info(axis) {
                                Ok(info) => info,
                                Err(err) => return Err(input_error_to_vfs_error(err)),
                            };
                            let abs = InputAbsInfo {
                                value: 0,
                                minimum: info.min,
                                maximum: info.max,
                                fuzz: info.fuzz,
                                flat: info.flat,
                                resolution: info.res,
                            };
                            let bytes = abs.as_bytes();
                            UserPtr::<u8>::from(arg).write_slice(current, bytes)?;
                            return Ok(bytes.len());
                        }
                        return Err(VfsError::InvalidInput);
                    }
                    _ => {}
                }

                Err(VfsError::InvalidInput)
            }
        }
    }
}

impl Pollable for EventDev {
    fn poll(&self) -> IoEvents {
        let mut events = IoEvents::empty();
        events.set(IoEvents::IN, self.inner.lock().has_event());
        events
    }

    unsafe fn register_shared(&self, sink: &mut dyn SharedRegistrationSink, events: IoEvents) {
        if !events.contains(IoEvents::IN) {
            return;
        }
        unsafe { sink.register_shared(&self.waiters, IoEvents::IN) };
        if self.inner.lock().has_event() {
            unsafe { self.waiters.wake(IoEvents::IN) };
        }
    }

    unsafe fn register_exclusive(
        &self,
        sink: &mut dyn ExclusiveRegistrationSink,
        events: IoEvents,
    ) {
        if !events.contains(IoEvents::IN) {
            return;
        }
        unsafe { sink.register_exclusive(&self.waiters, IoEvents::IN) };
        if self.inner.lock().has_event() {
            unsafe { self.waiters.wake(IoEvents::IN) };
        }
    }
}

pub fn input_devices(fs: Arc<SimpleFs>) -> DirMapping {
    let mut inputs = DirMapping::new();
    let mut mice_alias: Option<Arc<EventDev>> = None;
    let mut input_id: u32 = 0;
    let input_devices = ax_input::take_inputs();
    for mut device in input_devices.into_iter() {
        let mut keys = [0; 0x300usize.div_ceil(8)];
        assert!(device.get_event_bits(EventType::Key, &mut keys).unwrap());

        const BTN_MOUSE: usize = 0x110;
        let is_mouse = keys[BTN_MOUSE / 8] & (1 << (BTN_MOUSE % 8)) != 0;

        let event_dev = Arc::new(EventDev::new(device));
        event_dev.register_irq();
        let dev = Device::new(
            fs.clone(),
            NodeType::CharacterDevice,
            DeviceId::new(13, 64 + input_id),
            event_dev.clone(),
        );
        inputs.add(format!("event{input_id}"), dev);
        input_id += 1;

        if is_mouse && mice_alias.is_none() {
            mice_alias = Some(event_dev);
        }
    }

    if let Some(event_dev) = mice_alias {
        inputs.add(
            "mice",
            Device::new(
                fs,
                NodeType::CharacterDevice,
                DeviceId::new(13, 63),
                event_dev,
            ),
        );
    }

    EVENT_DEVICE_COUNT.store(input_id, Ordering::Release);
    inputs
}