use std::sync::atomic::Ordering;

use dashmap::DashMap;
use etherparse::Ethernet2Header;
use once_cell::sync::OnceCell;
use sdlan_sn_rs::{
    config::SDLAN_DEFAULT_TTL,
    utils::{aes_encrypt, get_current_timestamp, ip_to_string, Mac},
};

use tracing::debug;

use tracing::error;

use crate::{
    network::{form_ethernet_packet, send_packet_to_net},
    pb::{encode_to_udp_message, SdlData},
    tcp::PacketType,
    utils::mac_to_string,
};

use super::get_edge;

pub const MAX_WAIT_PACKETS: usize = 100;

pub trait TunTapPacketHandler {
    async fn handle_packet_from_net(&self, data: &[u8], key: &[u8]) -> std::io::Result<()>;
    async fn handle_packet_from_device(&self, data: Vec<u8>, key: &[u8]) -> std::io::Result<()>;
}

static ARP_WAIT_LIST: OnceCell<ArpWaitList> = OnceCell::new();

pub fn init_arp_wait_list() {
    let waitlist = ArpWaitList {
        content: DashMap::new(),
    };
    ARP_WAIT_LIST.set(waitlist).unwrap();
}

#[derive(Debug)]
pub struct ArpWaitInfo {
    timestamp: u64,
    // origin data is from the tun or tap device
    origin_data: Vec<u8>,
}

#[derive(Debug)]
pub struct ArpWaitList {
    content: DashMap<u32, Vec<ArpWaitInfo>>,
}

impl ArpWaitList {
    fn add_to_wait_list(&self, ip: u32, origin_data: Vec<u8>) {
        let mut entry = self.content.entry(ip).or_insert(vec![]);
        if entry.len() < MAX_WAIT_PACKETS {
            entry.push(ArpWaitInfo {
                timestamp: get_current_timestamp(),
                origin_data,
            })
        }
    }

    async fn arp_arrived(&self, ip: u32, mac: Mac) {
        debug!(
            "arp for {} arrived: {}",
            ip_to_string(&ip),
            mac_to_string(&mac)
        );
        let Some(items) = self.content.remove(&ip) else {
            return;
        };
        let edge = get_edge();
        // just remove the items
        if !edge.is_authorized() {
            return;
        }

        let encrypt_key = edge.get_encrypt_key();
        let network_id = edge.network_id.load(Ordering::Relaxed);

        let src_mac = edge.device_config.get_mac();
        let now = get_current_timestamp();
        for item in items.1 {
            if (now - item.timestamp) > 5 {
                continue;
            }
            let packet = form_ethernet_packet(src_mac, mac, &item.origin_data);

            let pkt_size = packet.len();

            let Ok(encrypted) = aes_encrypt(&encrypt_key, &packet) else {
                error!("failed to encrypt packet request");
                return;
            };
            let data = SdlData {
                is_p2p: true,
                network_id,
                ttl: SDLAN_DEFAULT_TTL as u32,
                src_mac: Vec::from(src_mac),
                dst_mac: Vec::from(mac),
                data: Vec::from(encrypted),
            };
            let msg = encode_to_udp_message(Some(data), PacketType::Data as u8).unwrap();
            send_packet_to_net(edge, mac, &msg, pkt_size as u64).await;
        }
    }
}

pub fn add_to_arp_wait_list(ip: u32, origin_data: Vec<u8>) {
    let waitlist = ARP_WAIT_LIST
        .get()
        .expect("ARP_WAIT_LIST has not been inited");

    waitlist.add_to_wait_list(ip, origin_data);
}

pub async fn arp_arrived(ip: u32, mac: Mac) {
    let waitlist = ARP_WAIT_LIST
        .get()
        .expect("ARP_WAIT_LIST has not been inited");

    waitlist.arp_arrived(ip, mac).await;
}