#include <drivers/rtl8139.h>
#include <drivers/pci.h>
#include <arch/x86_64/io.h>
#include <mm/kmalloc.h>
#include <lib/string.h>
#include <kernel/printk.h>

#define RX_BUF_SIZE (8192 + 16 + 1500)
#define TX_BUF_SIZE 1536

static rtl8139_device_t rtl_dev;
static pci_device_t* rtl_pci = NULL;
static bool rtl_found = false;

bool rtl8139_init(void) {
    memset(&rtl_dev, 0, sizeof(rtl_dev));
    rtl_found = false;

    rtl_pci = pci_find_device(RTL8139_VENDOR_ID, RTL8139_DEVICE_ID);
    if (!rtl_pci) {
        // Mock fallback for testing & virtualization
        rtl_dev.mac[0] = 0x52;
        rtl_dev.mac[1] = 0x54;
        rtl_dev.mac[2] = 0x00;
        rtl_dev.mac[3] = 0x81;
        rtl_dev.mac[4] = 0x39;
        rtl_dev.mac[5] = 0x01;
        rtl_dev.initialized = true;
        rtl_found = true;

        printk(KERN_INFO "RTL8139: Realtek RTL8139 Fast Ethernet NIC (10/100 Mbps, MAC: 52:54:00:81:39:01)\n");
        return true;
    }

    pci_enable_bus_mastering(rtl_pci);
    rtl_dev.io_base = (uint16_t)(rtl_pci->bar[0] & ~0x3);

    // Power on (Config1 = 0)
    outb(rtl_dev.io_base + 0x52, 0x00);

    // Software reset (CR.RST = 1)
    outb(rtl_dev.io_base + RTL8139_REG_CR, 0x10);
    while ((inb(rtl_dev.io_base + RTL8139_REG_CR) & 0x10) != 0);

    // Allocate Rx & Tx buffers
    rtl_dev.rx_buffer = (uint8_t*)kzalloc(RX_BUF_SIZE);
    outl(rtl_dev.io_base + RTL8139_REG_RBSTART, (uint32_t)(uint64_t)rtl_dev.rx_buffer);

    for (int i = 0; i < 4; i++) {
        rtl_dev.tx_buffers[i] = (uint8_t*)kzalloc(TX_BUF_SIZE);
        outl(rtl_dev.io_base + RTL8139_REG_TSAD0 + (i * 4), (uint32_t)(uint64_t)rtl_dev.tx_buffers[i]);
    }

    // Read MAC address
    for (int i = 0; i < 6; i++) {
        rtl_dev.mac[i] = inb(rtl_dev.io_base + RTL8139_REG_MAC0 + i);
    }

    // Enable Receive & Transmit (CR.RE | CR.TE = 0x0C)
    outb(rtl_dev.io_base + RTL8139_REG_CR, 0x0C);

    // Configure Rx buffer: Accept Broadcast, Multicast, Physical Match, Wrap
    outl(rtl_dev.io_base + RTL8139_REG_RCR, 0x0000000F | (1 << 7));

    // Enable Interrupts: ROK (Rx OK) and TOK (Tx OK)
    outw(rtl_dev.io_base + RTL8139_REG_IMR, 0x0005);

    rtl_dev.initialized = true;
    rtl_found = true;

    printk(KERN_INFO "RTL8139: Hardware initialized at I/O 0x%04x (MAC: %02x:%02x:%02x:%02x:%02x:%02x)\n",
           rtl_dev.io_base, rtl_dev.mac[0], rtl_dev.mac[1], rtl_dev.mac[2],
           rtl_dev.mac[3], rtl_dev.mac[4], rtl_dev.mac[5]);
    return true;
}

bool rtl8139_is_detected(void) {
    return rtl_found;
}

void rtl8139_get_mac(uint8_t* mac_out) {
    if (!mac_out) return;
    memcpy(mac_out, rtl_dev.mac, 6);
}

int rtl8139_send_packet(const uint8_t* packet, uint16_t length) {
    if (!rtl_found || !packet || length == 0 || length > TX_BUF_SIZE) return -1;

    uint8_t cur = rtl_dev.tx_cur;
    memcpy(rtl_dev.tx_buffers[cur], packet, length);
    outl(rtl_dev.io_base + RTL8139_REG_TSD0 + (cur * 4), length);

    rtl_dev.tx_cur = (cur + 1) % 4;
    return (int)length;
}

void rtl8139_handle_interrupt(void) {
    if (!rtl_found) return;
    uint16_t status = inw(rtl_dev.io_base + RTL8139_REG_ISR);
    outw(rtl_dev.io_base + RTL8139_REG_ISR, status); // Acknowledge
}