//! K434 ARP and K437 ICMP: separate images reuse one owned RX→TX DMA path.
//! This is bounded fixed-IP traffic, not a general network stack.
use super::*;
use sha2::{Digest, Sha256};
#[cfg(not(feature = "virtio-net-icmp-echo-test"))]
use turkos_kernel::net_arp;
#[cfg(feature = "virtio-net-icmp-echo-test")]
use turkos_kernel::net_icmp;
use turkos_kernel::{net_frame, net_mac, program_hash, virtio_mmio, virtio_ring};

#[cfg(all(
    not(feature = "virtio-net-icmp-echo-test"),
    not(feature = "virtio-net-reset-traffic-test")
))]
const MARKER: &str = "TURKOS/NET-ARP-REPLY";
#[cfg(feature = "virtio-net-icmp-echo-test")]
const MARKER: &str = "TURKOS/NET-ICMP";
#[cfg(feature = "virtio-net-reset-traffic-test")]
const MARKER: &str = "TURKOS/NET-RESET-TRAFFIC";

fn counter() -> u64 {
    let value: u64;
    unsafe {
        asm!("mrs {value}, CNTPCT_EL0", value = out(reg) value, options(nomem, nostack, preserves_flags))
    };
    value
}

unsafe fn read(base: u64, offset: u64) -> u32 {
    unsafe { read_volatile((base + offset) as *const u32) }
}
unsafe fn write(base: u64, offset: u64, value: u32) {
    unsafe { write_volatile((base + offset) as *mut u32, value) }
}
unsafe fn write64(base: u64, offset: u64, address: u64) {
    unsafe {
        write(base, offset, address as u32);
        write(base, offset + 4, (address >> 32) as u32);
    }
}
fn fail(reason: &str) -> ! {
    write_str(MARKER);
    write_str("-ERROR: ");
    write_line(reason);
    park()
}

fn wait_used(used: u64, label: &str) -> Option<(u16, u32, u32)> {
    let frequency: u64;
    unsafe {
        asm!("mrs {value}, CNTFRQ_EL0", value = out(reg) frequency, options(nomem, nostack, preserves_flags))
    };
    let start = counter();
    let budget = frequency
        .checked_mul(3)
        .unwrap_or_else(|| fail("deadline overflow"));
    let deadline = virtio_ring::deadline_after(start, budget).unwrap_or_else(|_| fail("deadline"));
    let index = loop {
        let index = unsafe { u16::from_le(read_volatile((used + 2) as *const u16)) };
        if index != 0 {
            break index;
        }
        let observed = counter();
        if virtio_ring::deadline_expired(observed, deadline) {
            write_str(MARKER);
            write_str(": TIMEOUT KIND=");
            write_str(if label == "RX timeout" { "RX" } else { "TX" });
            write_str(" START=0x");
            write_hex(start as usize);
            write_str(" DEADLINE=0x");
            write_hex(deadline as usize);
            write_str(" OBSERVED=0x");
            write_hex(observed as usize);
            write_str(" FREQUENCY=0x");
            write_hex(frequency as usize);
            write_str("\r\n");
            return None;
        }
        core::hint::spin_loop();
    };
    let (id, len) = unsafe {
        asm!("dmb oshld", options(nostack, preserves_flags));
        (
            u32::from_le(read_volatile((used + 4) as *const u32)),
            u32::from_le(read_volatile((used + 8) as *const u32)),
        )
    };
    Some((index, id, len))
}

fn stop_and_release(
    base: u64,
    rx: &mut virtio_ring::DmaLease,
    tx: &mut virtio_ring::DmaLease,
    allocator: &mut memory::FrameAllocator,
    free_before: u64,
    rx_buffer: u64,
    tx_buffer: u64,
) {
    unsafe {
        write(base, 0x70, 0);
        asm!("dsb sy", options(nostack, preserves_flags));
        if read(base, 0x70) != 0 {
            fail("reset status");
        }
        for index in 0..2 {
            write(base, 0x30, index);
            if read(base, 0x44) != 0 {
                fail("reset QueueReady");
            }
        }
    }
    rx.mark_reset().unwrap_or_else(|_| fail("RX reset state"));
    tx.mark_reset().unwrap_or_else(|_| fail("TX reset state"));
    if unsafe { read_volatile((rx_buffer + 2048) as *const u8) } != 0xa5
        || unsafe { read_volatile((tx_buffer + 72) as *const u8) } != 0xa5
    {
        fail("DMA guards");
    }
    if rx.release(allocator).is_err()
        || tx.release(allocator).is_err()
        || allocator.remaining_pages() != free_before
    {
        fail("DMA return accounting");
    }
}

fn print_hash(label: &str, frame: &[u8; 60]) {
    let digest: [u8; 32] = Sha256::digest(frame).into();
    let hex = program_hash::hex_lower(&digest);
    write_str(label);
    write_str(core::str::from_utf8(&hex).unwrap_or_else(|_| fail("digest hex")));
    write_str("\r\n");
}

pub fn smoke_test(boot: &BootData, allocator: &mut memory::FrameAllocator) -> ! {
    #[cfg(feature = "virtio-net-reset-traffic-test")]
    {
        run_cycle(boot, allocator, 1);
        run_cycle(boot, allocator, 2);
        write_line("TURKOS/NET-RESET-TRAFFIC: PASS CYCLES=2 RX=2 TX=2 RESETS=2 DMA-RELEASED=16");
        park()
    }
    #[cfg(not(feature = "virtio-net-reset-traffic-test"))]
    {
        run_cycle(boot, allocator, 0);
        park()
    }
}

fn run_cycle(boot: &BootData, allocator: &mut memory::FrameAllocator, cycle: usize) {
    #[cfg(not(feature = "virtio-net-reset-traffic-test"))]
    let _ = cycle;
    let mut net = None;
    for node in boot.virtio_mmio.as_slice() {
        let magic = unsafe { read(node.base, 0) };
        let version = unsafe { read(node.base, 4) };
        let id = unsafe { read(node.base, 8) };
        if virtio_mmio::validate_magic(magic).is_err() {
            fail("magic");
        }
        if version != 2 {
            fail("modern version required");
        }
        if virtio_mmio::require_net_device(id).is_ok() && net.replace(node.base).is_some() {
            fail("ambiguous net");
        }
    }
    let base = net.unwrap_or_else(|| fail("missing net"));
    let free_before = allocator.remaining_pages();
    let mut rx =
        virtio_ring::DmaLease::allocate(allocator, boot.hardware.ram_base, boot.hardware.ram_size)
            .unwrap_or_else(|_| fail("RX DMA allocation"));
    let mut tx =
        virtio_ring::DmaLease::allocate(allocator, boot.hardware.ram_base, boot.hardware.ram_size)
            .unwrap_or_else(|_| fail("TX DMA allocation"));
    let rx_frames = rx.frames().unwrap_or_else(|_| fail("RX frames"));
    let tx_frames = tx.frames().unwrap_or_else(|_| fail("TX frames"));
    for (index, frame) in rx_frames.iter().chain(tx_frames.iter()).enumerate() {
        if rx_frames
            .iter()
            .chain(tx_frames.iter())
            .take(index)
            .any(|other| other == frame)
        {
            fail("RX/TX DMA alias");
        }
    }
    for [desc, avail, used, _] in [rx_frames, tx_frames] {
        virtio_ring::split_ring_pages(
            desc,
            avail,
            used,
            8,
            boot.hardware.ram_base,
            boot.hardware.ram_size,
        )
        .unwrap_or_else(|_| fail("ring geometry"));
    }
    let [rx_desc, rx_avail, rx_used, rx_buffer] = rx_frames;
    let [tx_desc, tx_avail, tx_used, tx_buffer] = tx_frames;
    unsafe {
        for page in rx_frames.iter().chain(tx_frames.iter()) {
            core::ptr::write_bytes(*page as *mut u8, 0, memory::PAGE_SIZE as usize);
        }
        write_volatile(
            rx_desc as *mut virtio_ring::Descriptor,
            virtio_ring::Descriptor {
                addr_le: rx_buffer.to_le(),
                len_le: 2048u32.to_le(),
                flags_le: virtio_ring::DESC_F_WRITE.to_le(),
                next_le: 0,
            },
        );
        write_volatile(
            tx_desc as *mut virtio_ring::Descriptor,
            virtio_ring::Descriptor {
                addr_le: tx_buffer.to_le(),
                len_le: 72u32.to_le(),
                flags_le: 0,
                next_le: 0,
            },
        );
        write_volatile((rx_avail + 4) as *mut u16, 0u16.to_le());
        write_volatile((tx_avail + 4) as *mut u16, 0u16.to_le());
        write_volatile((rx_buffer + 2048) as *mut u8, 0xa5);
        write_volatile((tx_buffer + 72) as *mut u8, 0xa5);
    }
    if free_before.checked_sub(allocator.remaining_pages()) != Some(8) {
        fail("DMA accounting");
    }
    let mut own = [0u8; 6];
    unsafe {
        write(base, 0x70, 0);
        if read(base, 0x70) != 0 {
            fail("initial reset");
        }
        write(base, 0x70, 3);
        write(base, 0x14, 0);
        let low = read(base, 0x10);
        write(base, 0x14, 1);
        let high = read(base, 0x10);
        if low & net_mac::VIRTIO_NET_F_MAC == 0 || high & virtio_mmio::MODERN_F_VERSION_1 == 0 {
            fail("MAC/VERSION_1 missing");
        }
        write(base, 0x24, 0);
        write(base, 0x20, net_mac::VIRTIO_NET_F_MAC);
        write(base, 0x24, 1);
        write(base, 0x20, virtio_mmio::MODERN_F_VERSION_1);
        write(base, 0x70, 0x0b);
        if read(base, 0x70) != 0x0b {
            fail("FEATURES_OK rejected");
        }
        for (index, [desc, avail, used, _]) in [rx_frames, tx_frames].into_iter().enumerate() {
            write(base, 0x30, index as u32);
            if read(base, 0x34) < 8 || read(base, 0x44) != 0 {
                fail("queue capacity/stale ready");
            }
            write(base, 0x38, 8);
            write64(base, 0x80, desc);
            write64(base, 0x90, avail);
            write64(base, 0xa0, used);
            asm!("dmb oshst", options(nostack, preserves_flags));
            write(base, 0x44, 1);
            if read(base, 0x44) != 1 {
                fail("QueueReady readback");
            }
        }
        rx.mark_ready().unwrap_or_else(|_| fail("RX ready"));
        tx.mark_ready().unwrap_or_else(|_| fail("TX ready"));
        write(base, 0x70, 0x0f);
        if read(base, 0x70) != 0x0f {
            fail("DRIVER_OK rejected");
        }
        for (index, byte) in own.iter_mut().enumerate() {
            *byte = read_volatile((base + 0x100 + index as u64) as *const u8);
        }
    }
    if own != [0x52, 0x54, 0, 0x12, 0x34, 0x56]
        || net_mac::validate(net_mac::VIRTIO_NET_F_MAC, own).is_err()
    {
        fail("own MAC");
    }
    if rx.release(allocator) != Err(virtio_ring::Error::Active)
        || tx.release(allocator) != Err(virtio_ring::Error::Active)
    {
        fail("early DMA release");
    }
    unsafe {
        write_volatile((rx_avail + 2) as *mut u16, 1u16.to_le());
        asm!("dmb oshst", options(nostack, preserves_flags));
        write(base, 0x50, 0);
    }
    #[cfg(all(
        not(feature = "virtio-net-icmp-echo-test"),
        not(feature = "virtio-net-reset-traffic-test")
    ))]
    write_line("TURKOS/NET-ARP-REPLY: WAITING-REQUEST");
    #[cfg(feature = "virtio-net-icmp-echo-test")]
    write_line("TURKOS/NET-ICMP: WAITING-REQUEST");
    #[cfg(feature = "virtio-net-reset-traffic-test")]
    {
        write_str("TURKOS/NET-RESET-TRAFFIC: WAITING-REQUEST CYCLE=0x");
        write_hex(cycle);
        write_str("\r\n");
    }
    let (rx_index, rx_id, rx_len) = wait_used(rx_used, "RX timeout").unwrap_or_else(|| {
        stop_and_release(
            base,
            &mut rx,
            &mut tx,
            allocator,
            free_before,
            rx_buffer,
            tx_buffer,
        );
        write_str(MARKER);
        write_line(": PASS TIMEOUT=RX RESET=1 DMA-RELEASED=8 TX=0");
        park()
    });
    if net_frame::validate_fixture_completion(rx_index, rx_id, rx_len, 60).is_err() {
        fail("RX used index/id/length");
    }
    let mut request_frame = [0u8; 60];
    for (index, byte) in request_frame.iter_mut().enumerate() {
        *byte = unsafe { read_volatile((rx_buffer + 12 + index as u64) as *const u8) };
    }
    #[cfg(not(feature = "virtio-net-icmp-echo-test"))]
    let reply = {
        let ethernet = net_frame::copy_ethernet_header(&request_frame)
            .unwrap_or_else(|_| fail("Ethernet header"));
        if ethernet[..6] != [0xff; 6] || ethernet[12..14] != [0x08, 0x06] {
            fail("Ethernet broadcast/ARP type");
        }
        let request = net_arp::parse_request(&request_frame[14..])
            .unwrap_or_else(|_| fail("ARP request header"));
        if request.sender_mac != ethernet[6..12]
            || request.target_mac != [0; 6]
            || request_frame[42..].iter().any(|byte| *byte != 0)
        {
            fail("ARP request fixture bytes");
        }
        net_arp::build_reply(request, own, [192, 0, 2, 2])
            .unwrap_or_else(|_| fail("ARP reply target/IP/MAC policy"))
    };
    #[cfg(feature = "virtio-net-icmp-echo-test")]
    let reply = net_icmp::build_echo_reply(&request_frame, own, [192, 0, 2, 2])
        .unwrap_or_else(|_| fail("ICMP echo request policy"));
    #[cfg(all(
        not(feature = "virtio-net-icmp-echo-test"),
        not(feature = "virtio-net-reset-traffic-test")
    ))]
    {
        print_hash("TURKOS/NET-ARP-REPLY: REQUEST-SHA256=", &request_frame);
        print_hash("TURKOS/NET-ARP-REPLY: REPLY-SHA256=", &reply);
    }
    #[cfg(feature = "virtio-net-reset-traffic-test")]
    {
        write_str("TURKOS/NET-RESET-TRAFFIC: CYCLE=0x");
        write_hex(cycle);
        write_str(" ");
        print_hash("REQUEST-SHA256=", &request_frame);
        write_str("TURKOS/NET-RESET-TRAFFIC: CYCLE=0x");
        write_hex(cycle);
        write_str(" ");
        print_hash("REPLY-SHA256=", &reply);
    }
    #[cfg(feature = "virtio-net-icmp-echo-test")]
    {
        print_hash("TURKOS/NET-ICMP: REQUEST-SHA256=", &request_frame);
        print_hash("TURKOS/NET-ICMP: REPLY-SHA256=", &reply);
    }
    unsafe {
        // The 12-byte modern virtio-net header stays zero: no offload/GSO.
        for (index, byte) in reply.iter().enumerate() {
            write_volatile((tx_buffer + 12 + index as u64) as *mut u8, *byte);
        }
        write_volatile((tx_avail + 2) as *mut u16, 1u16.to_le());
        asm!("dmb oshst", options(nostack, preserves_flags));
        write(base, 0x50, 1);
    }
    #[cfg(all(
        not(feature = "virtio-net-icmp-echo-test"),
        not(feature = "virtio-net-reset-traffic-test")
    ))]
    write_line("TURKOS/NET-ARP-REPLY: WAITING-TX-USED");
    #[cfg(feature = "virtio-net-icmp-echo-test")]
    write_line("TURKOS/NET-ICMP: WAITING-TX-USED");
    #[cfg(feature = "virtio-net-reset-traffic-test")]
    {
        write_str("TURKOS/NET-RESET-TRAFFIC: WAITING-TX-USED CYCLE=0x");
        write_hex(cycle);
        write_str("\r\n");
    }
    let (tx_index, tx_id, tx_len) = wait_used(tx_used, "TX timeout").unwrap_or_else(|| {
        stop_and_release(
            base,
            &mut rx,
            &mut tx,
            allocator,
            free_before,
            rx_buffer,
            tx_buffer,
        );
        write_str(MARKER);
        write_line(": PASS TIMEOUT=TX RESET=1 DMA-RELEASED=8 TX=SUBMITTED");
        park()
    });
    if (tx_index, tx_id, tx_len) != (1, 0, 0) {
        fail("TX used index/id/length");
    }
    if rx.release(allocator) != Err(virtio_ring::Error::Active)
        || tx.release(allocator) != Err(virtio_ring::Error::Active)
    {
        fail("post-used early release");
    }
    for (index, byte) in reply.iter().enumerate() {
        if unsafe { read_volatile((tx_buffer + 12 + index as u64) as *const u8) } != *byte {
            fail("TX buffer modified");
        }
    }
    stop_and_release(
        base,
        &mut rx,
        &mut tx,
        allocator,
        free_before,
        rx_buffer,
        tx_buffer,
    );
    #[cfg(all(
        not(feature = "virtio-net-icmp-echo-test"),
        not(feature = "virtio-net-reset-traffic-test")
    ))]
    write_line(
        "TURKOS/NET-ARP-REPLY: RX-USED=1/0/72 TX-USED=1/0/0 OWN=525400123456 SPA=c0000202 THA=020000000001 TPA=c0000201 DMA-RELEASED=8",
    );
    #[cfg(all(
        not(feature = "virtio-net-icmp-echo-test"),
        not(feature = "virtio-net-reset-traffic-test")
    ))]
    write_line("TURKOS/NET-ARP-REPLY: PASS RX=1 TX=1 RESET=1 DMA-RELEASED=8");
    #[cfg(feature = "virtio-net-reset-traffic-test")]
    {
        write_str("TURKOS/NET-RESET-TRAFFIC: CYCLE=0x");
        write_hex(cycle);
        write_line(" RX-USED=1/0/72 TX-USED=1/0/0 RESET=1 DMA-RELEASED=8");
    }
    #[cfg(feature = "virtio-net-icmp-echo-test")]
    write_line(
        "TURKOS/NET-ICMP: RX-USED=1/0/72 TX-USED=1/0/0 TYPE=8/0 PROTOCOL=1 ID=1234 SEQ=0001 DMA-RELEASED=8",
    );
    #[cfg(feature = "virtio-net-icmp-echo-test")]
    write_line("TURKOS/NET-ICMP: PASS RX=1 TX=1 RESET=1 DMA-RELEASED=8");
}