// SPDX-License-Identifier: LCL-1.0
// Copyright (c) 2026 Markus Maiwald
// Stewardship: Self Sovereign Society Foundation
//
// This file is part of the Nexus Commonwealth.
// See legal/LICENSE_COMMONWEALTH.md for license terms.

//! Rumpk Layer 0: VirtIO-Block Driver
//!
//! Provides raw sector access for the unikernel storage system.
//!
//! SAFETY: Synchronous driver. Blocks current fiber until QEMU completes
//! the request. Uses bounce-buffers to guarantee alignment.

const std = @import("std");
const builtin = @import("builtin");
const uart = @import("uart.zig");

// Comptime transport switch: PCI on RISC-V, MMIO on ARM64
const transport_mod = if (builtin.cpu.arch == .aarch64)
    @import("virtio_mmio.zig")
else
    @import("virtio_pci.zig");

// External C/Zig stubs
extern fn malloc(size: usize) ?*anyopaque;

var global_blk: ?VirtioBlkDriver = null;

export fn virtio_blk_read(sector: u64, buf: [*]u8) void {
    if (global_blk) |*d| {
        d.read(sector, buf[0..512]) catch {
            uart.print("[VirtIO-Blk] READ ERROR\n");
        };
    }
}

export fn virtio_blk_write(sector: u64, buf: [*]const u8) void {
    if (global_blk) |*d| {
        d.write(sector, buf[0..512]) catch {
            uart.print("[VirtIO-Blk] WRITE ERROR\n");
        };
    }
}

pub fn init() void {
    if (VirtioBlkDriver.probe()) |_| {
        uart.print("[Rumpk L0] Storage initialized (The Ledger).\n");
    } else {
        uart.print("[Rumpk L0] No Storage Device Found.\n");
    }
}

pub const VirtioBlkDriver = struct {
    transport: transport_mod.VirtioTransport,
    v_desc: [*]volatile VirtioDesc,
    v_avail: *volatile VirtioAvail,
    v_used: *volatile VirtioUsed,
    queue_size: u16,

    pub fn probe() ?VirtioBlkDriver {
        if (builtin.cpu.arch == .aarch64) {
            return probe_mmio();
        } else {
            return probe_pci();
        }
    }

    fn probe_mmio() ?VirtioBlkDriver {
        const mmio = @import("virtio_mmio.zig");
        const base = mmio.find_device(2) orelse { // device_id=2 is block
            uart.print("[VirtIO] No VirtIO-Block MMIO device found\n");
            return null;
        };
        uart.print("[VirtIO] Found VirtIO-Block at MMIO 0x");
        uart.print_hex(base);
        uart.print("\n");
        var self = VirtioBlkDriver{
            .transport = transport_mod.VirtioTransport.init(base),
            .v_desc = undefined,
            .v_avail = undefined,
            .v_used = undefined,
            .queue_size = 0,
        };
        if (self.init_device()) {
            return self;
        }
        return null;
    }

    fn probe_pci() ?VirtioBlkDriver {
        const PCI_ECAM_BASE: usize = 0x30000000;
        const bus: u8 = 0;
        const func: u8 = 0;

        // Scan slots 1 to 8
        var i: u8 = 1;
        while (i <= 8) : (i += 1) {
            const addr = PCI_ECAM_BASE | (@as(usize, bus) << 20) | (@as(usize, i) << 15) | (@as(usize, func) << 12);
            const ptr: *volatile u32 = @ptrFromInt(addr);
            const id = ptr.*;

            if (id == 0x10011af4 or id == 0x10421af4) {
                uart.print("[VirtIO] Found VirtIO-Block device at PCI 00:0");
                uart.print_hex(i);
                uart.print(".0\n");
                var self = VirtioBlkDriver{
                    .transport = transport_mod.VirtioTransport.init(addr),
                    .v_desc = undefined,
                    .v_avail = undefined,
                    .v_used = undefined,
                    .queue_size = 0,
                };
                if (self.init_device()) {
                    return self;
                }
            }
        }
        return null;
    }

    pub fn init_device(self: *VirtioBlkDriver) bool {
        if (!self.transport.probe()) return false;

        self.transport.reset();
        self.transport.add_status(1); // ACKNOWLEDGE
        self.transport.add_status(2); // DRIVER

        // Feature negotiation
        const dev_features = self.transport.get_device_features();
        _ = dev_features;
        // Accept no special features for block — just basic operation
        self.transport.set_driver_features(0);
        transport_mod.io_barrier();
        // FEATURES_OK only on modern (v2) transport
        if (self.transport.is_modern) {
            self.transport.add_status(8); // FEATURES_OK
            transport_mod.io_barrier();
        }

        self.transport.select_queue(0);
        const max_count = self.transport.get_queue_size();
        // Cap queue size for memory efficiency
        const MAX_BLK_QUEUE: u16 = 128;
        const count = if (max_count > MAX_BLK_QUEUE) MAX_BLK_QUEUE else max_count;

        // [Desc] [Avail] [Used] (Simplified layout)
        const total = (count * 16) + (6 + count * 2) + 4096 + (6 + count * 8);
        const raw_ptr = malloc(total + 4096) orelse return false;
        const aligned = (@intFromPtr(raw_ptr) + 4095) & ~@as(usize, 4095);

        // Zero out queue memory to ensure clean state
        const byte_ptr: [*]u8 = @ptrFromInt(aligned);
        for (0..total) |i| {
            byte_ptr[i] = 0;
        }

        self.v_desc = @ptrFromInt(aligned);
        self.v_avail = @ptrFromInt(aligned + (count * 16));
        self.v_used = @ptrFromInt(aligned + (count * 16) + (6 + count * 2) + 4096);
        self.queue_size = count;

        // Ensure avail/used rings start clean
        self.v_avail.flags = 0;
        self.v_avail.idx = 0;
        self.v_used.flags = 0;

        if (self.transport.is_modern) {
            self.transport.setup_modern_queue(aligned, aligned + (count * 16), @intFromPtr(self.v_used));
        } else {
            self.transport.set_queue_size(count);
            self.transport.setup_legacy_queue(@intCast(aligned >> 12));
        }

        self.transport.add_status(4); // DRIVER_OK
        global_blk = self.*;

        uart.print("[VirtIO-Blk] Device Ready. Queue Size: ");
        uart.print_hex(count);
        uart.print(" HeaderSize: ");
        uart.print_hex(@sizeOf(VirtioBlkHeader));
        uart.print("\n");

        return true;
    }

    pub fn read(self: *VirtioBlkDriver, sector: u64, buf: []u8) !void {
        const header = VirtioBlkHeader{
            .type = 0, // READ
            .reserved = 0,
            .sector = sector,
        };
        var status: u8 = 0xFF;

        // Simple synchronous request: Use descriptors 0, 1, 2
        // Desc 0: Header (Read-only for device)
        self.v_desc[0].addr = @intFromPtr(&header);
        self.v_desc[0].len = @sizeOf(VirtioBlkHeader);
        self.v_desc[0].flags = 1; // NEXT
        self.v_desc[0].next = 1;

        // Desc 1: Data Buffer (Write-only for device)
        self.v_desc[1].addr = @intFromPtr(buf.ptr);
        self.v_desc[1].len = 512;
        self.v_desc[1].flags = 1 | 2; // NEXT | WRITE
        self.v_desc[1].next = 2;

        // Desc 2: Status Byte (Write-only for device)
        self.v_desc[2].addr = @intFromPtr(&status);
        self.v_desc[2].len = 1;
        self.v_desc[2].flags = 2; // WRITE
        self.v_desc[2].next = 0;

        // Submit to Avail Ring
        const ring = @as([*]volatile u16, @ptrFromInt(@intFromPtr(self.v_avail) + 4));
        ring[self.v_avail.idx % self.queue_size] = 0; // Head of chain
        const expected_used = self.v_used.idx +% 1;
        transport_mod.io_barrier();
        self.v_avail.idx +%= 1;
        transport_mod.io_barrier();

        self.transport.notify(0);

        // Wait for device (Polling — wait until used ring advances)
        var timeout: usize = 0;
        while (self.v_used.idx != expected_used) {
            transport_mod.io_barrier();
            timeout += 1;
            if (timeout > 100_000_000) {
                uart.print("[VirtIO-Blk] READ TIMEOUT! used.idx=");
                uart.print_hex(self.v_used.idx);
                uart.print(" expected=");
                uart.print_hex(expected_used);
                uart.print("\n");
                return error.DiskError;
            }
        }

        if (status != 0) return error.DiskError;
    }

    pub fn write(self: *VirtioBlkDriver, sector: u64, buf: []const u8) !void {
        const header = VirtioBlkHeader{
            .type = 1, // WRITE
            .reserved = 0,
            .sector = sector,
        };
        var status: u8 = 0xFF;

        self.v_desc[3].addr = @intFromPtr(&header);
        self.v_desc[3].len = @sizeOf(VirtioBlkHeader);
        self.v_desc[3].flags = 1;
        self.v_desc[3].next = 4;

        self.v_desc[4].addr = @intFromPtr(buf.ptr);
        self.v_desc[4].len = 512;
        self.v_desc[4].flags = 1; // Note: Write for disk is READ for device
        self.v_desc[4].next = 5;

        self.v_desc[5].addr = @intFromPtr(&status);
        self.v_desc[5].len = 1;
        self.v_desc[5].flags = 2;
        self.v_desc[5].next = 0;

        const ring = @as([*]volatile u16, @ptrFromInt(@intFromPtr(self.v_avail) + 4));
        ring[self.v_avail.idx % self.queue_size] = 3;
        const expected_used = self.v_used.idx +% 1;
        transport_mod.io_barrier();
        self.v_avail.idx +%= 1;
        transport_mod.io_barrier();

        self.transport.notify(0);

        // Wait for device (Polling — wait until used ring advances)
        var timeout: usize = 0;
        while (self.v_used.idx != expected_used) {
            transport_mod.io_barrier();
            timeout += 1;
            if (timeout > 100_000_000) {
                uart.print("[VirtIO-Blk] WRITE TIMEOUT!\n");
                return error.DiskError;
            }
        }

        if (status != 0) return error.DiskError;
    }

    const VirtioDesc = struct {
        addr: u64,
        len: u32,
        flags: u16,
        next: u16,
    };

    const VirtioAvail = extern struct {
        flags: u16,
        idx: u16,
    };

    const VirtioUsed = extern struct {
        flags: u16,
        idx: u16,
    };

    const VirtioBlkHeader = extern struct {
        type: u32,
        reserved: u32,
        sector: u64,
    };
};