// core/rumpk/hal/virtio_pci.zig
// Sovereign VirtIO Transport Layer
// Handles PCI Capability Discovery and Universal Access

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

// PCI Config Offsets
const PCI_COMMAND = 0x04;
const PCI_STATUS = 0x06;
const PCI_CAP_PTR = 0x34;

// Global Allocator for I/O Ports
var next_io_port: u32 = 0x1000;

// VirtIO Capability Types
const VIRTIO_PCI_CAP_COMMON_CFG = 1;
const VIRTIO_PCI_CAP_NOTIFY_CFG = 2;
const VIRTIO_PCI_CAP_ISR_CFG = 3;
const VIRTIO_PCI_CAP_DEVICE_CFG = 4;
const VIRTIO_PCI_CAP_PCI_CFG = 5;

pub const VirtioTransport = struct {
    base_addr: usize, // ECAM Base
    is_modern: bool,

    // Legacy Interface
    legacy_bar: usize,

    // Modern Interface (Mapped Addresses)
    common_cfg: ?*volatile VirtioPciCommonCfg,
    notify_cfg: ?usize, // Base of notification region
    notify_off_multiplier: u32,
    isr_cfg: ?*volatile u8,

    pub fn init(ecam_base: usize) VirtioTransport {
        return .{
            .base_addr = ecam_base,
            .is_modern = false,
            .legacy_bar = 0,
            .common_cfg = null,
            .notify_cfg = null,
            .notify_off_multiplier = 0,
            .isr_cfg = null,
        };
    }

    pub fn probe(self: *VirtioTransport) bool {
        uart.print("[VirtIO-PCI] Probing capabilities...\n");

        // 1. Enable Bus Master & Memory Space & IO Space
        const cmd_ptr: *volatile u16 = @ptrFromInt(self.base_addr + PCI_COMMAND);
        cmd_ptr.* |= 0x07; // IO | MEM | BUS_MASTER

        // 2. Check for Capabilities
        const status_ptr: *volatile u16 = @ptrFromInt(self.base_addr + PCI_STATUS);
        if ((status_ptr.* & 0x10) != 0) {
            // Has Capabilities
            var cap_offset = @as(*volatile u8, @ptrFromInt(self.base_addr + PCI_CAP_PTR)).*;

            while (cap_offset != 0) {
                const cap_addr = self.base_addr + cap_offset;
                const cap_id = @as(*volatile u8, @ptrFromInt(cap_addr)).*;
                const cap_next = @as(*volatile u8, @ptrFromInt(cap_addr + 1)).*;

                if (cap_id == 0x09) { // Vendor Specific (VirtIO)
                    const cap_type = @as(*volatile u8, @ptrFromInt(cap_addr + 3)).*;
                    const bar_idx = @as(*volatile u8, @ptrFromInt(cap_addr + 4)).*;
                    const offset = @as(*volatile u32, @ptrFromInt(cap_addr + 8)).*;

                    // Resolve BAR Address
                    const bar_reg = @as(*volatile u32, @ptrFromInt(self.base_addr + 0x10 + (@as(usize, bar_idx) * 4)));

                    // Basic BAR resolution (Memory only for Modern)
                    // We assume Modern BARs are Memory Mapped
                    const bar_base = bar_reg.* & 0xFFFFFFF0;

                    if (cap_type == VIRTIO_PCI_CAP_COMMON_CFG) {
                        uart.print("[VirtIO-PCI] Found Modern Common Config\n");
                        self.common_cfg = @ptrFromInt(bar_base + offset);
                        self.is_modern = true;
                    }
                    if (cap_type == VIRTIO_PCI_CAP_NOTIFY_CFG) {
                        uart.print("[VirtIO-PCI] Found Modern Notify Config\n");
                        self.notify_cfg = bar_base + offset;
                        self.notify_off_multiplier = @as(*volatile u32, @ptrFromInt(cap_addr + 16)).*;
                    }
                    if (cap_type == VIRTIO_PCI_CAP_ISR_CFG) {
                        self.isr_cfg = @ptrFromInt(bar_base + offset);
                    }
                }
                cap_offset = cap_next;
            }
        }

        if (self.is_modern) {
            uart.print("[VirtIO] Using Modern Interface\n");
            return true;
        }

        // 3. Fallback to Legacy
        uart.print("[VirtIO] Capabilities not found. Falling back to Legacy.\n");
        const bar0_ptr = @as(*volatile u32, @ptrFromInt(self.base_addr + 0x10));
        var bar0 = bar0_ptr.*;

        // DEBUG: Print BAR0 value from PCI config
        uart.print("[VirtIO] BAR0 from PCI config: ");
        uart.print_hex(bar0);
        uart.print("\n");

        // Handle I/O vs Mem
        const is_io = (bar0 & 0x1) == 1;

        // QEMU RISC-V Constants
        const PIO_BASE = 0x03000000;
        const MMIO_BASE = 0x40000000;

        if ((bar0 & 0xFFFFFFF0) == 0) {
            if (is_io) {
                // Assign I/O port address dynamically
                const io_port: u32 = next_io_port | 0x1;
                uart.print("[VirtIO] Legacy I/O BAR unassigned. Assigning ");
                uart.print_hex(next_io_port);
                uart.print("...\n");

                bar0_ptr.* = io_port;
                next_io_port += 0x100; // Increment 256 bytes (Safer alignment)

                // Readback to verify QEMU accepted the assignment
                const readback = bar0_ptr.*;
                uart.print("[VirtIO] BAR0 readback after write: ");
                uart.print_hex(readback);
                uart.print("\n");
                bar0 = readback;
            } else {
                uart.print("[VirtIO] Legacy Mem BAR unassigned. Assigning 0x40000000...\n");
                bar0_ptr.* = MMIO_BASE;
                bar0 = MMIO_BASE;
            }
        }

        if (is_io) {
            // IO Space translation
            self.legacy_bar = PIO_BASE + (bar0 & 0xFFFFFFFC);
        } else {
            self.legacy_bar = bar0 & 0xFFFFFFF0;
        }

        // Re-enable I/O | MEM | BUS_MASTER after BAR assignment
        const cmd_reg: *volatile u16 = @ptrFromInt(self.base_addr + PCI_COMMAND);
        cmd_reg.* |= 0x07;

        return true;
    }

    // Unified Interface
    pub fn reset(self: *VirtioTransport) void {
        self.set_status(0);
    }

    pub fn get_status(self: *VirtioTransport) u8 {
        if (self.is_modern) {
            return self.common_cfg.?.device_status;
        } else {
            return @as(*volatile u8, @ptrFromInt(self.legacy_bar + 0x12)).*;
        }
    }

    pub fn set_status(self: *VirtioTransport, status: u8) void {
        if (self.is_modern) {
            self.common_cfg.?.device_status = status;
        } else {
            @as(*volatile u8, @ptrFromInt(self.legacy_bar + 0x12)).* = status;
        }
    }

    pub fn add_status(self: *VirtioTransport, status: u8) void {
        self.set_status(self.get_status() | status);
    }

    pub fn select_queue(self: *VirtioTransport, idx: u16) void {
        if (self.is_modern) {
            self.common_cfg.?.queue_select = idx;
        } else {
            @as(*volatile u16, @ptrFromInt(self.legacy_bar + 0x0E)).* = idx;
        }
    }

    pub fn get_queue_size(self: *VirtioTransport) u16 {
        if (self.is_modern) {
            return self.common_cfg.?.queue_size;
        } else {
            return @as(*volatile u16, @ptrFromInt(self.legacy_bar + 0x0C)).*;
        }
    }

    pub fn setup_legacy_queue(self: *VirtioTransport, pfn: u32) void {
        // Only for legacy
        @as(*volatile u32, @ptrFromInt(self.legacy_bar + 0x08)).* = pfn;
    }

    pub fn setup_modern_queue(self: *VirtioTransport, desc: u64, avail: u64, used: u64) void {
        if (self.common_cfg) |cfg| {
            cfg.queue_desc = desc;
            cfg.queue_avail = avail;
            cfg.queue_used = used;
            cfg.queue_enable = 1;
        }
    }

    pub fn notify(self: *VirtioTransport, queue_idx: u16) void {
        if (self.is_modern) {
            if (self.common_cfg) |cfg| {
                if (self.notify_cfg) |notify_base| {
                    cfg.queue_select = queue_idx;
                    const offset = @as(usize, cfg.queue_notify_off) * self.notify_off_multiplier;
                    const ptr: *volatile u16 = @ptrFromInt(notify_base + offset);
                    ptr.* = queue_idx;
                }
            }
        } else {
            const notify_ptr: *volatile u16 = @ptrFromInt(self.legacy_bar + 0x10);
            notify_ptr.* = queue_idx;
        }
    }
};

// Modern Config Structure Layout
pub const VirtioPciCommonCfg = extern struct {
    device_feature_select: u32,
    device_feature: u32,
    driver_feature_select: u32,
    driver_feature: u32,
    msix_config: u16,
    num_queues: u16,
    device_status: u8,
    config_generation: u8,

    // Queue configuration
    queue_select: u16,
    queue_size: u16,
    queue_msix_vector: u16,
    queue_enable: u16,
    queue_notify_off: u16,
    queue_desc: u64,
    queue_avail: u64,
    queue_used: u64,
};
