use core::{
    future::Future,
    mem::size_of,
    net::{IpAddr, Ipv4Addr},
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};

use alloc::{boxed::Box, collections::BTreeMap};
use libk::task::runtime;
use libk_util::{sync::spin_rwlock::IrqSafeRwLock, waker::QueueWaker};
use yggdrasil_abi::{
    error::Error,
    net::{
        protocols::{ArpFrame, EtherType},
        types::NetValueImpl,
        MacAddress,
    },
};

use crate::{ethernet, interface::NetworkInterface, L2Packet};

struct Inner<A: Ord + Eq + Copy> {
    entries: BTreeMap<(u32, A), (MacAddress, bool)>,
    reverse: BTreeMap<(u32, MacAddress), (A, bool)>,
}

pub struct ArpTable {
    v4: IrqSafeRwLock<Inner<Ipv4Addr>>,
    notify: QueueWaker,
}

impl<A: Ord + Eq + Copy> Inner<A> {
    const fn new() -> Self {
        Self {
            entries: BTreeMap::new(),
            reverse: BTreeMap::new(),
        }
    }

    fn query_mac(&self, interface: u32, address: A) -> Option<(MacAddress, bool)> {
        self.entries.get(&(interface, address)).copied()
    }

    // fn query_address(&self, interface: u32, mac: MacAddress) -> Option<(A, bool)> {
    //     self.reverse.get(&(interface, mac)).copied()
    // }

    fn insert(&mut self, interface: u32, mac: MacAddress, address: A, owned: bool) -> bool {
        let new = self
            .entries
            .insert((interface, address), (mac, owned))
            .is_none();
        self.reverse.insert((interface, mac), (address, owned));
        new
    }

    fn flush(&mut self, interface: u32, address: A) -> bool {
        if let Some((mac, _)) = self.entries.remove(&(interface, address)) {
            self.reverse.remove(&(interface, mac));
            true
        } else {
            false
        }
    }
}

impl ArpTable {
    pub const fn new() -> Self {
        Self {
            v4: IrqSafeRwLock::new(Inner::new()),
            notify: QueueWaker::new(),
        }
    }

    pub fn lookup_cache_v4(interface: u32, address: Ipv4Addr) -> Option<(MacAddress, bool)> {
        ARP_TABLE.v4.read().query_mac(interface, address)
    }

    pub fn lookup_cache(interface: u32, address: IpAddr) -> Option<MacAddress> {
        let (address, _) = match address {
            IpAddr::V4(address) => Self::lookup_cache_v4(interface, address),
            IpAddr::V6(v6) => {
                log::warn!("TODO: ArpTable v6 lookup: {v6}");
                return None;
            }
        }?;
        Some(address)
    }

    pub fn flush_address_v4(interface: u32, address: Ipv4Addr) -> bool {
        ARP_TABLE.v4.write().flush(interface, address)
    }

    pub fn flush_address(interface: u32, address: IpAddr) -> bool {
        match address {
            IpAddr::V4(address) => Self::flush_address_v4(interface, address),
            IpAddr::V6(v6) => {
                log::warn!("TODO: ArpTable v6 flush: {v6}");
                false
            }
        }
    }

    pub fn insert_address_v4(interface: u32, mac: MacAddress, address: Ipv4Addr, owned: bool) {
        ARP_TABLE.v4.write().insert(interface, mac, address, owned);
    }

    pub fn insert_address(interface: u32, mac: MacAddress, address: IpAddr, owned: bool) {
        match address {
            IpAddr::V4(address) => Self::insert_address_v4(interface, mac, address, owned),
            IpAddr::V6(v6) => {
                log::warn!("TODO: ArpTable v6 insert: {v6}");
                return;
            }
        }
        ARP_TABLE.notify.wake_all();
    }

    fn poll_address(
        interface: u32,
        address: IpAddr,
        timeout: Duration,
    ) -> impl Future<Output = Option<MacAddress>> {
        struct F<T: Future<Output = ()>> {
            interface: u32,
            address: IpAddr,
            timeout: Pin<Box<T>>,
        }

        impl<T: Future<Output = ()>> Future for F<T> {
            type Output = Option<MacAddress>;

            fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                if self.timeout.as_mut().poll(cx).is_ready() {
                    ARP_TABLE.notify.remove(cx.waker());
                    return Poll::Ready(None);
                }
                ARP_TABLE.notify.register(cx.waker());
                if let Some(mac) = ArpTable::lookup_cache(self.interface, self.address) {
                    ARP_TABLE.notify.remove(cx.waker());
                    Poll::Ready(Some(mac))
                } else {
                    Poll::Pending
                }
            }
        }

        F {
            interface,
            address,
            timeout: Box::pin(runtime::sleep(timeout)),
        }
    }
}

static ARP_TABLE: ArpTable = ArpTable::new();

pub async fn lookup(interface: u32, ip: IpAddr, perform_query: bool) -> Result<MacAddress, Error> {
    if let Some(mac) = ArpTable::lookup_cache(interface, ip) {
        return Ok(mac);
    }
    if !perform_query {
        return Err(Error::HostUnreachable);
    }
    query(interface, ip, 5, Duration::from_millis(200)).await
}

pub fn lookup_nonblocking(
    interface: u32,
    ip: IpAddr,
    perform_query: bool,
) -> Result<MacAddress, Error> {
    if let Some(mac) = ArpTable::lookup_cache(interface, ip) {
        return Ok(mac);
    }
    if perform_query {
        Err(Error::WouldBlock)
    } else {
        Err(Error::HostUnreachable)
    }
}

async fn query(
    interface: u32,
    ip: IpAddr,
    retries: usize,
    retry_timeout: Duration,
) -> Result<MacAddress, Error> {
    let interface = NetworkInterface::get(interface)?;
    for _ in 0..retries {
        send_request(&interface, ip)?;
        if let Some(mac) = ArpTable::poll_address(interface.id, ip, retry_timeout).await {
            return Ok(mac);
        }
    }
    Err(Error::HostUnreachable)
}

fn send_request_v4(interface: &NetworkInterface, query_address: Ipv4Addr) -> Result<(), Error> {
    let request = ArpFrame {
        protocol: EtherType::IPV4.to_network_order(),
        protocol_size: 4,
        hardware_type: u16::to_network_order(1),
        hardware_size: 6,
        opcode: 1u16.to_network_order(),

        sender_mac: interface.mac,
        // TODO maybe would be nice to specify
        sender_ip: 0u32.to_network_order(),

        target_ip: u32::to_network_order(query_address.into()),
        target_mac: MacAddress::UNSPECIFIED,
    };

    ethernet::send_l2(
        interface,
        interface.mac,
        MacAddress::BROADCAST,
        EtherType::ARP,
        &request,
    )
}

fn send_request(interface: &NetworkInterface, query_address: IpAddr) -> Result<(), Error> {
    log::debug!("Querying address of {}", query_address);

    match query_address {
        IpAddr::V4(address) => send_request_v4(interface, address),
        IpAddr::V6(v6) => {
            log::warn!("TODO: ARP IPv6 query: {v6}");
            Err(Error::NotImplemented)
        }
    }
}

fn send_reply(interface_id: u32, arp: &ArpFrame, target_mac: MacAddress) -> Result<(), Error> {
    let interface = NetworkInterface::get(interface_id)?;
    let reply = ArpFrame {
        protocol: arp.protocol,
        hardware_type: arp.hardware_type,
        hardware_size: arp.hardware_size,
        protocol_size: arp.protocol_size,
        opcode: 2u16.to_network_order(),

        sender_mac: target_mac,
        sender_ip: arp.target_ip,

        target_ip: arp.sender_ip,
        target_mac: arp.sender_mac,
    };

    ethernet::send_l2(
        &interface,
        target_mac,
        arp.sender_mac,
        EtherType::ARP,
        &reply,
    )
}

pub fn handle_packet(packet: L2Packet) {
    let arp: &ArpFrame = bytemuck::from_bytes(&packet.l2_data()[..size_of::<ArpFrame>()]);
    let proto = EtherType::from_network_order(arp.protocol);
    let opcode = u16::from_network_order(arp.opcode);

    let (target_address, sender_address) = match proto {
        EtherType::IPV4 => (
            Ipv4Addr::from(u32::from_network_order(arp.target_ip)),
            Ipv4Addr::from(u32::from_network_order(arp.sender_ip)),
        ),
        _ => {
            log::warn!("TODO: unhandled ARP proto: {:#x?}", proto);
            return;
        }
    };

    log::debug!("ARP: {} -> {}", sender_address, target_address);

    ArpTable::insert_address_v4(packet.interface_id, arp.sender_mac, sender_address, false);

    if opcode == 1 {
        // Don't answer with non-owned addresses
        if let Some((mac, true)) = ArpTable::lookup_cache_v4(packet.interface_id, target_address) {
            // Reply with own address
            send_reply(packet.interface_id, arp, mac).ok();
        }
    }
}
