#include "net.h"
#include "arch.h"
#include "fb.h"
#include "util.h"
#include "printf.h"
#include "dbg.h"

/* ──────────────────────────────────────────────
 * PCI helpers
 * ────────────────────────────────────────────── */
#define PCI_CONFIG_ADDR 0xCF8
#define PCI_CONFIG_DATA 0xCFC

static uint32_t pci_read(uint8_t bus, uint8_t dev, uint8_t func, uint8_t off)
{
    uint32_t addr = (1u << 31) | ((uint32_t)bus << 16) | ((uint32_t)dev << 11) | ((uint32_t)func << 8) | (off & ~3u);
    outl(PCI_CONFIG_ADDR, addr);
    return inl(PCI_CONFIG_DATA);
}

static void pci_write(uint8_t bus, uint8_t dev, uint8_t func, uint8_t off, uint32_t val)
{
    uint32_t addr = (1u << 31) | ((uint32_t)bus << 16) | ((uint32_t)dev << 11) | ((uint32_t)func << 8) | (off & ~3u);
    outl(PCI_CONFIG_ADDR, addr);
    outl(PCI_CONFIG_DATA, val);
}

/* ──────────────────────────────────────────────
 * e1000 state
 * ────────────────────────────────────────────── */
static volatile uint32_t *e1000_mmio;
static uint8_t e1000_bus, e1000_dev;

#define E1000_REG_CTRL     0x0000
#define E1000_REG_STATUS   0x0008
#define E1000_REG_EERD     0x0014
#define E1000_REG_ICR      0x00C0
#define E1000_REG_IMS      0x00D0
#define E1000_REG_RCTL     0x0100
#define E1000_REG_TCTL     0x0400
#define E1000_REG_RDBAL    0x02800
#define E1000_REG_RDBAH    0x02804
#define E1000_REG_RDLEN    0x02808
#define E1000_REG_RDH      0x02810
#define E1000_REG_RDT      0x02818
#define E1000_REG_TDBAL    0x03800
#define E1000_REG_TDBAH    0x03804
#define E1000_REG_TDLEN    0x03808
#define E1000_REG_TDH      0x03810
#define E1000_REG_TDT      0x03818
#define E1000_REG_TIPG     0x00410
#define E1000_REG_MTA      0x5200
#define E1000_REG_RAL      0x5400
#define E1000_REG_RAH      0x5404

static inline void e1000_write(uint32_t reg, uint32_t val) { e1000_mmio[reg / 4] = val; }
static inline uint32_t e1000_read(uint32_t reg) { return e1000_mmio[reg / 4]; }

/* ──────────────────────────────────────────────
 * TX/RX descriptors
 * ────────────────────────────────────────────── */
struct e1000_desc {
    uint64_t addr;
    uint16_t length;
    uint16_t cso_csum; /* TX: bytes 10-11 = [CSO, CMD]; RX: csum */
    uint8_t  status;   /* byte 12: STATUS (DD bit 0) for both TX and RX */
    uint8_t  css_err;  /* byte 13: TX: CSS; RX: errors */
    uint16_t special;  /* bytes 14-15: TX: special; RX: vlan */
} __attribute__((packed));

#define E1000_CMD_EOP   0x01
#define E1000_CMD_IFCS  0x02
#define E1000_CMD_RS    0x08 /* Report Status — required for DD writeback */
#define E1000_STS_DD    0x01

    static struct e1000_desc *tx_descs;
static uint8_t           *tx_bufs[NET_TX_BUFS];
static uint8_t            tx_tail_idx;

static struct e1000_desc *rx_descs;
static uint8_t           *rx_bufs[NET_RX_BUFS];

/* ──────────────────────────────────────────────
 * Network state
 * ────────────────────────────────────────────── */
static uint8_t net_mac[NET_MAC_LEN];
static uint8_t net_ip_addr[NET_IP_LEN];
static uint8_t net_mask_addr[NET_IP_LEN];
static uint8_t net_gw_addr[NET_IP_LEN];

#define ARP_CACHE_SIZE 16
struct arp_entry {
    uint8_t ip[4];
    uint8_t mac[6];
    uint8_t valid;
};
static struct arp_entry arp_cache[ARP_CACHE_SIZE];

static uint8_t tx_buf_area[NET_TX_BUFS * NET_BUF_SIZE] __attribute__((aligned(4096)));
static uint8_t rx_buf_area[NET_RX_BUFS * NET_BUF_SIZE] __attribute__((aligned(4096)));
static struct e1000_desc tx_desc_area[NET_TX_BUFS] __attribute__((aligned(4096)));
static struct e1000_desc rx_desc_area[NET_RX_BUFS] __attribute__((aligned(4096)));

/* ──────────────────────────────────────────────
 * Checksum
 * ────────────────────────────────────────────── */
uint16_t net_checksum(const void *data, uint32_t len)
{
    const uint16_t *p = data;
    uint32_t sum = 0;
    while(len > 1) { sum += *p++; len -= 2; }
    if(len) sum += *(const uint8_t *)p;
    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    return ~sum;
}

static uint16_t ip_hdr_checksum(const void *hdr, uint32_t len)
{
    struct net_ip_hdr *ip = (struct net_ip_hdr *)hdr;
    uint16_t saved = ip->checksum;
    ip->checksum = 0;
    uint16_t c = net_checksum(hdr, len);
    ip->checksum = saved;
    return c;
}

static uint16_t tcp_udp_checksum(const uint8_t *src_ip, const uint8_t *dst_ip,
                                  uint8_t proto, const void *hdr, uint32_t hdr_len)
{
    uint32_t sum = 0;
    const uint16_t *p;

    p = (const uint16_t *)src_ip;
    sum += p[0]; sum += p[1];
    p = (const uint16_t *)dst_ip;
    sum += p[0]; sum += p[1];
    sum += __builtin_bswap16(proto);
    sum += __builtin_bswap16((uint16_t)hdr_len);

    p = hdr;
    while(hdr_len > 1) { sum += *p++; hdr_len -= 2; }
    if(hdr_len) sum += *(const uint8_t *)p;

    while(sum >> 16) sum = (sum & 0xFFFF) + (sum >> 16);
    return ~sum;
}

/* ──────────────────────────────────────────────
 * e1000 init
 * ────────────────────────────────────────────── */
static void e1000_read_mac(void)
{
    uint32_t ral = e1000_read(E1000_REG_RAL);
    uint32_t rah = e1000_read(E1000_REG_RAH);
    net_mac[0] = ral & 0xFF;
    net_mac[1] = (ral >> 8) & 0xFF;
    net_mac[2] = (ral >> 16) & 0xFF;
    net_mac[3] = (ral >> 24) & 0xFF;
    net_mac[4] = rah & 0xFF;
    net_mac[5] = (rah >> 8) & 0xFF;

    if(net_mac[0] == 0 && net_mac[1] == 0 && net_mac[2] == 0 &&
       net_mac[3] == 0 && net_mac[4] == 0 && net_mac[5] == 0) {
        for(int i = 0; i < 3; i++) {
            e1000_write(E1000_REG_EERD, (i << 2) | 1);
            for(volatile int w = 0; w < 10000; w++);
            uint16_t val = (e1000_read(E1000_REG_EERD) >> 16) & 0xFFFF;
            net_mac[i*2]     = val & 0xFF;
            net_mac[i*2 + 1] = (val >> 8) & 0xFF;
        }
    }
}

static void e1000_init_mmio(void)
{
    dbg_printf("[net] e1000 reset...\n");

    /* Software reset — set RST bit (bit 26) */
    e1000_write(E1000_REG_CTRL, 0x04000000);
    /* Wait for reset to complete — RST bit auto-clears */
    for(volatile int i = 0; i < 5000000; i++) {
        if(!(e1000_read(E1000_REG_CTRL) & 0x04000000)) break;
    }
    dbg_printf("[net] ctrl after reset: 0x%x\n", e1000_read(E1000_REG_CTRL));

    /* Set Auto-Speed Detect + Link Up */
    e1000_write(E1000_REG_CTRL, e1000_read(E1000_REG_CTRL) | 0x20 | 0x40);

    /* Disable interrupts */
    e1000_write(E1000_REG_IMS, 0);

    /* Clear multicast table */
    for(int i = 0; i < 128; i++)
        e1000_write(E1000_REG_MTA + i*4, 0);

    e1000_read_mac();
    dbg_printf("[net] MAC read OK\n");

    /* Write MAC address to RA[0] with AV bit */
    e1000_write(E1000_REG_RAL, net_mac[0] | (net_mac[1] << 8) | (net_mac[2] << 16) | (net_mac[3] << 24));
    e1000_write(E1000_REG_RAH, (net_mac[4] | (net_mac[5] << 8)) | (1u << 31));
    dbg_printf("[net] RA[0] written with AV\n");

    /* RX ring */
    uint32_t rx_phys = (uint32_t)(uintptr_t)rx_descs;
    e1000_write(E1000_REG_RDBAL, rx_phys);
    e1000_write(E1000_REG_RDBAH, 0);
    e1000_write(E1000_REG_RDLEN, NET_RX_BUFS * sizeof(struct e1000_desc));
    e1000_write(E1000_REG_RDH, 0);
    e1000_write(E1000_REG_RDT, NET_RX_BUFS - 1);
    /* EN=1, BSIZE=00 (2048), SECRC=1 (strip CRC) */
    /* EN=1 (bit1), BAM=1 (bit15, accept broadcasts), SZ_2048=0, SECRC=1 (bit26, strip CRC) — matches xv6 reference */
    e1000_write(E1000_REG_RCTL, (1 << 1) | (1 << 15) | (1 << 26));
    dbg_printf("[net] RX ring: RDBAL=0x%x RDH=%d RDT=%d RCTL=0x%x\n",
               rx_phys, e1000_read(E1000_REG_RDH),
               e1000_read(E1000_REG_RDT), e1000_read(E1000_REG_RCTL));

    /* TX ring */
    uint32_t tx_phys = (uint32_t)(uintptr_t)tx_descs;
    e1000_write(E1000_REG_TDBAL, tx_phys);
    e1000_write(E1000_REG_TDBAH, 0);
    e1000_write(E1000_REG_TDLEN, NET_TX_BUFS * sizeof(struct e1000_desc));
    e1000_write(E1000_REG_TDH, 0);
    e1000_write(E1000_REG_TDT, 0);
    /* TE=1 (bit1), PSP=1 (bit3), CT=0x10<<4, COLD=0x40<<12 — matches xv6 reference */
    e1000_write(E1000_REG_TCTL, (1 << 1) | (1 << 3) | (0x10 << 4) | (0x40 << 12));
    e1000_write(E1000_REG_TIPG, 10 | (8 << 10) | (6 << 20));
    dbg_printf("[net] TX ring: TDBAL=0x%x TDH=%d TDT=%d TCTL=0x%x\n",
               tx_phys, e1000_read(E1000_REG_TDH),
               e1000_read(E1000_REG_TDT), e1000_read(E1000_REG_TCTL));

    /* Wait for link up */
    dbg_printf("[net] waiting for link...");
    for(volatile int i = 0; i < 10000000; i++) {
        uint32_t status = e1000_read(E1000_REG_STATUS);
        if(status & 0x02) { /* LU bit */
            dbg_printf(" UP (status=0x%x)\n", status);
            goto link_ok;
        }
    }
    dbg_printf(" TIMEOUT (status=0x%x)\n", e1000_read(E1000_REG_STATUS));
link_ok:;
}

/* ──────────────────────────────────────────────
 * TX
 * ────────────────────────────────────────────── */
int net_send_raw(uint16_t eth_type, const uint8_t *dst_mac,
                  const void *data, uint32_t len)
{
    uint32_t idx = tx_tail_idx;

    for(volatile int w = 0; w < 500000; w++) {
        if(tx_descs[idx].status & E1000_STS_DD) break;
    }
    if(!(tx_descs[idx].status & E1000_STS_DD)) {
        dbg_printf("[net] TX TIMEOUT desc[%d] status=0x%x\n", idx, tx_descs[idx].status);
        return -1;
    }

    tx_tail_idx = (tx_tail_idx + 1) % NET_TX_BUFS;

    uint8_t *buf = tx_bufs[idx];
    struct net_eth_hdr *eth = (struct net_eth_hdr *)buf;
    memcpy(eth->dst, dst_mac, 6);
    memcpy(eth->src, net_mac, 6);
    eth->type = __builtin_bswap16(eth_type);
    memcpy(buf + NET_ETH_HDR_LEN, data, len);

    tx_descs[idx].addr = (uint64_t)(uintptr_t)buf;
    tx_descs[idx].length = NET_ETH_HDR_LEN + len;
    tx_descs[idx].cso_csum = (uint16_t)(E1000_CMD_EOP | E1000_CMD_RS) << 8;
    tx_descs[idx].status = 0;
    tx_descs[idx].css_err = 0;
    tx_descs[idx].special = 0;

    e1000_write(E1000_REG_TDT, tx_tail_idx);
    return 0;
}

/* ──────────────────────────────────────────────
 * ARP
 * ────────────────────────────────────────────── */
static void arp_send(uint16_t opcode, const uint8_t *target_mac,
                     const uint8_t *target_ip)
{
    struct net_arp_hdr arp;
    memset(&arp, 0, sizeof(arp));
    arp.hw_type = __builtin_bswap16(1);
    arp.proto_type = __builtin_bswap16(ETH_P_IP);
    arp.hw_len = 6;
    arp.proto_len = 4;
    arp.opcode = __builtin_bswap16(opcode);
    memcpy(arp.sender_mac, net_mac, 6);
    memcpy(arp.sender_ip, net_ip_addr, 4);
    memcpy(arp.target_mac, target_mac, 6);
    memcpy(arp.target_ip, target_ip, 4);

    uint8_t broadcast[6] = {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF};
    net_send_raw(ETH_P_ARP, broadcast, &arp, sizeof(arp));
}

static void arp_handle(const void *data)
{
    const struct net_arp_hdr *arp = data;
    uint16_t opcode = __builtin_bswap16(arp->opcode);

    for(int i = 0; i < ARP_CACHE_SIZE; i++) {
        if(!arp_cache[i].valid) {
            memcpy(arp_cache[i].ip, arp->sender_ip, 4);
            memcpy(arp_cache[i].mac, arp->sender_mac, 6);
            arp_cache[i].valid = 1;
            break;
        }
        if(memcmp(arp_cache[i].ip, arp->sender_ip, 4) == 0) {
            memcpy(arp_cache[i].mac, arp->sender_mac, 6);
            break;
        }
    }

    if(opcode == ARP_OP_REQUEST && memcmp(arp->target_ip, net_ip_addr, 4) == 0) {
        arp_send(ARP_OP_REPLY, arp->sender_mac, arp->sender_ip);
    }
}

int net_arp_resolve(const uint8_t *ip, uint8_t *mac_out)
{
    for(int i = 0; i < ARP_CACHE_SIZE; i++) {
        if(arp_cache[i].valid && memcmp(arp_cache[i].ip, ip, 4) == 0) {
            memcpy(mac_out, arp_cache[i].mac, 6);
            return 0;
        }
    }

    dbg_printf("[net] ARP request for %d.%d.%d.%d\n", ip[0], ip[1], ip[2], ip[3]);
    uint8_t zero[6] = {0};
    arp_send(ARP_OP_REQUEST, zero, ip);

    uint64_t arp_start = rdtsc();
    while(rdtsc() - arp_start < 5000000000ULL) {
        net_poll();
        for(int i = 0; i < ARP_CACHE_SIZE; i++) {
            if(arp_cache[i].valid && memcmp(arp_cache[i].ip, ip, 4) == 0) {
                memcpy(mac_out, arp_cache[i].mac, 6);
                return 0;
            }
        }
    }
    dbg_printf("[net] ARP timeout for %d.%d.%d.%d\n", ip[0], ip[1], ip[2], ip[3]);
    return -1;
}

/* ──────────────────────────────────────────────
 * IP
 * ────────────────────────────────────────────── */
int net_send_ip(const uint8_t *dst_ip, uint8_t proto,
                 const void *data, uint32_t len)
{
    uint8_t buf[NET_BUF_SIZE];
    memset(buf, 0, NET_IP_HDR_LEN);
    struct net_ip_hdr *ip = (struct net_ip_hdr *)buf;

    ip->ver_ihl = 0x45;
    ip->total_len = __builtin_bswap16(NET_IP_HDR_LEN + len);
    ip->ttl = 64;
    ip->proto = proto;
    memcpy(ip->src_ip, net_ip_addr, 4);
    memcpy(ip->dst_ip, dst_ip, 4);
    ip->checksum = ip_hdr_checksum(buf, NET_IP_HDR_LEN);

    memcpy(buf + NET_IP_HDR_LEN, data, len);

    uint8_t target_ip[4];
    int same_subnet = 1;
    for(int i = 0; i < 4; i++) {
        if((net_ip_addr[i] & net_mask_addr[i]) != (dst_ip[i] & net_mask_addr[i])) {
            same_subnet = 0;
            break;
        }
    }
    memcpy(target_ip, same_subnet ? dst_ip : net_gw_addr, 4);

    uint8_t dst_mac[6];
    if(net_arp_resolve(target_ip, dst_mac) < 0) {
        dbg_printf("[net] IP send fail: no ARP for %d.%d.%d.%d\n",
                   target_ip[0], target_ip[1], target_ip[2], target_ip[3]);
        return -1;
    }

    return net_send_raw(ETH_P_IP, dst_mac, buf, NET_IP_HDR_LEN + len);
}

/* ──────────────────────────────────────────────
 * ICMP
 * ────────────────────────────────────────────── */
static uint8_t ping_id;
static uint16_t ping_seq;
static volatile int ping_got_reply;
static uint64_t ping_send_tsc;

int net_ping(const uint8_t *ip, int count)
{
    ping_id = 0x34 + (uint8_t)(rdtsc() & 0xFF);
    ping_seq = 0;

    for(int i = 0; i < count; i++) {
        uint8_t pkt[64];
        memset(pkt, 0, sizeof(pkt));
        struct net_icmp_hdr *icmp = (struct net_icmp_hdr *)pkt;
        icmp->type = ICMP_ECHO_REQUEST;
        icmp->id = __builtin_bswap16(ping_id);
        icmp->seq = __builtin_bswap16(++ping_seq);
        icmp->checksum = net_checksum(pkt, sizeof(pkt));

        ping_got_reply = 0;
        ping_send_tsc = rdtsc();
        net_send_ip(ip, IP_PROTO_ICMP, pkt, sizeof(pkt));

        uint64_t start = rdtsc();
        while(rdtsc() - start < 5000000000ULL) {
            net_poll();
            if(ping_got_reply) break;
        }
        if(!ping_got_reply)
            fb_printf("Request timeout for %d.%d.%d.%d\n",
                      ip[0], ip[1], ip[2], ip[3]);
    }
    return 0;
}

/* ──────────────────────────────────────────────
 * UDP
 * ────────────────────────────────────────────── */
int net_send_udp(const uint8_t *dst_ip, uint16_t dst_port,
                  uint16_t src_port, const void *data, uint32_t len)
{
    uint8_t buf[NET_BUF_SIZE];
    struct net_udp_hdr *udp = (struct net_udp_hdr *)buf;

    udp->src_port = __builtin_bswap16(src_port);
    udp->dst_port = __builtin_bswap16(dst_port);
    udp->length   = __builtin_bswap16(8 + len);
    udp->checksum = 0;

    memcpy(buf + 8, data, len);

    udp->checksum = 0;

    return net_send_ip(dst_ip, IP_PROTO_UDP, buf, 8 + len);
}

/* ──────────────────────────────────────────────
 * DNS
 * ────────────────────────────────────────────── */
static uint8_t dns_server[4] = {10, 0, 2, 3};
static volatile uint8_t dns_reply_ip[4];
static volatile int dns_got_reply;

static void dns_build_query(uint8_t *buf, const char *hostname, uint16_t *out_len)
{
    uint16_t id = (uint16_t)(rdtsc() & 0xFFFF);

    buf[0] = id >> 8;
    buf[1] = id & 0xFF;
    buf[2] = 0x01; buf[3] = 0x00; /* flags: standard query */
    buf[4] = 0x00; buf[5] = 0x01; /* 1 question */
    buf[6] = 0x00; buf[7] = 0x00; /* 0 answers */
    buf[8] = 0x00; buf[9] = 0x00; /* 0 authority */
    buf[10] = 0x00; buf[11] = 0x00; /* 0 additional */

    int pos = 12;
    const char *p = hostname;
    while(*p) {
        const char *dot = p;
        while(*dot && *dot != '.') dot++;
        int seglen = (int)(dot - p);
        buf[pos++] = (uint8_t)seglen;
        for(int i = 0; i < seglen; i++)
            buf[pos++] = (uint8_t)p[i];
        p = dot;
        if(*p == '.') p++;
    }
    buf[pos++] = 0; /* end name */

    buf[pos++] = 0x00; buf[pos++] = 0x01; /* type A */
    buf[pos++] = 0x00; buf[pos++] = 0x01; /* class IN */

    *out_len = pos;
}

static void dns_handle_response(const uint8_t *data, uint32_t len)
{
    if(len < 12) return;
    int qdcount = (data[4] << 8) | data[5];
    int ancount = (data[6] << 8) | data[7];

    int pos = 12;
    /* skip questions */
    for(int i = 0; i < qdcount && pos < (int)len; i++) {
        while(pos < (int)len && data[pos] != 0) {
            if((data[pos] & 0xC0) == 0xC0) { pos += 2; goto next_q; }
            pos += data[pos] + 1;
        }
        pos++; /* null */
        pos += 4; /* type + class */
    next_q:;
    }

    for(int i = 0; i < ancount && pos < (int)len; i++) {
        if((data[pos] & 0xC0) == 0xC0) pos += 2;
        else { while(pos < (int)len && data[pos] != 0) pos++; pos++; }

        if(pos + 10 > (int)len) break;
        uint16_t atype = (data[pos] << 8) | data[pos+1];
        pos += 2;
        pos += 2; /* class */
        pos += 4; /* ttl */
        uint16_t rdlen = (data[pos] << 8) | data[pos+1];
        pos += 2;

        if(atype == 1 && rdlen == 4) {
            dns_reply_ip[0] = data[pos];
            dns_reply_ip[1] = data[pos+1];
            dns_reply_ip[2] = data[pos+2];
            dns_reply_ip[3] = data[pos+3];
            dns_got_reply = 1;
            return;
        }
        pos += rdlen;
    }
}

int dns_resolve(const char *hostname, uint8_t *ip_out)
{
    uint8_t pkt[512];
    uint16_t pkt_len;

    dbg_printf("[dns] resolve '%s'\n", hostname);
    dns_build_query(pkt, hostname, &pkt_len);
    dns_got_reply = 0;
    net_send_udp(dns_server, 53, 12345, pkt, pkt_len);

    uint64_t start = rdtsc();
    while(rdtsc() - start < 2000000000ULL) {
        net_poll();
        if(dns_got_reply) {
            dbg_printf("[dns] -> %d.%d.%d.%d\n",
                       dns_reply_ip[0], dns_reply_ip[1],
                       dns_reply_ip[2], dns_reply_ip[3]);
            memcpy(ip_out, (const void *)dns_reply_ip, 4);
            return 0;
        }
    }
    dbg_printf("[dns] TIMEOUT\n");
    return -1;
}

/* ──────────────────────────────────────────────
 * TCP
 * ────────────────────────────────────────────── */
static struct tcp_conn tcp;
static volatile int tcp_got_fin;
static volatile int tcp_got_data;

static uint16_t tcp_next_port(void)
{
    static uint16_t port = 49152;
    return port++;
}

static void tcp_send_raw(uint32_t seq, uint32_t ack, uint8_t flags,
                           const void *data, uint32_t data_len)
{
    uint8_t pkt[NET_BUF_SIZE];
    struct net_tcp_hdr *h = (struct net_tcp_hdr *)pkt;

    memset(pkt, 0, 20);
    h->src_port = __builtin_bswap16(tcp.src_port);
    h->dst_port = __builtin_bswap16(tcp.dst_port);
    h->seq = __builtin_bswap32(seq);
    h->ack = __builtin_bswap32(ack);
    h->doff_flags = (5 << 4);
    h->flags = flags;
    h->window = __builtin_bswap16(16384);
    h->urgent = 0;

    if(data_len > 0)
        memcpy(pkt + 20, data, data_len);

    h->checksum = tcp_udp_checksum(net_ip_addr, tcp.dst_ip, IP_PROTO_TCP, pkt, 20 + data_len);
    if(h->checksum == 0) h->checksum = 0xFFFF;

    net_send_ip(tcp.dst_ip, IP_PROTO_TCP, pkt, 20 + data_len);
}

int tcp_connect(const uint8_t *ip, uint16_t port)
{
    dbg_printf("[tcp] connect %d.%d.%d.%d:%d\n", ip[0], ip[1], ip[2], ip[3], port);
    memset(&tcp, 0, sizeof(tcp));
    tcp.state = TCP_CLOSED;
    tcp.dst_port = port;
    tcp.src_port = tcp_next_port();
    memcpy(tcp.dst_ip, ip, 4);
    tcp.rx_head = 0;
    tcp.rx_tail = 0;
    tcp_got_fin = 0;
    tcp_got_data = 0;

    uint32_t ISN = (uint32_t)(rdtsc() & 0xFFFFFFFF);
    tcp.snd_nxt = ISN + 1;

    tcp_send_raw(ISN, 0, TCP_SYN, 0, 0);
    tcp.state = TCP_SYN_SENT;

    uint64_t start = rdtsc();
    while(tcp.state == TCP_SYN_SENT && rdtsc() - start < 2000000000ULL)
        net_poll();

    if(tcp.state != TCP_ESTABLISHED) {
        dbg_printf("[tcp] connect FAILED (state=%d)\n", tcp.state);
        tcp.state = TCP_CLOSED;
        return -1;
    }
    dbg_printf("[tcp] ESTABLISHED\n");
    return 0;
}

int tcp_send(const void *data, uint32_t len)
{
    if(tcp.state != TCP_ESTABLISHED || len == 0) return -1;

    uint32_t seq = tcp.snd_nxt;
    tcp_send_raw(seq, tcp.rcv_nxt, TCP_ACK | TCP_PSH, data, len);
    tcp.snd_nxt += len;
    return 0;
}

int tcp_recv(void *buf, uint32_t max_len, uint32_t timeout_ms)
{
    uint64_t start = rdtsc();
    uint64_t timeout_tsc = (uint64_t)timeout_ms * 2000000ULL;

    while(tcp.rx_head == tcp.rx_tail) {
        if(rdtsc() - start > timeout_tsc) return 0;
        net_poll();
        if(tcp.state == TCP_CLOSED) return -1;
    }

    uint32_t avail = (tcp.rx_head - tcp.rx_tail) & (TCP_BUF_SIZE - 1);
    uint32_t to_read = avail < max_len ? avail : max_len;
    uint32_t tail = tcp.rx_tail;

    uint8_t *dst = buf;
    for(uint32_t i = 0; i < to_read; i++) {
        dst[i] = tcp.rx_buf[(tail + i) & (TCP_BUF_SIZE - 1)];
    }
    tcp.rx_tail = (tcp.rx_tail + to_read) & (TCP_BUF_SIZE - 1);
    return (int)to_read;
}

void tcp_close(void)
{
    dbg_printf("[tcp] close (state=%d)\n", tcp.state);
    if(tcp.state == TCP_ESTABLISHED) {
        tcp_send_raw(tcp.snd_nxt, tcp.rcv_nxt, TCP_FIN | TCP_ACK, 0, 0);
        tcp.snd_nxt++;
        tcp.state = TCP_FIN_WAIT_1;
    } else if(tcp.state == TCP_CLOSE_WAIT) {
        tcp_send_raw(tcp.snd_nxt, tcp.rcv_nxt, TCP_FIN | TCP_ACK, 0, 0);
        tcp.snd_nxt++;
        tcp.state = TCP_LAST_ACK;
    }
    tcp.state = TCP_CLOSED;
}

static void tcp_handle(const uint8_t *data, uint32_t len, const uint8_t *src_ip)
{
    if(len < 20) return;
    const struct net_tcp_hdr *h = (const struct net_tcp_hdr *)data;

    uint16_t dst_port = __builtin_bswap16(h->dst_port);
    if(dst_port != tcp.src_port) return;
    if(memcmp(src_ip, tcp.dst_ip, 4) != 0) return;

    uint32_t seq = __builtin_bswap32(h->seq);
    uint8_t flags = h->flags;
    uint8_t hdr_len = (h->doff_flags >> 4) * 4;
    uint32_t payload_len = len - hdr_len;

    switch(tcp.state) {
    case TCP_SYN_SENT:
        if((flags & TCP_SYN) && (flags & TCP_ACK)) {
            tcp.rcv_nxt = seq + 1;
            tcp_send_raw(tcp.snd_nxt, tcp.rcv_nxt, TCP_ACK, 0, 0);
            tcp.state = TCP_ESTABLISHED;
        }
        break;

    case TCP_ESTABLISHED:
        if(flags & TCP_FIN) {
            tcp.rcv_nxt = seq + 1;
            tcp_send_raw(tcp.snd_nxt, tcp.rcv_nxt, TCP_ACK, 0, 0);
            tcp.state = TCP_CLOSE_WAIT;
            tcp_got_fin = 1;
        }
        if(payload_len > 0 && (flags & (TCP_ACK | TCP_PSH))) {
            if(seq == tcp.rcv_nxt) {
                uint32_t space = TCP_BUF_SIZE - ((tcp.rx_head - tcp.rx_tail) & (TCP_BUF_SIZE - 1));
                uint32_t to_copy = payload_len < space ? payload_len : space;
                for(uint32_t i = 0; i < to_copy; i++) {
                    tcp.rx_buf[(tcp.rx_head + i) & (TCP_BUF_SIZE - 1)] = data[hdr_len + i];
                }
                tcp.rx_head = (tcp.rx_head + to_copy) & (TCP_BUF_SIZE - 1);
                tcp.rcv_nxt += payload_len;
                tcp_got_data = 1;
            }
            tcp_send_raw(tcp.snd_nxt, tcp.rcv_nxt, TCP_ACK, 0, 0);
        }
        break;

    case TCP_FIN_WAIT_1:
        if(flags & TCP_ACK) {
            tcp.state = TCP_FIN_WAIT_2;
        }
        break;

    case TCP_FIN_WAIT_2:
        if(flags & TCP_FIN) {
            tcp.rcv_nxt = seq + 1;
            tcp_send_raw(tcp.snd_nxt, tcp.rcv_nxt, TCP_ACK, 0, 0);
            tcp.state = TCP_CLOSED;
        }
        break;

    case TCP_CLOSE_WAIT:
        break;

    case TCP_LAST_ACK:
        if(flags & TCP_ACK) {
            tcp.state = TCP_CLOSED;
        }
        break;
    default:
        break;
    }
}

/* ──────────────────────────────────────────────
 * RX poll
 * ────────────────────────────────────────────── */
static uint8_t dns_rx_buf[512];
static uint32_t dns_rx_len;

void net_poll(void)
{
    if(!e1000_mmio) return;

    /* Clear any pending interrupts */
    e1000_read(E1000_REG_ICR);

    /* Process received packets — exact xv6 reference pattern */
    while(1) {
        uint32_t idx = (e1000_read(E1000_REG_RDT) + 1) % NET_RX_BUFS;

        if(!(rx_descs[idx].status & E1000_STS_DD)) break;

        uint32_t len = rx_descs[idx].length;
        uint8_t *pkt = rx_bufs[idx];
        struct net_eth_hdr *eth = (struct net_eth_hdr *)pkt;
        uint16_t etype = __builtin_bswap16(eth->type);

        if(etype == ETH_P_ARP && len >= NET_ETH_HDR_LEN + sizeof(struct net_arp_hdr)) {
            arp_handle(pkt + NET_ETH_HDR_LEN);
        }
        else if(etype == ETH_P_IP && len >= NET_ETH_HDR_LEN + NET_IP_HDR_LEN) {
            struct net_ip_hdr *iph = (struct net_ip_hdr *)(pkt + NET_ETH_HDR_LEN);
            uint32_t ip_total = __builtin_bswap16(iph->total_len);

            if(iph->proto == IP_PROTO_ICMP) {
                struct net_icmp_hdr *ic = (struct net_icmp_hdr *)((uint8_t *)iph + NET_IP_HDR_LEN);
                if(ic->type == ICMP_ECHO_REPLY && ic->id == __builtin_bswap16(ping_id)) {
                    ping_got_reply = 1;
                }
            }
            else if(iph->proto == IP_PROTO_UDP) {
                struct net_udp_hdr *udp = (struct net_udp_hdr *)((uint8_t *)iph + NET_IP_HDR_LEN);
                uint16_t dport = __builtin_bswap16(udp->dst_port);

                if(dport == 12345) {
                    uint32_t udp_len = __builtin_bswap16(udp->length);
                    uint32_t payload = udp_len - 8;
                    if(payload > 512) payload = 512;
                    dns_rx_len = payload;
                    memcpy(dns_rx_buf, (uint8_t *)udp + 8, payload);
                    dbg_printf("[rx] DNS reply %d bytes\n", payload);
                    dns_handle_response(dns_rx_buf, dns_rx_len);
                }
            }
            else if(iph->proto == IP_PROTO_TCP) {
                tcp_handle((uint8_t *)iph + NET_IP_HDR_LEN,
                           ip_total - NET_IP_HDR_LEN, iph->src_ip);
            }
        }

        rx_descs[idx].status = 0;
        e1000_write(E1000_REG_RDT, idx);
    }
}

/* ──────────────────────────────────────────────
 * Public
 * ────────────────────────────────────────────── */
void net_get_mac(uint8_t *mac) { memcpy(mac, net_mac, 6); }
void net_get_ip(uint8_t *ip)   { memcpy(ip, net_ip_addr, 4); }
void net_get_mask(uint8_t *m)  { memcpy(m, net_mask_addr, 4); }
void net_get_gw(uint8_t *gw)   { memcpy(gw, net_gw_addr, 4); }
void net_set_ip(const uint8_t *ip)   { memcpy(net_ip_addr, ip, 4); }
void net_set_mask(const uint8_t *m)  { memcpy(net_mask_addr, m, 4); }
void net_set_gw(const uint8_t *gw)   { memcpy(net_gw_addr, gw, 4); }

int net_init(void)
{
    uint8_t bus = 0, dev = 0;
    for(uint16_t b = 0; b < 256; b++) {
        for(uint8_t d = 0; d < 32; d++) {
            uint32_t id = pci_read(b, d, 0, 0);
            uint16_t ven = id & 0xFFFF;
            uint16_t dev_id = (id >> 16) & 0xFFFF;
            if(ven == 0x8086 && (dev_id == 0x100E || dev_id == 0x1004 ||
                                  dev_id == 0x100F || dev_id == 0x10EA)) {
                bus = b; dev = d;
                goto found;
            }
        }
    }
    fb_printf("net: no e1000 found\n");
    return -1;
found:
    e1000_bus = bus;
    e1000_dev = dev;
    dbg_printf("[net] e1000 at PCI %d:%d\n", bus, dev);

    tx_descs = tx_desc_area;
    rx_descs = rx_desc_area;
    for(int i = 0; i < NET_TX_BUFS; i++) {
        tx_bufs[i] = tx_buf_area + i * NET_BUF_SIZE;
        tx_descs[i].addr = (uint64_t)(uintptr_t)tx_bufs[i];
        tx_descs[i].status = E1000_STS_DD;
    }
    for(int i = 0; i < NET_RX_BUFS; i++) {
        rx_bufs[i] = rx_buf_area + i * NET_BUF_SIZE;
        rx_descs[i].addr = (uint64_t)(uintptr_t)rx_bufs[i];
    }

    uint32_t cmd = pci_read(bus, dev, 0, 0x04);
    pci_write(bus, dev, 0, 0x04, cmd | 0x04 | 0x02);
    uint32_t cmd2 = pci_read(bus, dev, 0, 0x04);
    dbg_printf("[net] PCI CMD old=0x%x new=0x%x busmaster=%d\n",
               cmd, cmd2, (cmd2 & 0x04) != 0);

    uint32_t bar0 = pci_read(bus, dev, 0, 0x10);
    dbg_printf("[net] BAR0=0x%x\n", bar0);
    e1000_mmio = (volatile uint32_t *)(uintptr_t)(bar0 & ~0xF);
    dbg_printf("[net] MMIO base=0x%x\n", (uint32_t)(uintptr_t)e1000_mmio);

    e1000_init_mmio();

    uint8_t default_ip[4] = {10, 0, 2, 15};
    uint8_t default_mask[4] = {255, 255, 255, 0};
    uint8_t default_gw[4] = {10, 0, 2, 2};
    memcpy(net_ip_addr, default_ip, 4);
    memcpy(net_mask_addr, default_mask, 4);
    memcpy(net_gw_addr, default_gw, 4);

    /* Hardcode SLiRP virtual MACs — SLiRP uses ONE NIC, all IPs share same MAC */
    {
        uint8_t mac[6] = {0x52, 0x55, 0x0a, 0x00, 0x02, 0x02};
        memcpy(arp_cache[0].ip, net_gw_addr, 4);
        memcpy(arp_cache[0].mac, mac, 6);
        arp_cache[0].valid = 1;
    }
    {
        uint8_t mac[6] = {0x52, 0x55, 0x0a, 0x00, 0x02, 0x02};
        memcpy(arp_cache[1].ip, dns_server, 4);
        memcpy(arp_cache[1].mac, mac, 6);
        arp_cache[1].valid = 1;
    }

    fb_printf("e1000: MAC %02x:%02x:%02x:%02x:%02x:%02x\n",
              net_mac[0], net_mac[1], net_mac[2],
              net_mac[3], net_mac[4], net_mac[5]);
    fb_printf("net: IP %d.%d.%d.%d\n",
              net_ip_addr[0], net_ip_addr[1], net_ip_addr[2], net_ip_addr[3]);
    return 0;
}
