//! Intel VT-d backend for the IOMMU core, behind the architecture boundary.
//!
//! Provides the architecture-neutral core (system/kernel/iommu.zig) with the VT-d
//! hardware specifics behind the `Backend` vtable (iommu.zig beside this file): second-level page-table entry bits, the root/context table structure, the
//! translation-enable and invalidation register sequences, and the fault drain. The
//! core owns the domain table and the page-table walk; this file owns the registers.
//!
//! Register model (VT-d spec §10-11): offsets from the DRHD register base. The unit is
//! programmed once at enable (root table + Translation Enable), then touched only for
//! per-device context changes, per-domain invalidations, and fault draining. Interrupt
//! remapping is deliberately left OFF (GCMD.IRE stays 0): with it off, upstream writes
//! to 0xFEE0_0000-0xFEEF_FFFF are treated as interrupt requests and bypass second-level
//! translation, so the existing MSI contract survives unchanged.

const std = @import("std");
const abi = @import("abi");
const boot_handoff = @import("boot-handoff");
const paging = @import("paging.zig");
const iommu = @import("iommu.zig");

const page_size = abi.page_size;

// Register offsets from the unit's base.
const reg_cap = 0x08; // Capability (64)
const reg_ecap = 0x10; // Extended Capability (64)
const reg_gcmd = 0x18; // Global Command (32, write-only)
const reg_gsts = 0x1C; // Global Status (32, read-only)
const reg_rtaddr = 0x20; // Root Table Address (64)
const reg_ccmd = 0x28; // Context Command (64)
const reg_fsts = 0x34; // Fault Status (32)

const gcmd_te: u32 = 1 << 31; // Translation Enable
const gcmd_srtp: u32 = 1 << 30; // Set Root Table Pointer
const gsts_tes: u32 = 1 << 31; // Translation Enable Status
const gsts_rtps: u32 = 1 << 30; // Root Table Pointer Status

const cap_cm: u64 = 1 << 7; // Caching Mode
const cap_sagaw_shift = 8; // Supported Adjusted Guest Address Widths, bits 12:8
const cap_sagaw_39bit: u64 = 1 << 9; // 3-level
const cap_sagaw_48bit: u64 = 1 << 10; // 4-level
const cap_fro_shift = 24; // Fault-Recording Register Offset, bits 33:24 (×16)
const cap_nfr_shift = 40; // Number of Fault Recording regs, bits 47:40 (+1)
const ecap_coherent: u64 = 1 << 0; // hardware snoops CPU caches reading its structures
const ecap_iro_shift = 8; // IOTLB Register Offset, bits 17:8 (×16)

const ccmd_icc: u64 = 1 << 63; // Invalidate Context-Cache
const ccmd_cirg_global: u64 = @as(u64, 1) << 61; // global granularity
const ccmd_cirg_device: u64 = @as(u64, 3) << 61; // device-selective

const iotlb_ivt: u64 = 1 << 63; // Invalidate IOTLB
const iotlb_iirg_global: u64 = @as(u64, 1) << 60;
const iotlb_iirg_domain: u64 = @as(u64, 2) << 60;
const iotlb_dr: u64 = 1 << 49; // drain reads
const iotlb_dw: u64 = 1 << 48; // drain writes

const fsts_ppf: u32 = 1 << 1; // Primary Pending Fault

// Second-level PTE bits.
const slpte_read: u64 = 1 << 0;
const slpte_write: u64 = 1 << 1;
const slpte_page_size: u64 = 1 << 7; // a 2 MiB leaf (== iommu.huge_leaf_bit)

const address_mask: u64 = 0x000F_FFFF_FFFF_F000;

var register_base: usize = 0;
var version: u32 = 0;
var capabilities: u64 = 0;
var extended_capabilities: u64 = 0;
var coherent: bool = true; // ECAP.C — whether clflush is unnecessary
var levels: u8 = 4;
var context_aw: u64 = 2; // context-entry AW field (001=3-level, 010=4-level)
var gcmd_shadow: u32 = 0; // sticky GCMD bits (TE etc.), for the write-only register

var root_table: u64 = 0; // physical base of the 256-entry root table
var context_table: [256]u64 = .{0} ** 256; // per-bus context table physical, 0 = none

var fault_log_budget: u32 = 32; // rate-limit: log this many faults, then just count
var faults_suppressed: u64 = 0;

fn read32(offset: usize) u32 {
    return @as(*const volatile u32, @ptrFromInt(register_base + offset)).*;
}
fn write32(offset: usize, value: u32) void {
    @as(*volatile u32, @ptrFromInt(register_base + offset)).* = value;
}
fn read64(offset: usize) u64 {
    return @as(*const volatile u64, @ptrFromInt(register_base + offset)).*;
}
fn write64(offset: usize, value: u64) void {
    @as(*volatile u64, @ptrFromInt(register_base + offset)).* = value;
}

fn tableAt(physical: u64) [*]volatile u64 {
    return @ptrFromInt(boot_handoff.physicalToVirtual(physical));
}

/// Map the register window, read caps, pick the address width. Returns the vtable, or
/// null when the unit is not live or advertises no address width danos can drive.
pub fn detect(discovery: iommu.Discovery) ?iommu.Backend {
    // Map 16 KiB: FRCD and IOTLB registers can sit past the first page (CAP.FRO /
    // ECAP.IRO are 16-byte-unit offsets).
    register_base = paging.mapMmio(discovery.register_base, 16 * 1024, true);
    // The Version register's low byte is major.minor; reading it back nonzero is
    // the live-mappable-unit sanity check (previously a kernel-test assertion).
    version = read32(0x00);
    if (version == 0) return null;
    capabilities = read64(reg_cap);
    extended_capabilities = read64(reg_ecap);
    coherent = (extended_capabilities & ecap_coherent) != 0;

    const sagaw = capabilities >> cap_sagaw_shift;
    if (sagaw & cap_sagaw_48bit != 0) {
        levels = 4;
        context_aw = 2; // 010b
    } else if (sagaw & cap_sagaw_39bit != 0) {
        levels = 3;
        context_aw = 1; // 001b
    } else {
        return null; // no width we build tables for
    }

    root_table = allocZeroed() orelse return null;

    return iommu.Backend{
        .levels = levels,
        .supports_huge_pages = true,
        .enable = enable,
        .makeLeaf = makeLeaf,
        .makeTable = makeTable,
        .isPresent = isPresent,
        .flushStructure = flushStructure,
        .attach = attach,
        .detach = detach,
        .invalidateDomain = invalidateDomain,
        .faultDrain = faultDrain,
    };
}

/// Program the root table and turn Translation Enable on. The core has already created
/// and populated the RMRR domains (their context entries are live via `attach`), so at
/// this instant every OTHER device's context entry is not-present and will fault — which
/// for stale firmware bus-mastering is the desired evidence, not a bug.
fn enable() void {
    write64(reg_rtaddr, root_table); // legacy mode (bits 11:10 = 00)
    setGlobalCommand(gcmd_srtp);
    spinStatus(gsts_rtps);
    globalInvalidate();
    setGlobalCommand(gcmd_te);
    spinStatus(gsts_tes);
    gcmd_shadow |= gcmd_te;

    iommu.environment.write("/system/kernel: iommu online (Intel VT-d)\n");
    var buffer: [64]u8 = undefined;
    if (std.fmt.bufPrint(&buffer, "  version    : 0x{x}\n", .{version})) |line|
        iommu.environment.write(line)
    else |_| {}
    if (std.fmt.bufPrint(&buffer, "  agaw       : {d} levels\n", .{levels})) |line|
        iommu.environment.write(line)
    else |_| {}
}

// --- Backend vtable ------------------------------------------------------------------

fn makeLeaf(physical: u64, huge: bool) u64 {
    return (physical & address_mask) | slpte_read | slpte_write | (if (huge) slpte_page_size else 0);
}
fn makeTable(table_physical: u64, level: u8) u64 {
    _ = level;
    return (table_physical & address_mask) | slpte_read | slpte_write;
}
fn isPresent(entry: u64) bool {
    return (entry & (slpte_read | slpte_write)) != 0;
}

fn flushStructure(address: usize) void {
    if (coherent) return; // the unit snoops CPU caches; no flush needed (QEMU)
    asm volatile ("clflush (%[p])"
        :
        : [p] "r" (address),
        : .{ .memory = true });
}

fn attach(bdf: u16, domain: u16, page_table_root: u64) void {
    const bus: u8 = @intCast(bdf >> 8);
    const devfn: u8 = @intCast(bdf & 0xFF);

    // Lazily allocate this bus's context table and link it into the root table.
    if (context_table[bus] == 0) {
        const table = allocZeroed() orelse return;
        context_table[bus] = table;
        const root_entry = &tableAt(root_table)[@as(usize, bus) * 2]; // 16-byte entries
        root_entry.* = (table & address_mask) | 1; // present
        flushStructure(@intFromPtr(root_entry));
    }

    const context = tableAt(context_table[bus]);
    const low = &context[@as(usize, devfn) * 2];
    const high = &context[@as(usize, devfn) * 2 + 1];
    high.* = (context_aw & 0x7) | (@as(u64, domain) << 8); // AW + DID
    low.* = (page_table_root & address_mask) | 1; // present, TT=00 (use second-level)
    flushStructure(@intFromPtr(high));
    flushStructure(@intFromPtr(low));

    invalidateContextDevice(bdf, domain);
    invalidateDomain(domain);
}

fn detach(bdf: u16) void {
    const bus: u8 = @intCast(bdf >> 8);
    const devfn: u8 = @intCast(bdf & 0xFF);
    if (context_table[bus] == 0) return;
    const context = tableAt(context_table[bus]);
    context[@as(usize, devfn) * 2] = 0; // not present
    context[@as(usize, devfn) * 2 + 1] = 0;
    flushStructure(@intFromPtr(&context[@as(usize, devfn) * 2]));
    invalidateContextDevice(bdf, 0);
    globalIotlb();
}

fn invalidateDomain(domain: u16) void {
    const iotlb_offset = iotlbOffset();
    write64(iotlb_offset, iotlb_ivt | iotlb_iirg_domain | iotlb_dr | iotlb_dw | (@as(u64, domain) << 32));
    spin64(iotlb_offset, iotlb_ivt);
}

fn faultDrain() usize {
    const fsts = read32(reg_fsts);
    if (fsts & fsts_ppf == 0) return 0;

    const fro = (capabilities >> cap_fro_shift) & 0x3FF;
    const nfr = ((capabilities >> cap_nfr_shift) & 0xFF) + 1;
    const frcd_base = @as(usize, @intCast(fro)) * 16;

    var seen: usize = 0;
    var i: usize = 0;
    while (i < nfr) : (i += 1) {
        const off = frcd_base + i * 16;
        const high = read64(off + 8);
        if (high & (@as(u64, 1) << 63) == 0) continue; // F: no fault recorded here
        const low = read64(off);
        const address = low & ~@as(u64, 0xFFF);
        const source: u16 = @intCast(high & 0xFFFF);
        const reason: u8 = @intCast((high >> 32) & 0xFF);
        const is_read = (high >> 62) & 1; // T: 1 = read request
        logFault(source, address, reason, is_read == 1);
        write64(off + 8, @as(u64, 1) << 63); // RW1C: clear F
        seen += 1;
    }
    write32(reg_fsts, fsts); // clear PPF/PFO (write-1-to-clear)
    return seen;
}

fn logFault(source: u16, address: u64, reason: u8, is_read: bool) void {
    if (fault_log_budget > 0) {
        fault_log_budget -= 1;
        var buffer: [128]u8 = undefined;
        if (std.fmt.bufPrint(&buffer, "DANOS-IOMMU-FAULT: bdf={x:0>2}:{x:0>2}.{d} addr=0x{x} reason=0x{x} write={d}\n", .{
            source >> 8,
            (source >> 3) & 0x1F,
            source & 0x7,
            address,
            reason,
            @intFromBool(!is_read),
        })) |line| iommu.environment.write(line) else |_| {}
        if (fault_log_budget == 0)
            iommu.environment.write("DANOS-IOMMU-FAULT: (further faults suppressed)\n");
    } else {
        faults_suppressed += 1;
    }
}

// --- register helpers ----------------------------------------------------------------

fn setGlobalCommand(one_shot: u32) void {
    // GCMD is write-only: every write must carry the full sticky state plus the one-shot
    // bit being requested, or a set sticky bit (TE) would be cleared as a side effect.
    write32(reg_gcmd, gcmd_shadow | one_shot);
}

fn spinStatus(bit: u32) void {
    var spins: u64 = 0;
    while (read32(reg_gsts) & bit == 0) {
        spins += 1;
        if (spins > 10_000_000) {
            iommu.environment.write("/system/kernel: WARNING VT-d status bit never set — translation may be incomplete\n");
            return;
        }
    }
}

fn spin64(offset: usize, bit: u64) void {
    var spins: u64 = 0;
    while (read64(offset) & bit != 0) {
        spins += 1;
        if (spins > 10_000_000) return;
    }
}

fn globalInvalidate() void {
    write64(reg_ccmd, ccmd_icc | ccmd_cirg_global);
    spin64(reg_ccmd, ccmd_icc);
    globalIotlb();
}

fn globalIotlb() void {
    const iotlb_offset = iotlbOffset();
    write64(iotlb_offset, iotlb_ivt | iotlb_iirg_global | iotlb_dr | iotlb_dw);
    spin64(iotlb_offset, iotlb_ivt);
}

fn invalidateContextDevice(bdf: u16, domain: u16) void {
    write64(reg_ccmd, ccmd_icc | ccmd_cirg_device | (@as(u64, bdf) << 16) | domain);
    spin64(reg_ccmd, ccmd_icc);
}

fn iotlbOffset() usize {
    const iro = (extended_capabilities >> ecap_iro_shift) & 0x3FF;
    return @as(usize, @intCast(iro)) * 16 + 8; // IOTLB register sits at IRO*16 + 8
}

fn allocZeroed() ?u64 {
    const frame = iommu.environment.allocateFrame() orelse return null;
    const table = tableAt(frame);
    var i: usize = 0;
    while (i < 512) : (i += 1) table[i] = 0;
    return frame;
}