// MARKUS MAIWALD (ARCHITECT) | VOXIS FORGE (AI)
// Rumpk Layer 0: VirtIO-Block Driver (The Ledger)
// - Provides persistent storage access (Sector I/O)

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

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_sync(sector, buf);
    } else {
        uart.print("[VirtIO-Blk] Error: Driver not initialized.\n");
    }
}

export fn virtio_blk_write(sector: u64, buf: [*]const u8) void {
    if (global_blk) |*d| {
        d.write_sync(sector, buf);
    } else {
        uart.print("[VirtIO-Blk] Error: Driver not initialized.\n");
    }
}

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

const VIRTIO_BLK_T_IN: u32 = 0;
const VIRTIO_BLK_T_OUT: u32 = 1;
const SECTOR_SIZE: usize = 512;

pub const VirtioBlkDriver = struct {
    transport: pci.VirtioTransport,
    req_queue: ?*Virtqueue = null,

    pub fn init(base: usize) VirtioBlkDriver {
        return .{
            .transport = pci.VirtioTransport.init(base),
        };
    }

    pub fn probe() ?VirtioBlkDriver {
        uart.print("[VirtIO] Probing PCI for Block device...\n");
        const PCI_ECAM_BASE: usize = 0x30000000;
        // Scan a few slots. Usually 00:02.0 if 00:01.0 is Net.
        // Or implement real PCI scan logic later.
        // For now, check slot 2 (dev=2).
        const bus: u8 = 0;
        const dev: u8 = 2; // Assuming second device
        const func: u8 = 0;

        const addr = PCI_ECAM_BASE | (@as(usize, bus) << 20) | (@as(usize, dev) << 15) | (@as(usize, func) << 12);
        const ptr: *volatile u32 = @ptrFromInt(addr);
        const id = ptr.*;

        // Device ID 0x1001 (Legacy Block) or 0x1042 (Modern Block)
        // 0x1042 = 0x1040 + 2
        if (id == 0x10011af4 or id == 0x10421af4) {
            uart.print("[VirtIO] Found VirtIO-Block device at PCI 00:02.0\n");
            return VirtioBlkDriver.init(addr);
        }

        // Try Slot 3 just in case
        const dev3: u8 = 3;
        const addr3 = PCI_ECAM_BASE | (@as(usize, bus) << 20) | (@as(usize, dev3) << 15) | (@as(usize, func) << 12);
        const ptr3: *volatile u32 = @ptrFromInt(addr3);
        const id3 = ptr3.*;
        if (id3 == 0x10011af4 or id3 == 0x10421af4) {
            uart.print("[VirtIO] Found VirtIO-Block device at PCI 00:03.0\n");
            return VirtioBlkDriver.init(addr3);
        }

        return null;
    }

    pub fn init_device(self: *VirtioBlkDriver) bool {
        // 0. Probe Transport (Legacy/Modern)
        if (!self.transport.probe()) {
            uart.print("[VirtIO-Blk] Transport Probe Failed.\n");
            return false;
        }

        // 1. Reset
        self.transport.reset();
        // 2. ACK + DRIVER
        self.transport.add_status(3);

        // 3. Queue Setup (Queue 0 is Request Queue)
        self.transport.select_queue(0);
        const q_size = self.transport.get_queue_size();
        if (q_size == 0) return false;

        self.req_queue = self.setup_queue(0, q_size) catch return false;

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

        uart.print("[VirtIO-Blk] Device Ready. Queue Size: ");
        uart.print_hex(q_size);
        uart.print("\n");
        return true;
    }

    pub fn read_sync(self: *VirtioBlkDriver, sector: u64, buf: [*]u8) void {
        self.submit_request(VIRTIO_BLK_T_IN, sector, buf, 512);
    }

    pub fn write_sync(self: *VirtioBlkDriver, sector: u64, buf: [*]const u8) void {
        // Cast const away because submit_request buffer logic is generic, but T_OUT implies read from buf
        self.submit_request(VIRTIO_BLK_T_OUT, sector, @constCast(buf), 512);
    }

    fn submit_request(self: *VirtioBlkDriver, type_: u32, sector: u64, buf: [*]u8, len: u32) void {
        const q = self.req_queue orelse return;
        const idx = q.avail.idx % q.num;

        // We need 3 descriptors: Header, Buffer, Status
        // For simplicity, we assume we have 3 consecutive descriptors available.
        // A robust driver would allocate from a free list.
        // We will just take 3 linearly? No, 'desc' is a ring.
        // We need to find 3 free descriptors.
        // Simplification: We assume traffic is low and we consume valid indices.
        // Currently 'virtio_net' used a simple 1:1 map.
        // We will just use `idx * 3`, `idx * 3 + 1`, `idx * 3 + 2`?
        // No, `idx` is from avail ring.
        // Let's implement a simple free list or just bump a counter?
        // To be safe, let's just pick head = idx.
        // Wait, standard `idx` tracks the avail ring index, not descriptor index.
        // We can pick descriptor index = idx (modulo q.num/3?).
        // Let's maintain a `next_free_desc` in the driver or Queue?
        // Since this is Sync, we can just use descriptors 0, 1, 2 always?
        // NO. Concurrency issues if called from multiple fibers?
        // Since we are single-threaded (mostly) in ION fiber for now, maybe.
        // But cleaner: use `idx` as base.
        // Descriptor table size = q_size. If q_size=128, we can support 128/3 concurrent requests.
        // Let's use `head_desc = (idx * 3) % q_size`.
        // Ensure q_size is large enough.

        const head = (idx * 3) % q.num;
        const d1 = head;
        const d2 = (head + 1) % q.num;
        const d3 = (head + 2) % q.num;

        // 1. Header
        const req_header = malloc(@sizeOf(VirtioBlkReq)) orelse return;
        const header: *VirtioBlkReq = @ptrCast(@alignCast(req_header));
        header.type = type_;
        header.reserved = 0;
        header.sector = sector;

        // 2. Status
        const status_ptr = malloc(1) orelse return;
        const status: *u8 = @ptrCast(@alignCast(status_ptr));
        status.* = 0xFF; // Init with error

        // Setup Desc 1 (Header)
        q.desc[d1].addr = @intFromPtr(header);
        q.desc[d1].len = @sizeOf(VirtioBlkReq);
        q.desc[d1].flags = 1; // NEXT
        q.desc[d1].next = d2;

        // Setup Desc 2 (Buffer)
        q.desc[d2].addr = @intFromPtr(buf);
        q.desc[d2].len = len;
        // If READ (IN), device writes to buffer -> VRING_DESC_F_WRITE (2)
        // If WRITE (OUT), device reads from buffer -> 0
        q.desc[d2].flags = if (type_ == VIRTIO_BLK_T_IN) @as(u16, 1 | 2) else @as(u16, 1); // NEXT | (WRITE?)
        q.desc[d2].next = d3;

        // Setup Desc 3 (Status)
        q.desc[d3].addr = @intFromPtr(status);
        q.desc[d3].len = 1;
        q.desc[d3].flags = 2; // WRITE (Device writes status)
        q.desc[d3].next = 0;

        asm volatile ("fence" ::: .{ .memory = true });

        // Put head in Avail Ring
        const avail_ring = get_avail_ring(q.avail);
        avail_ring[idx] = d1;

        asm volatile ("fence" ::: .{ .memory = true });
        q.avail.idx +%= 1;
        asm volatile ("fence" ::: .{ .memory = true });

        self.transport.notify(0);

        // Busy Wait for Completion (Sync)
        // We poll Used Ring.
        // We need to track 'last_used_idx'.
        // Simplified: Wait until status changes?
        // No, status write might happen last.
        // Wait for status to be 0 (OK) or 1 (Error).

        // Safety timeout
        var timeout: usize = 10000000;
        while (status.* == 0xFF and timeout > 0) : (timeout -= 1) {
            // asm volatile ("pause");
            // Invalidate cache?
            asm volatile ("fence" ::: .{ .memory = true });
        }

        if (timeout == 0) {
            uart.print("[VirtIO-Blk] Timeout on Sector ");
            uart.print_hex(sector);
            uart.print("\n");
        } else if (status.* != 0) {
            uart.print("[VirtIO-Blk] I/O Error: ");
            uart.print_hex(status.*);
            uart.print("\n");
        }

        // Cleanup?
        // We used malloc, we should free.
        // But implementing 'free' is hard if we don't have it exposed from stubs.
        // 'virtio_net' used 'ion_alloc_raw' and 'ion_free_raw'.
        // Here we simulated malloc.
        // Assumption: malloc usage here is leaky in this MVP unless we implement free.
        // For Phase 10: "The Ledger", leaking 16 bytes per block op is acceptable for a demo,
        // OR we use a static buffer for headers if single threaded.
        // Let's use a static global header buffer since we are sync.
    }

    fn setup_queue(self: *VirtioBlkDriver, index: u16, count: u16) !*Virtqueue {
        // ...(Similar to Net)...
        // Allocate Memory
        const desc_size = 16 * @as(usize, count);
        const avail_size = 6 + 2 * @as(usize, count);
        const used_offset = (desc_size + avail_size + 4095) & ~@as(usize, 4095);
        const used_size = 6 + 8 * @as(usize, count);
        const total_size = used_offset + used_size;

        const raw_ptr = malloc(total_size + 4096) orelse return error.OutOfMemory;
        const aligned_addr = (@intFromPtr(raw_ptr) + 4095) & ~@as(usize, 4095);

        const q_ptr_raw = malloc(@sizeOf(Virtqueue)) orelse return error.OutOfMemory;
        const q_ptr: *Virtqueue = @ptrCast(@alignCast(q_ptr_raw));

        q_ptr.num = count;
        q_ptr.desc = @ptrFromInt(aligned_addr);
        q_ptr.avail = @ptrFromInt(aligned_addr + desc_size);
        q_ptr.used = @ptrFromInt(aligned_addr + used_offset);

        // Notify Device
        const phys_addr = aligned_addr;
        self.transport.select_queue(index);
        if (self.transport.is_modern) {
            self.transport.setup_modern_queue(phys_addr, phys_addr + desc_size, phys_addr + used_offset);
        } else {
            const pfn = @as(u32, @intCast(phys_addr >> 12));
            self.transport.setup_legacy_queue(pfn);
        }

        return q_ptr;
    }

    // structs ...
    const VirtioBlkReq = extern struct {
        type: u32,
        reserved: u32,
        sector: u64,
    };

    const Virtqueue = struct {
        desc: [*]volatile VirtioDesc,
        avail: *volatile VirtioAvail,
        used: *volatile VirtioUsed,
        num: u16,
    };

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

    inline fn get_avail_ring(avail: *volatile VirtioAvail) [*]volatile u16 {
        return @ptrFromInt(@intFromPtr(avail) + 4);
    }
};