// 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-Net Driver (Sovereign Edition)
//!
//! Implements a zero-copy network interface using ION slabs for RX/TX.
//! Supports both Legacy and Modern VirtIO via the universal transport layer.
//!
//! SAFETY: Uses volatile pointers for hardware rings and memory barriers (fences)
//! to ensure correct synchronization with the virtual device.

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

// VirtIO Feature Bits
const VIRTIO_F_VERSION_1 = 32;
const VIRTIO_NET_F_MAC = 5;
const VIRTIO_NET_F_MRG_RXBUF = 15;

// Status Bits
const VIRTIO_CONFIG_S_ACKNOWLEDGE = 1;
const VIRTIO_CONFIG_S_DRIVER = 2;
const VIRTIO_CONFIG_S_DRIVER_OK = 4;
const VIRTIO_CONFIG_S_FEATURES_OK = 8;

// 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_shared(out_id: *u16) u64;
extern fn ion_free_raw(id: u16) void;
extern fn ion_ingress(id: u16, len: u16, offset: 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;

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

    // Periodic debug: show queue state
    if (poll_count == 1 or (poll_count % 100000 == 0)) {
        if (global_driver) |*d| {
            if (d.rx_queue) |q| {
                const hw_idx = q.used.idx;
                const drv_idx = q.index;
                uart.print("[VirtIO] Poll #");
                uart.print_hex(poll_count);
                uart.print(" RX HW:");
                uart.print_hex(hw_idx);
                uart.print(" DRV:");
                uart.print_hex(drv_idx);
                uart.print(" Avail:");
                uart.print_hex(q.avail.idx);
                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 {
    uart.print("[VirtIO] virtio_net_send called, len=");
    uart.print_hex(len);
    uart.print("\n");
    if (global_driver) |*d| {
        d.send_packet(data, len);
    }
}

pub export fn virtio_net_get_mac(out_mac: [*]u8) void {
    if (global_driver) |*d| {
        d.get_mac(out_mac);
    } else {
        // Default fallback if no driver
        out_mac[0] = 0x00;
        out_mac[1] = 0x00;
        out_mac[2] = 0x00;
        out_mac[3] = 0x00;
        out_mac[4] = 0x00;
        out_mac[5] = 0x00;
    }
}

pub export fn rumpk_net_init() void {
    if (VirtioNetDriver.probe()) |_| {
        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 get_mac(self: *VirtioNetDriver, out: [*]u8) void {
        uart.print("[VirtIO-Net] Reading MAC from device_cfg...\n");
        if (self.transport.is_modern) {
            // Use device_cfg directly - this is the VirtIO-Net specific config
            if (self.transport.device_cfg) |cfg| {
                const ptr: [*]volatile u8 = @ptrCast(cfg);
                uart.print("  DeviceCfg at: ");
                uart.print_hex(@intFromPtr(cfg));
                uart.print("\n  MAC bytes: ");

                for (0..6) |i| {
                    out[i] = ptr[i];
                    uart.print_hex8(ptr[i]);
                    if (i < 5) uart.print(":");
                }
                uart.print("\n");
            } else {
                uart.print("  ERROR: device_cfg is null!\n");
                // Fallback to zeros
                for (0..6) |i| {
                    out[i] = 0;
                }
            }
        } else {
            // Legacy
            // Device Config starts at offset 20.
            const base = self.transport.legacy_bar + 20;
            for (0..6) |i| {
                out[i] = @as(*volatile u8, @ptrFromInt(base + i)).*;
            }
        }
    }

    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 func: u8 = 0;

        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.*;

            uart.print("[VirtIO-Net] Probing dev ");
            uart.print_hex(i);
            uart.print(", ID: ");
            uart.print_hex(id);
            uart.print("\n");

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

        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(VIRTIO_CONFIG_S_ACKNOWLEDGE);
        self.transport.add_status(VIRTIO_CONFIG_S_DRIVER);

        // 4. Feature Negotiation
        if (self.transport.is_modern) {
            uart.print("[VirtIO] Starting feature negotiation...\n");

            if (self.transport.common_cfg == null) {
                uart.print("[VirtIO] ERROR: common_cfg is null!\n");
                return false;
            }

            const cfg = self.transport.common_cfg.?;
            uart.print("[VirtIO] common_cfg addr: ");
            uart.print_hex(@intFromPtr(cfg));
            uart.print("\n");

            uart.print("[VirtIO] Reading device features...\n");
            // Read Device Features (Page 0)
            cfg.device_feature_select = 0;
            asm volatile ("fence" ::: .{ .memory = true });
            const f_low = cfg.device_feature;

            // Read Device Features (Page 1)
            cfg.device_feature_select = 1;
            asm volatile ("fence" ::: .{ .memory = true });
            const f_high = cfg.device_feature;

            uart.print("[VirtIO] Device Features: ");
            uart.print_hex(f_low);
            uart.print(" ");
            uart.print_hex(f_high);
            uart.print("\n");

            // Accept VERSION_1 (Modern) and MAC
            const accept_low: u32 = (1 << VIRTIO_NET_F_MAC);
            const accept_high: u32 = (1 << (VIRTIO_F_VERSION_1 - 32));

            uart.print("[VirtIO] Writing driver features...\n");
            cfg.driver_feature_select = 0;
            cfg.driver_feature = accept_low;
            asm volatile ("fence" ::: .{ .memory = true });
            cfg.driver_feature_select = 1;
            cfg.driver_feature = accept_high;
            asm volatile ("fence" ::: .{ .memory = true });

            uart.print("[VirtIO] Checking feature negotiation...\n");
            self.transport.add_status(VIRTIO_CONFIG_S_FEATURES_OK);
            asm volatile ("fence" ::: .{ .memory = true });
            if ((self.transport.get_status() & VIRTIO_CONFIG_S_FEATURES_OK) == 0) {
                uart.print("[VirtIO] Feature negotiation failed!\n");
                return false;
            }
            uart.print("[VirtIO] Features accepted.\n");
        }
        // 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);

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

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

        uart.print("  [Queue Setup] Base: ");
        uart.print_hex(aligned_addr);
        uart.print(" Desc: ");
        uart.print_hex(@intFromPtr(q_ptr.desc));
        uart.print(" Avail: ");
        uart.print_hex(@intFromPtr(q_ptr.avail));
        uart.print(" Used: ");
        uart.print_hex(@intFromPtr(q_ptr.used));
        uart.print("\n");

        // 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_shared(&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");

            // Modern VirtIO-net header: 10 bytes (legacy) + 2 bytes (num_buffers) = 12 bytes
            const header_len: u32 = 12;
            if (len > header_len) {
                // Call ION - Pass only the Ethernet Frame (Skip VirtIO Header)
                // ion_ingress receives slab_id which contains full buffer.
                // We need to tell it the offset.
                // Hack: Pass `len - header_len` as the actual Ethernet frame length.
                // The NPL must then offset into the buffer by 10 to get to Ethernet.
                // OR: We adjust here. Let's adjust here by storing offset.
                // Simplest: Pass len directly, NPL will skip first 10 bytes.
                ion_ingress(slab_id, @intCast(len - header_len), @intCast(header_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_shared(&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 virt_addr = ion_get_virt(slab_id);
        @memset(virt_addr[0..10], 0); // Zero out VirtIO Header (Legacy/Modern 10-byte)

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

        uart.print("[VirtIO TX] Packet Data: ");
        for (0..16) |i| {
            if (i < len) {
                uart.print_hex8(data[i]);
                uart.print(" ");
            }
        }
        uart.print("\n");

        var slab_id: u16 = 0;
        const phys = ion_alloc_shared(&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;

        // Modern VirtIO-net header: 10 bytes (legacy) + 2 bytes (num_buffers) = 12 bytes
        const header_len: usize = 12;
        @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);
        uart.print("[VirtIO TX] Queued & Notified Len=");
        uart.print_hex(header_len + copy_len);
        uart.print("\n");
    }

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