//! The PS/2 Controller is located on the mainboard.
//! In the early days the controller was a single chip (Intel 8042).
//! As of today it is part of the Advanced Integrated Peripheral.
//!
//! It shows up in the device discovery as:
//! KBD_ [acpi_device] hid=PNP0303 (PS/2 Keyboard)
//!       - io_port 0x60 len 0x1
//!        - io_port 0x64 len 0x1
//!        - irq 0x1 len 0x1
//!  MOU_ [acpi_device] hid=PNP0F13 (PS/2 Mouse)
//!        - irq 0xc len 0x1
const std = @import("std");
const runtime = @import("runtime");
const acpi_ids = @import("acpi-ids");
const ps2 = @import("ps2-library.zig");
const device = runtime.device;
const ipc = runtime.ipc;

/// Ask the device on `port` what it is, then spawn the matching driver from the
/// initial-ramdisk, handing it the device's HID as argv[1]. The driver is chosen
/// from what the device reports, not from the port number. Returns the identified
/// type so the forwarding loop can route that port's bytes to the driver once it
/// attaches, or null if nothing was spawned.
fn spawnIdentifiedDriver(controller: ps2.Controller, port: ps2.Port) ?ps2.DeviceType {
    const device_type = controller.identifyDevice(port) orelse {
        std.log.info("identify timed out on port {s}", .{@tagName(port)});
        return null;
    };
    const driver_name = device_type.driverName() orelse {
        std.log.info("unrecognized device on port {s}", .{@tagName(port)});
        return null;
    };
    const hid = device_type.hid() orelse "";
    if (runtime.system.spawnWithArguments(driver_name, &.{hid}) != null) {
        std.log.info("port {s} is a {s}, spawned {s}", .{ @tagName(port), hid, driver_name });
        return device_type;
    }
    std.log.info("failed to spawn {s}", .{driver_name});
    return null;
}

/// Resource index of the controller's IRQ (IRQ1) on the PNP0303 descriptor, found
/// the way the ports are found in `Controller.init`.
fn findInterruptResourceIndex(descriptor: device.DeviceDescriptor) ?u64 {
    for (0..descriptor.resource_count) |index| {
        if (descriptor.resources[index].kind == @intFromEnum(device.ResourceKind.irq)) return index;
    }
    return null;
}

/// Forwarding endpoints of the attached child drivers, indexed by `ps2.Port`.
/// Written when a child's `AttachRequest` arrives, read on every forwarded byte.
var port_endpoints = [_]?ipc.Handle{ null, null };

/// Which device type each port identified as, so an attaching child (which knows
/// its type, not its port) can be matched to the right port's byte stream.
var port_device_types = [_]?ps2.DeviceType{ null, null };

/// Handle a child driver's `AttachRequest`: record the endpoint capability it
/// passed as the forwarding target for the port whose device matches its type.
/// Writes an `AttachReply` into `out` and returns its length.
fn handleAttach(message: []const u8, got: ipc.Received, out: []u8) usize {
    const reply = struct {
        fn write(buffer: []u8, status: ps2.AttachStatus) usize {
            const header = ps2.AttachReply{ .status = @intFromEnum(status) };
            @memcpy(buffer[0..@sizeOf(ps2.AttachReply)], std.mem.asBytes(&header));
            return @sizeOf(ps2.AttachReply);
        }
    };

    if (message.len < @sizeOf(ps2.AttachRequest)) return reply.write(out, .invalid_request);
    const request = std.mem.bytesToValue(ps2.AttachRequest, message[0..@sizeOf(ps2.AttachRequest)]);
    const endpoint = got.cap orelse return reply.write(out, .missing_endpoint);

    for (&port_device_types, 0..) |maybe_type, port_index| {
        const device_type = maybe_type orelse continue;
        if (@intFromEnum(device_type) != request.device_type) continue;
        port_endpoints[port_index] = endpoint;
        std.log.info("{s} driver attached", .{@tagName(device_type)});
        return reply.write(out, .ok);
    }
    return reply.write(out, .no_such_device);
}

pub fn main() void {
    const buffer = runtime.allocator().alloc(device.DeviceDescriptor, 64) catch {
        _ = runtime.system.write("/system/drivers/ps2-bus: out of memory\n");
        return;
    };

    var has_two_channels = false;
    var maybe_controller: ?ps2.Controller = null;
    var maybe_interrupt_index: ?u64 = null;
    // The 8042's IO ports (0x60/0x64) are enumerated under the keyboard ACPI node
    // (PNP0303), so we init the controller from that descriptor — but which device
    // is on which port is decided later by identify, not by this HID.
    const maybe_controller_device_descriptor = device.findDeviceDescriptorByHid(buffer, acpi_ids.HardwareId.ps2_keyboard.hid());
    if (maybe_controller_device_descriptor) |controller_device_descriptor| {
        _ = runtime.system.write("/system/drivers/ps2-bus: found PS/2 controller\n");
        _ = runtime.system.write("/system/drivers/ps2-bus: initializing controller\n");

        if (!device.claim(controller_device_descriptor.id)) {
            _ = runtime.system.write("/system/drivers/ps2-bus: unable to claim controller \n");
            return;
        }

        const controller = ps2.Controller.init(controller_device_descriptor) orelse {
            _ = runtime.system.write("/system/drivers/ps2-bus: controller is missing its IO ports\n");
            return;
        };
        maybe_controller = controller;
        maybe_interrupt_index = findInterruptResourceIndex(controller_device_descriptor);

        controller.disablePort(.one);
        controller.disablePort(.two);
        controller.flushOutputBuffer();

        const current = controller.readConfigurationByte() orelse {
            _ = runtime.system.write("/system/drivers/ps2-bus: controller configuration timed out\n");
            return;
        };

        const update = current & ~(ps2.configuration_first_port_interrupt |
            ps2.configuration_second_port_interrupt |
            ps2.configuration_first_port_translation);

        if (controller.writeConfigurationByte(update) == null) {
            _ = runtime.system.write("/system/drivers/ps2-bus: controller configuration timed out\n");
            return;
        }

        if (controller.performSelfTest()) |reply| {
            if (reply != ps2.response_controller_test_passed) {
                _ = runtime.system.write("/system/drivers/ps2-bus: perform controller self test failed\n");
                return;
            }
        } else {
            _ = runtime.system.write("/system/drivers/ps2-bus: controller self test timed out\n");
            return;
        }

        has_two_channels = controller.hasTwoChannels() orelse {
            _ = runtime.system.write("/system/drivers/ps2-bus: controller channels timed out\n");
            return;
        };

        if (has_two_channels) {
            _ = runtime.system.write("/system/drivers/ps2-bus: has two channels\n");
            // keep the bus quiet until we have tested the ports and are ready to use them
            controller.disablePort(.two);
        } else {
            _ = runtime.system.write("/system/drivers/ps2-bus: has one channel\n");
        }

        // interface tests: always test port 1, test port 2 only if it exists
        const port_one_works = (controller.testPort(.one) orelse {
            _ = runtime.system.write("/system/drivers/ps2-bus: port 1 test timed out\n");
            return;
        }) == ps2.response_port_test_passed;

        var port_two_works = false;
        if (has_two_channels) {
            port_two_works = (controller.testPort(.two) orelse {
                _ = runtime.system.write("/system/drivers/ps2-bus: port 2 test timed out\n");
                return;
            }) == ps2.response_port_test_passed;
        }

        if (!port_one_works and !port_two_works) {
            _ = runtime.system.write("/system/drivers/ps2-bus: no usable ports\n");
            return;
        }

        // Enable the working ports. Their interrupts stay off until IRQ1 is bound
        // below — reset and identify use polled reads, which must never race the
        // interrupt-driven drain loop for bytes.
        controller.enablePort(.one);
        if (port_two_works) controller.enablePort(.two);

        // reset each working device; a failing device is logged but does not
        // abort bring-up of the other one
        if (port_one_works) {
            if (controller.resetDevice(.one)) |passed| {
                if (!passed) _ = runtime.system.write("/system/drivers/ps2-bus: port 1 device reset failed\n");
            } else {
                _ = runtime.system.write("/system/drivers/ps2-bus: port 1 device reset timed out\n");
            }
        }
        if (port_two_works) {
            if (controller.resetDevice(.two)) |passed| {
                if (!passed) _ = runtime.system.write("/system/drivers/ps2-bus: port 2 device reset failed\n");
            } else {
                _ = runtime.system.write("/system/drivers/ps2-bus: port 2 device reset timed out\n");
            }
        }

        // Identify the device on each working port and hand it off to the driver
        // that matches what it reported — a port is not assumed to be a keyboard
        // or a mouse by its number.
        if (port_one_works) port_device_types[@intFromEnum(ps2.Port.one)] = spawnIdentifiedDriver(controller, .one);
        if (port_two_works) port_device_types[@intFromEnum(ps2.Port.two)] = spawnIdentifiedDriver(controller, .two);
    } else {
        _ = runtime.system.write("/system/drivers/ps2-bus: no PS/2 controller found\n");
        return;
    }

    const controller = maybe_controller.?;
    const interrupt_index = maybe_interrupt_index orelse {
        _ = runtime.system.write("/system/drivers/ps2-bus: controller is missing its IRQ\n");
        return;
    };

    // The endpoint the child drivers attach to and IRQ1 wakes. Registered under a
    // well-known id so the children can find it, the way input subscribers find
    // the input service.
    const endpoint = ipc.createIpcEndpoint() orelse {
        _ = runtime.system.write("/system/drivers/ps2-bus: no endpoint\n");
        return;
    };
    if (!ipc.register(.ps2_bus, endpoint)) {
        _ = runtime.system.write("/system/drivers/ps2-bus: register failed\n");
        return;
    }

    // From here on, only the interrupt path reads the data port. Drop anything a
    // device sent between enable-scanning and now, bind the IRQs, and only then
    // let the controller raise them — an interrupt with nobody bound is lost.
    controller.drainOutputBuffer();
    if (!device.irqBind(controller.device_id, interrupt_index, endpoint)) {
        _ = runtime.system.write("/system/drivers/ps2-bus: irq_bind failed\n");
        return;
    }

    // Port 2's interrupt (IRQ12) is enumerated on the auxiliary device's own ACPI
    // node (PNP0F13), not on the controller's — so if port 2 carries a device,
    // claim that node too and route its IRQ to the same endpoint. The IRQ belongs
    // to the *port*, whatever device identify found on it.
    var maybe_auxiliary_interrupt: ?struct { device_id: u64, interrupt_index: u64, gsi: u64 } = null;
    if (port_device_types[@intFromEnum(ps2.Port.two)] != null) {
        if (ps2.findMouseDescriptor(buffer)) |descriptor| {
            if (findInterruptResourceIndex(descriptor)) |auxiliary_index| {
                if (device.claim(descriptor.id) and device.irqBind(descriptor.id, auxiliary_index, endpoint)) {
                    maybe_auxiliary_interrupt = .{
                        .device_id = descriptor.id,
                        .interrupt_index = auxiliary_index,
                        .gsi = descriptor.resources[auxiliary_index].start,
                    };
                } else {
                    _ = runtime.system.write("/system/drivers/ps2-bus: auxiliary irq_bind failed\n");
                }
            }
        }
    }

    var configuration = controller.readConfigurationByte() orelse {
        _ = runtime.system.write("/system/drivers/ps2-bus: controller configuration timed out\n");
        return;
    };
    if (port_device_types[@intFromEnum(ps2.Port.one)] != null) configuration |= ps2.Port.one.interruptBit();
    if (maybe_auxiliary_interrupt != null) configuration |= ps2.Port.two.interruptBit();
    _ = controller.writeConfigurationByte(configuration);

    _ = runtime.system.write("/system/drivers/ps2-bus: ok\n");

    // The forwarding loop: an IRQ1 notification drains the output buffer, routing
    // each byte to the attached driver of the port it came from; a client message
    // is a child driver's AttachRequest.
    var reply_buffer: [@sizeOf(ps2.AttachReply)]u8 = undefined;
    var reply_len: usize = 0;
    var receive: [@sizeOf(ps2.AttachRequest)]u8 = undefined;
    while (true) {
        const got = ipc.replyWait(endpoint, reply_buffer[0..reply_len], &receive, null);
        if (got.isNotification()) {
            reply_len = 0;
            if (got.isMessage() or got.isChildExit()) continue; // nothing sends us these
            while (true) {
                const current_status = ps2.status(controller.device_id, controller.status_index);
                if (current_status & ps2.status_output_buffer_full == 0) break;
                const byte = device.ioRead(controller.device_id, controller.data_index, 0, 1) orelse break;
                const port: ps2.Port = if (current_status & ps2.status_auxiliary_output != 0) .two else .one;
                if (port_endpoints[@intFromEnum(port)]) |child| {
                    const forwarded = ps2.ForwardedByte{ .port = @intFromEnum(port), .byte = byte };
                    _ = ipc.send(child, std.mem.asBytes(&forwarded));
                }
                // An unattached port's byte is dropped — e.g. a keystroke before
                // the keyboard driver has attached.
            }
            // Re-arm the line that woke us: the notification badge carries the
            // GSI, and IRQ1 and IRQ12 are acked through different device claims.
            if (maybe_auxiliary_interrupt) |auxiliary| {
                if (got.source() == auxiliary.gsi) {
                    _ = device.irqAck(auxiliary.device_id, auxiliary.interrupt_index);
                } else {
                    _ = device.irqAck(controller.device_id, interrupt_index);
                }
            } else {
                _ = device.irqAck(controller.device_id, interrupt_index);
            }
            continue;
        }
        reply_len = handleAttach(receive[0..got.len], got, &reply_buffer);
    }
}
