//! library/device/pci/pci.zig — a device driver's view of the ONE PCI function it has
//! claimed. Config space is mapped as resource 0 (a full 4 KiB ECAM page); this gives
//! header-field accessors, BAR decode + map, capability walks (legacy and extended),
//! MSI/MSI-X programming, power-state handling, and function-level reset, so a driver
//! never re-derives the config-space layout by hand.
//!
//! This is the *device-owned* view: read my own function's live config, map my own BARs.
//! The bus enumerator's view — probing arbitrary, not-yet-claimed functions and sizing
//! their BARs — is a different mechanism and lives in the pci-bus driver. The pure
//! config-space layout both need (offsets, BAR bit fields) is named once in the `pci-class`
//! data module; this logic module adds the parts that need `mmio` + the `driver` client.

const std = @import("std");
const mmio = @import("mmio");
const pci_class = @import("pci-class");
const device = @import("driver");
const time = @import("time");

/// Spec recovery time after a D3hot -> D0 transition.
const d0_recovery_millis: u64 = 10;
/// How long to wait for in-flight transactions to drain before a function-level reset
/// (then reset anyway — resetting a stuck function is the point of FLR).
const flr_pending_timeout_millis: u64 = 100;
/// The spec's maximum FLR completion time.
const flr_settle_millis: u64 = 100;
/// How long to wait for the function to become readable again after an FLR.
const flr_ready_timeout_millis: u64 = 1000;
const flr_poll_interval_millis: u64 = 10;

/// A claimed PCI function whose configuration space is mapped (resource 0). `descriptor`
/// must outlive the Function — the driver's `device.enumerate` buffer does, for the whole
/// bring-up. Header reads and the capability walk hit live config space; `mapBar` caches.
pub const Function = struct {
    device_id: u64,
    descriptor: *const device.DeviceDescriptor,
    config: usize, // virtual base of mapped resource 0
    bar_virtual: [6]usize = .{ 0, 0, 0, 0, 0, 0 }, // per-BAR mmio_map cache

    /// Map config space (resource 0) of the already-claimed `device_id`. null if the map
    /// fails (not claimed, or no config resource).
    pub fn map(device_id: u64, descriptor: *const device.DeviceDescriptor) ?Function {
        const base = device.mmioMap(device_id, 0) orelse return null;
        return .{ .device_id = device_id, .descriptor = descriptor, .config = base };
    }

    pub fn vendorId(self: *const Function) u16 {
        return mmio.readRegister(u16, self.config + pci_class.config_vendor_id);
    }
    pub fn deviceId(self: *const Function) u16 {
        return mmio.readRegister(u16, self.config + pci_class.config_device_id);
    }
    pub fn command(self: *const Function) u16 {
        return mmio.readRegister(u16, self.config + pci_class.config_command);
    }
    pub fn status(self: *const Function) u16 {
        return mmio.readRegister(u16, self.config + pci_class.config_status);
    }
    pub fn revisionId(self: *const Function) u8 {
        return mmio.readRegister(u8, self.config + pci_class.config_revision_id);
    }
    /// Subsystem vendor ID (config 0x2C) — with `subsystemId`, the standard key for
    /// board-level quirk matching.
    pub fn subsystemVendorId(self: *const Function) u16 {
        return mmio.readRegister(u16, self.config + pci_class.config_subsystem_vendor_id);
    }
    pub fn subsystemId(self: *const Function) u16 {
        return mmio.readRegister(u16, self.config + pci_class.config_subsystem_id);
    }
    /// Interrupt pin (config 0x3D): 0 = none, 1..4 = INTA..INTD.
    pub fn interruptPin(self: *const Function) u8 {
        return mmio.readRegister(u8, self.config + pci_class.config_interrupt_pin);
    }
    /// The live class-code triple (config 0x09..0x0B), same shape discovery records.
    pub fn classCode(self: *const Function) pci_class.ClassCode {
        return .{
            .prog_if = mmio.readRegister(u8, self.config + pci_class.config_class_code),
            .subclass = mmio.readRegister(u8, self.config + pci_class.config_class_code + 1),
            .base = mmio.readRegister(u8, self.config + pci_class.config_class_code + 2),
        };
    }

    fn commandSetBits(self: *const Function, bits: u16) void {
        const at = self.config + pci_class.config_command;
        mmio.writeRegister(u16, at, mmio.readRegister(u16, at) | bits);
    }
    fn commandClearBits(self: *const Function, bits: u16) void {
        const at = self.config + pci_class.config_command;
        mmio.writeRegister(u16, at, mmio.readRegister(u16, at) & ~bits);
    }

    /// Set Memory-Space + Bus-Master Enable in the command register. Firmware only enables
    /// memory decode on devices it used at boot; any other device has dead BARs until its
    /// driver sets it. Bus mastering is separately required for the device to do DMA.
    pub fn enableMemoryAndBusMaster(self: *const Function) void {
        self.commandSetBits(pci_class.command_memory_and_bus_master);
    }

    /// Clear Bus-Master Enable — stop the device initiating DMA. The quiesce half of a
    /// driver's shutdown (or a supervisor restart): after this the device can no longer
    /// write memory the process is about to stop owning.
    pub fn disableBusMaster(self: *const Function) void {
        self.commandClearBits(pci_class.command_bus_master);
    }

    /// Set command bit 10: suppress legacy INTx assertion. MSI/MSI-X are unaffected —
    /// set this when enabling either, so the device cannot also raise the shared pin.
    pub fn setInterruptDisable(self: *const Function) void {
        self.commandSetBits(pci_class.command_interrupt_disable);
    }
    /// Clear command bit 10, re-allowing legacy INTx assertion.
    pub fn clearInterruptDisable(self: *const Function) void {
        self.commandClearBits(pci_class.command_interrupt_disable);
    }

    /// Decode BAR `bar` (0..5) and map it: read the BAR register, reject I/O-space BARs,
    /// combine the high dword for a 64-bit BAR, mask the base, then correlate that physical
    /// base with one of the descriptor's memory resources and `mmio_map` it — a BAR names a
    /// *number*, while `mmio_map` takes a *resource index*, and gaps/config-space shift the
    /// numbering. Cached per BAR. null if the BAR is I/O-space or is not a mapped resource.
    pub fn mapBar(self: *Function, bar: u8) ?usize {
        if (bar >= 6) return null;
        if (self.bar_virtual[bar] != 0) return self.bar_virtual[bar];

        const low = mmio.readRegister(u32, self.config + pci_class.config_bar0 + @as(usize, bar) * 4);
        if (low & pci_class.bar_io_space != 0) return null; // an I/O-space BAR
        var base: u64 = low & pci_class.bar_memory_base_mask;
        if ((low & pci_class.bar_type_mask) == pci_class.bar_type_64bit) { // 64-bit: high half is the next dword
            const high = mmio.readRegister(u32, self.config + pci_class.config_bar0 + (@as(usize, bar) + 1) * 4);
            base |= @as(u64, high) << 32;
        }

        for (self.descriptor.resources[0..@intCast(self.descriptor.resource_count)], 0..) |resource, index| {
            if (resource.kind == @intFromEnum(device.ResourceKind.memory) and resource.start == base) {
                const v = device.mmioMap(self.device_id, index) orelse return null;
                self.bar_virtual[bar] = v;
                return v;
            }
        }
        return null;
    }

    /// Iterate the capability list. Empty when the function advertises none.
    pub fn capabilities(self: *const Function) CapabilityIterator {
        const present = self.status() & pci_class.status_capabilities_list != 0;
        const first = if (present)
            mmio.readRegister(u8, self.config + pci_class.config_capabilities_pointer) & pci_class.capability_pointer_mask
        else
            0;
        return .{ .config = self.config, .cursor = first };
    }

    /// First capability with `id`, or null.
    pub fn findCapability(self: *const Function, id: pci_class.CapabilityId) ?Capability {
        var walk = self.capabilities();
        while (walk.next()) |capability| {
            if (capability.id == @intFromEnum(id)) return capability;
        }
        return null;
    }

    /// Program the MSI capability with the kernel's `msi_bind` result and enable it —
    /// one vector (multiple-message-enable 0, matching the kernel's single-vector
    /// grant), INTx suppressed. false if the function has no MSI capability.
    pub fn programMsi(self: *const Function, message: device.Msi) bool {
        const cap = self.findCapability(.msi) orelse return false;
        const control_at = cap.offset + pci_class.msi.control;
        const control = mmio.readRegister(u16, control_at);
        // Program the registers while the capability is disabled.
        mmio.writeRegister(u16, control_at, control & ~pci_class.msi.control_enable);
        mmio.writeRegister(u32, cap.offset + pci_class.msi.address, @truncate(message.address));
        const data_offset = if (control & pci_class.msi.control_64bit_capable != 0) offset: {
            mmio.writeRegister(u32, cap.offset + pci_class.msi.address_high, @intCast(message.address >> 32));
            break :offset pci_class.msi.data_64;
        } else pci_class.msi.data_32;
        // Message data is a 16-bit register in both layouts.
        mmio.writeRegister(u16, cap.offset + data_offset, @truncate(message.data));
        mmio.writeRegister(u16, control_at, (control & ~pci_class.msi.control_multiple_message_enable_mask) | pci_class.msi.control_enable);
        self.setInterruptDisable();
        return true;
    }

    /// Clear the MSI enable bit. No-op if the function has no MSI capability.
    pub fn disableMsi(self: *const Function) void {
        const cap = self.findCapability(.msi) orelse return;
        const control_at = cap.offset + pci_class.msi.control;
        mmio.writeRegister(u16, control_at, mmio.readRegister(u16, control_at) & ~pci_class.msi.control_enable);
    }

    /// The function's MSI-X capability with its vector table mapped: the table's BIR is
    /// resolved through `mapBar` (a free cache hit when it is a BAR the driver already
    /// mapped). null if the capability is absent or the table's BAR cannot be mapped.
    pub fn msix(self: *Function) ?MsiX {
        const cap = self.findCapability(.msix) orelse return null;
        const control = mmio.readRegister(u16, cap.offset + pci_class.msix.control);
        const word = mmio.readRegister(u32, cap.offset + pci_class.msix.table_offset_word);
        const location = pci_class.msix.tableLocation(word);
        const bar_base = self.mapBar(location.bar) orelse return null;
        return .{
            .capability = cap.offset,
            .table = bar_base + location.offset,
            .entry_count = pci_class.msix.tableSize(control),
        };
    }

    /// Bring the function to D0. Firmware can leave a non-boot device in D3hot, where
    /// its BARs and MSI registers do not decode; call this before touching either. No
    /// power-management capability means the function is always at D0: nothing to do.
    /// Preserves PME-Enable and never clears the write-1-to-clear PME-Status bit.
    pub fn ensurePowerStateD0(self: *const Function) void {
        const cap = self.findCapability(.power_management) orelse return;
        const at = cap.offset + pci_class.power_management.control_status;
        const pmcsr = mmio.readRegister(u16, at);
        if (pmcsr & pci_class.power_management.control_status_power_state_mask == pci_class.power_management.power_state_d0) return;
        // PME-Status is RW1C: echoing a read 1 back would clear it, so write it as 0.
        mmio.writeRegister(u16, at, (pmcsr & ~pci_class.power_management.control_status_power_state_mask & ~pci_class.power_management.control_status_pme_status) | pci_class.power_management.power_state_d0);
        time.sleepMillis(d0_recovery_millis);
    }

    /// Function Level Reset via the PCI Express capability: return the hardware to a
    /// known state (a supervisor re-claiming a device after its driver died, or a driver
    /// recovering a wedged function). The six BAR dwords are saved and restored — FLR
    /// clears them, and the bus enumerator's assignment must survive for the descriptor
    /// correlation and `mapBar` cache to stay valid. Everything else is reset: command
    /// enables and MSI/MSI-X programming are gone, so the caller re-runs its whole
    /// bring-up afterwards. false if the function has no PCI Express capability, does
    /// not advertise FLR (conventional-PCI Advanced Features FLR is a possible
    /// follow-up), or never became readable again. Blocks for at least 100 ms.
    pub fn functionLevelReset(self: *const Function) bool {
        const cap = self.findCapability(.pci_express) orelse return false;
        const device_capabilities = mmio.readRegister(u32, cap.offset + pci_class.pci_express.device_capabilities);
        if (device_capabilities & pci_class.pci_express.device_capabilities_flr == 0) return false;

        // Stop new DMA, then give in-flight transactions a bounded chance to drain —
        // and reset anyway on timeout, since resetting a stuck function is the point.
        self.disableBusMaster();
        var waited: u64 = 0;
        while (mmio.readRegister(u16, cap.offset + pci_class.pci_express.device_status) & pci_class.pci_express.device_status_transactions_pending != 0) {
            if (waited >= flr_pending_timeout_millis) break;
            time.sleepMillis(flr_poll_interval_millis);
            waited += flr_poll_interval_millis;
        }

        var bars: [6]u32 = undefined;
        for (&bars, 0..) |*bar, index| bar.* = mmio.readRegister(u32, self.config + pci_class.config_bar0 + index * 4);

        const control_at = cap.offset + pci_class.pci_express.device_control;
        mmio.writeRegister(u16, control_at, mmio.readRegister(u16, control_at) | pci_class.pci_express.device_control_initiate_flr);
        time.sleepMillis(flr_settle_millis);

        waited = 0;
        while (self.vendorId() == 0xFFFF) {
            if (waited >= flr_ready_timeout_millis) return false;
            time.sleepMillis(flr_poll_interval_millis);
            waited += flr_poll_interval_millis;
        }
        for (bars, 0..) |bar, index| mmio.writeRegister(u32, self.config + pci_class.config_bar0 + index * 4, bar);
        return true;
    }

    /// Iterate the extended (PCI Express) capability list at 0x100.. in the 4 KiB ECAM
    /// page. Empty on a conventional-PCI function (the space reads as all-ones).
    pub fn extendedCapabilities(self: *const Function) ExtendedCapabilityIterator {
        return .{ .config = self.config };
    }

    /// First extended capability with `id`, or null.
    pub fn findExtendedCapability(self: *const Function, id: u16) ?ExtendedCapability {
        var walk = self.extendedCapabilities();
        while (walk.next()) |capability| {
            if (capability.id == id) return capability;
        }
        return null;
    }
};

/// One capability header. `offset` is the ABSOLUTE virtual address of the header, so the
/// caller reads its body with `mmio.readRegister(T, cap.offset + n)`.
pub const Capability = struct { id: u8, offset: usize };

pub const CapabilityIterator = struct {
    config: usize,
    cursor: u8,
    guard: u32 = 0, // bounds a malformed/looping chain (48 = the 256-byte space in dwords)

    pub fn next(self: *CapabilityIterator) ?Capability {
        if (self.cursor == 0 or self.guard >= 48) return null;
        self.guard += 1;
        const at = self.config + self.cursor;
        const id = mmio.readRegister(u8, at + 0);
        self.cursor = mmio.readRegister(u8, at + 1) & pci_class.capability_pointer_mask;
        return .{ .id = id, .offset = at };
    }
};

/// A resolved MSI-X capability from `Function.msix`: `capability` is the absolute
/// virtual address of the config-space header, `table` of vector-table entry 0 (in BAR
/// space — table writes are MMIO, not config space). Entries reset masked; bring-up
/// order is programEntry per vector, unmaskEntry per used vector, `enable`, then
/// `Function.setInterruptDisable`.
pub const MsiX = struct {
    capability: usize,
    table: usize,
    entry_count: u16,

    /// Write `message` into table entry `entry`, leaving the entry masked (its reset
    /// state) — the spec requires masking while address/data change. false if `entry`
    /// is out of range.
    pub fn programEntry(self: *const MsiX, entry: u16, message: device.Msi) bool {
        if (entry >= self.entry_count) return false;
        const at = self.table + @as(usize, entry) * pci_class.msix.entry_size;
        mmio.writeRegister(u32, at + pci_class.msix.entry_vector_control, pci_class.msix.entry_vector_control_masked);
        mmio.writeRegister(u32, at + pci_class.msix.entry_address, @truncate(message.address));
        mmio.writeRegister(u32, at + pci_class.msix.entry_address_high, @intCast(message.address >> 32));
        mmio.writeRegister(u32, at + pci_class.msix.entry_data, message.data);
        return true;
    }

    /// Set the entry's vector-control mask bit — its interrupt is held off (pended in
    /// the PBA, not lost). false if `entry` is out of range.
    pub fn maskEntry(self: *const MsiX, entry: u16) bool {
        return self.writeEntryMask(entry, true);
    }
    /// Clear the entry's vector-control mask bit. false if `entry` is out of range.
    pub fn unmaskEntry(self: *const MsiX, entry: u16) bool {
        return self.writeEntryMask(entry, false);
    }
    fn writeEntryMask(self: *const MsiX, entry: u16, masked: bool) bool {
        if (entry >= self.entry_count) return false;
        const at = self.table + @as(usize, entry) * pci_class.msix.entry_size + pci_class.msix.entry_vector_control;
        const control = mmio.readRegister(u32, at);
        mmio.writeRegister(u32, at, if (masked)
            control | pci_class.msix.entry_vector_control_masked
        else
            control & ~pci_class.msix.entry_vector_control_masked);
        return true;
    }

    /// Set the function-mask control bit: every vector masked regardless of entry bits.
    pub fn setFunctionMask(self: *const MsiX) void {
        self.writeControl(pci_class.msix.control_function_mask, true);
    }
    /// Clear the function-mask control bit.
    pub fn clearFunctionMask(self: *const MsiX) void {
        self.writeControl(pci_class.msix.control_function_mask, false);
    }
    /// Set MSI-X Enable. The caller also calls `Function.setInterruptDisable` (INTx off).
    pub fn enable(self: *const MsiX) void {
        self.writeControl(pci_class.msix.control_enable, true);
    }
    /// Clear MSI-X Enable.
    pub fn disable(self: *const MsiX) void {
        self.writeControl(pci_class.msix.control_enable, false);
    }
    fn writeControl(self: *const MsiX, bit: u16, set: bool) void {
        const at = self.capability + pci_class.msix.control;
        const control = mmio.readRegister(u16, at);
        mmio.writeRegister(u16, at, if (set) control | bit else control & ~bit);
    }
};

/// One extended capability. `offset` is the ABSOLUTE virtual address of its header,
/// like `Capability.offset`.
pub const ExtendedCapability = struct { id: u16, version: u4, offset: usize };

pub const ExtendedCapabilityIterator = struct {
    config: usize,
    cursor: u16 = @intCast(pci_class.extended_capability_start),
    guard: u32 = 0, // bounds a malformed chain (480 = the 0xF00-byte space / 8-byte minimum spacing)

    pub fn next(self: *ExtendedCapabilityIterator) ?ExtendedCapability {
        if (self.cursor == 0 or self.guard >= 480) return null;
        self.guard += 1;
        const at = self.config + self.cursor;
        const header = pci_class.ExtendedCapabilityHeader.decode(mmio.readRegister(u32, at));
        // Id 0 marks an empty list; all-ones is a conventional-PCI function (no
        // extended space — reads come back as FFs).
        if (header.id == 0 or header.id == 0xFFFF) return null;
        // A next pointer below 0x100 would walk into the legacy header; treat it as the
        // terminator it must be (0 is the normal one). The 0xFFC decode mask already
        // keeps `config + cursor + 4` inside the 4 KiB page.
        self.cursor = if (header.next >= pci_class.extended_capability_start) header.next else 0;
        return .{ .id = header.id, .version = header.version, .offset = at };
    }
};