// Rumpk Layer 0: VirtIO-Net Driver (Sovereign Edition)
// - Uses VirtioTransport for PCI Capability Traversal
// - Supports both Legacy (I/O & Memory) and Modern VirtIO

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

// External Nim functions
extern fn net_ingest_packet(data: [*]const u8, len: usize) bool;

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

extern fn ion_alloc_raw(out_id: *u16) u64;
extern fn ion_free_raw(id: u16) void;
extern fn ion_ingress(id: u16, len: u16) void;
extern fn ion_get_virt(id: u16) [*]u8;
extern fn ion_get_phys(id: u16) u64;
extern fn ion_tx_pop(out_id: *u16, out_len: *u16) bool;

var global_driver: ?VirtioNetDriver = null;

var poll_count: u32 = 0;

export fn virtio_net_poll() void {
    poll_count += 1;

    // Periodic debug: show queue state
    if (poll_count == 1 or (poll_count % 1000000 == 0)) {
        if (global_driver) |*d| {
            if (d.rx_queue) |_| {
                asm volatile ("fence" ::: .{ .memory = true });
                uart.print("[VirtIO] Poll #");
                uart.print_hex(poll_count);
                uart.print("\n");
            }
        }
    }

    if (global_driver) |*d| {
        if (d.rx_queue) |q| {
            d.rx_poll(q);
        }
        if (d.tx_queue) |q| {
            d.tx_poll(q);

            // Fast Path Egress
            var budget: usize = 16;
            while (budget > 0) : (budget -= 1) {
                var slab_id: u16 = 0;
                var len: u16 = 0;
                if (ion_tx_pop(&slab_id, &len)) {
                    d.send_slab(slab_id, len);
                } else {
                    break;
                }
            }
        }
    }
}

export fn virtio_net_send(data: [*]const u8, len: usize) void {
    if (global_driver) |*d| {
        d.send_packet(data, len);
    }
}

export fn virtio_net_init() void {
    if (VirtioNetDriver.probe()) |*driver| {
        var d = driver.*;
        if (d.init_device()) {
            uart.print("[Rumpk L0] Networking initialized (Sovereign).\n");
        }
    }
}

pub const VirtioNetDriver = struct {
    transport: pci.VirtioTransport,
    irq: u32,
    rx_queue: ?*Virtqueue = null,
    tx_queue: ?*Virtqueue = null,

    pub fn init(base: usize, irq_num: u32) VirtioNetDriver {
        return .{
            .transport = pci.VirtioTransport.init(base),
            .irq = irq_num,
            .rx_queue = null,
            .tx_queue = null,
        };
    }

    pub fn probe() ?VirtioNetDriver {
        uart.print("[VirtIO] Probing PCI for networking device...\n");
        const PCI_ECAM_BASE: usize = 0x30000000;
        const bus: u8 = 0;
        const dev: u8 = 1;
        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.*;

        if (id == 0x10001af4 or id == 0x10411af4) {
            uart.print("[VirtIO] Found VirtIO-Net device at PCI 00:01.0\n");
            return VirtioNetDriver.init(addr, 33);
        }

        return null;
    }

    pub fn init_device(self: *VirtioNetDriver) bool {
        uart.print("[VirtIO] Initializing device (Sovereign Mode)...\n");

        // 1. Probe Capabilities
        if (!self.transport.probe()) {
            uart.print("[VirtIO] Probe Failed. Aborting.\n");
            return false;
        }

        // 2. Reset Device
        uart.print("[VirtIO] Resetting device...\n");
        self.transport.reset();

        // 3. Acknowledge & Sense Driver
        self.transport.add_status(1); // ACKNOWLEDGE
        self.transport.add_status(2); // DRIVER

        // 4. Feature Negotiation
        // 5. Setup RX Queue (0)
        self.transport.select_queue(0);
        const rx_count = self.transport.get_queue_size();
        uart.print("[VirtIO] RX Queue Size: ");
        uart.print_hex(rx_count);
        uart.print("\n");

        if (rx_count == 0 or rx_count == 0xFFFF) {
            uart.print("[VirtIO] Invalid RX Queue Size. Aborting.\n");
            return false;
        }

        self.rx_queue = self.setup_queue(0, rx_count) catch {
            uart.print("[VirtIO] Failed to setup RX queue\n");
            return false;
        };

        // KICK THE RX QUEUE
        uart.print("[VirtIO] Kicking RX Queue to activate...\n");
        self.transport.notify(0);

        // 6. Setup TX Queue (1)
        self.transport.select_queue(1);
        const tx_count = self.transport.get_queue_size();
        uart.print("[VirtIO] TX Queue Size: ");
        uart.print_hex(tx_count);
        uart.print("\n");

        if (tx_count == 0 or tx_count == 0xFFFF) {
            uart.print("[VirtIO] Invalid TX Queue Size. Aborting.\n");
            return false;
        }

        self.tx_queue = self.setup_queue(1, tx_count) catch {
            uart.print("[VirtIO] Failed to setup TX queue\n");
            return false;
        };

        // 7. Driver OK
        self.transport.add_status(4); // DRIVER_OK
        global_driver = self.*;

        const end_status = self.transport.get_status();
        uart.print("[VirtIO] Initialization Complete. Status: ");
        uart.print_hex(end_status);
        uart.print("\n");

        return true;
    }

    fn setup_queue(self: *VirtioNetDriver, index: u16, count: u16) !*Virtqueue {
        // Layout: [Descriptors (16*N)] [Avail (6 + 2*N)] [Padding] [Used (6 + 8*N)]
        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.index = 0;
        q_ptr.desc = @ptrFromInt(aligned_addr);
        q_ptr.avail = @ptrFromInt(aligned_addr + desc_size);
        q_ptr.used = @ptrFromInt(aligned_addr + used_offset);

        // Allocate ID tracking array
        const ids_size = @as(usize, count) * @sizeOf(u16);
        const ids_ptr = malloc(ids_size) orelse return error.OutOfMemory;
        q_ptr.ids = @ptrCast(@alignCast(ids_ptr));

        // Pre-allocate buffers for descriptors
        const is_rx = (index == 0);
        const avail_ring = get_avail_ring(q_ptr.avail);

        for (0..count) |i| {
            var slab_id: u16 = 0;
            var phys_addr: u64 = 0;

            if (is_rx) {
                // RX: Allocate Initial Slabs
                phys_addr = ion_alloc_raw(&slab_id);
                if (phys_addr == 0) {
                    uart.print("[VirtIO] RX ION Alloc Failed. OOM.\n");
                    return error.OutOfMemory;
                }
            } else {
                // TX: Start empty
                phys_addr = 0;
                slab_id = 0;
            }

            q_ptr.desc[i].addr = phys_addr;
            q_ptr.desc[i].len = 2048; // Slab Size
            q_ptr.desc[i].flags = if (is_rx) @as(u16, 2) else @as(u16, 0); // 2 = VRING_DESC_F_WRITE
            q_ptr.desc[i].next = 0;
            q_ptr.ids[i] = slab_id;

            if (is_rx) {
                avail_ring[i] = @intCast(i);
            }
        }

        // Initialize flags to 0 to avoid event suppression
        q_ptr.avail.flags = 0;
        q_ptr.used.flags = 0;

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

        if (is_rx) {
            q_ptr.avail.idx = count;
            asm volatile ("fence w, w" ::: .{ .memory = true });
        }

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

        uart.print("[VirtIO] Queue setup complete\n");
        return q_ptr;
    }

    pub fn rx_poll(self: *VirtioNetDriver, q: *Virtqueue) void {
        asm volatile ("fence" ::: .{ .memory = true });

        const used = q.used;
        const hw_idx = used.idx;
        const drv_idx = q.index;

        if (hw_idx == drv_idx) {
            return;
        }

        const avail_ring = get_avail_ring(q.avail);
        const used_ring = get_used_ring(used);
        var replenished: bool = false;

        while (q.index != hw_idx) {
            uart.print("[VirtIO RX] Processing Packet...\n");

            const elem = used_ring[q.index % q.num];
            const desc_idx = elem.id;
            const slab_id = q.ids[desc_idx];
            const len = elem.len;

            uart.print("  Desc: ");
            uart.print_hex(@intCast(desc_idx));
            uart.print("  Len: ");
            uart.print_hex(len);
            uart.print("  Slab: ");
            uart.print_hex(slab_id);
            uart.print("\n");

            const header_len: u32 = 10;
            if (len > header_len) {
                // Call ION
                ion_ingress(slab_id, @intCast(len));
            } else {
                uart.print("  [Warn] Packet too short/empty\n");
                ion_free_raw(slab_id);
            }

            // Replenish
            var new_id: u16 = 0;
            const new_phys = ion_alloc_raw(&new_id);
            if (new_phys != 0) {
                q.desc[desc_idx].addr = new_phys;
                q.ids[desc_idx] = new_id;

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

                avail_ring[q.avail.idx % q.num] = @intCast(desc_idx);
                q.avail.idx +%= 1;
                replenished = true;
            } else {
                uart.print("  [Crit] OOM during replenish!\n");
            }

            q.index +%= 1;
        }

        if (replenished) {
            asm volatile ("fence" ::: .{ .memory = true });
            self.transport.notify(0);
        }
    }

    pub fn tx_poll(self: *VirtioNetDriver, q: *Virtqueue) void {
        _ = self;
        asm volatile ("fence" ::: .{ .memory = true });
        const used = q.used;
        const used_idx = used.idx;
        const used_ring = get_used_ring(used);

        while (q.index != used_idx) {
            const elem = used_ring[q.index % q.num];
            const desc_idx = elem.id;
            const slab_id = q.ids[desc_idx];

            // Free the TX slab
            ion_free_raw(slab_id);

            q.index +%= 1;
        }
    }

    pub fn send_slab(self: *VirtioNetDriver, slab_id: u16, len: u16) void {
        // Zero-Copy Transmit
        const q = self.tx_queue orelse return;
        const avail_phase = q.avail.idx;
        const avail_ring = get_avail_ring(q.avail);
        const idx = avail_phase % q.num;

        const phys_addr = ion_get_phys(slab_id);

        const desc = &q.desc[idx];
        desc.addr = phys_addr;
        desc.len = @intCast(len);
        desc.flags = 0; // No NEXT, No WRITE

        q.ids[idx] = slab_id;

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

        self.transport.notify(1);
        uart.print("[VirtIO TX-Slab] Sent ");
        uart.print_hex(len);
        uart.print(" bytes\n");
    }

    pub fn send_packet(self: *VirtioNetDriver, data: [*]const u8, len: usize) void {
        const q = self.tx_queue orelse return;
        const avail_ring = get_avail_ring(q.avail);

        var slab_id: u16 = 0;
        const phys = ion_alloc_raw(&slab_id);
        if (phys == 0) {
            uart.print("[VirtIO] TX OOM\n");
            return;
        }

        const buf_ptr = ion_get_virt(slab_id);

        const idx = q.avail.idx % q.num;
        const desc_idx = idx;

        const desc = &q.desc[desc_idx];
        q.ids[desc_idx] = slab_id;

        const header_len: usize = 10;
        @memset(buf_ptr[0..header_len], 0);

        const copy_len = if (len > 2000) 2000 else len;
        @memcpy(buf_ptr[header_len .. header_len + copy_len], data[0..copy_len]);

        desc.addr = phys;
        desc.len = @intCast(header_len + copy_len);
        desc.flags = 0;

        asm volatile ("fence" ::: .{ .memory = true });
        avail_ring[idx] = @intCast(desc_idx);
        asm volatile ("fence" ::: .{ .memory = true });
        q.avail.idx +%= 1;
        asm volatile ("fence" ::: .{ .memory = true });

        self.transport.notify(1);
    }

    const Virtqueue = struct {
        desc: [*]volatile VirtioDesc,
        avail: *volatile VirtioAvail,
        used: *volatile VirtioUsed,
        ids: [*]u16, // Shadow array to track Slab ID per descriptor
        num: u16,
        index: u16,
    };

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

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

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

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

    inline fn get_used_ring(used: *volatile VirtioUsed) [*]volatile VirtioUsedElem {
        return @ptrFromInt(@intFromPtr(used) + 4);
    }

    const VirtioUsedElem = extern struct {
        id: u32,
        len: u32,
    };
};