// 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 uart = @import("uart.zig");
const pci = @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: pci.VirtioTransport,
    v_desc: [*]volatile VirtioDesc,
    v_avail: *volatile VirtioAvail,
    v_used: *volatile VirtioUsed,
    queue_size: u16,

    pub fn probe() ?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 = pci.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);
        self.transport.add_status(2);

        self.transport.select_queue(0);
        const count = self.transport.get_queue_size();

        // [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);

        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;

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

        self.transport.add_status(4);
        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
        asm volatile ("fence w, w" ::: .{ .memory = true });
        self.v_avail.idx +%= 1;
        asm volatile ("fence w, w" ::: .{ .memory = true });

        self.transport.notify(0);

        // Wait for device (Polling)
        while (self.v_used.idx == 0) {
            asm volatile ("nop");
        }

        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;
        asm volatile ("fence w, w" ::: .{ .memory = true });
        self.v_avail.idx +%= 1;
        asm volatile ("fence w, w" ::: .{ .memory = true });

        self.transport.notify(0);

        while (status == 0xFF) {
            asm volatile ("nop");
        }

        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,
    };
};