edge-schema 0.1.0

use once_cell::sync::Lazy;
use parking_lot::Mutex;
use rand_chacha::{rand_core::RngCore, ChaCha20Rng};
use serde::{Deserialize, Serialize};
use std::{
    fmt::Display,
    net::{IpAddr, Ipv6Addr},
    str::FromStr,
};

use crate::AppId;

const NETWORK_ID_ENCRYPT_KEY: [u8; 8] = [11u8, 22u8, 33u8, 44u8, 55u8, 66u8, 77u8, 88u8];

const PREFIX_RESERVED: u64 = 0x0000_0000_0000_0000;
const PREFIX_RESERVED_END: u64 = 0x0100_0000_0000_0000;
const PREFIX_APP_ID: u64 = 0x0100_0000_0000_0000;
const PREFIX_APP_ID_END: u64 = 0x0200_0000_0000_0000;
const PREFIX_RANDOM: u64 = 0x0200_0000_0000_0000;
const PREFIX_RANDOM_END: u64 = 0x0300_0000_0000_0000;

static GLOBAL_SECURE_AND_FAST_RANDOM: Lazy<Mutex<ChaCha20Rng>> =
    Lazy::new(|| Mutex::new(rand_chacha::rand_core::SeedableRng::from_entropy()));

#[derive(
    Debug,
    Clone,
    Copy,
    PartialEq,
    Eq,
    PartialOrd,
    Ord,
    Hash,
    Serialize,
    Deserialize,
    schemars::JsonSchema,
)]
pub enum NetworkId {
    Reserved(u64),
    AppDefault(AppId),
    Random(u64),
    Unknown(u64),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub enum NetworkIdEncodingMethod {
    PrivateProjection,
    PublicProjection,
}

impl NetworkId {
    pub fn new(val: u64) -> Self {
        match val {
            mut val if val < PREFIX_RESERVED_END => {
                val -= PREFIX_RESERVED;
                Self::Reserved(val & 0xFF_FFFF_FFFF_FFFF) // 2^56 bits
            }
            mut val if val >= PREFIX_APP_ID && val < PREFIX_APP_ID_END => {
                val -= PREFIX_APP_ID;
                Self::AppDefault(AppId::new(val & 0xFF_FFFF_FFFF_FFFF)) // 2^56 bits
            }
            mut val if val >= PREFIX_RANDOM && val < PREFIX_RANDOM_END => {
                val -= PREFIX_RANDOM;
                Self::Random(val & 0xFF_FFFF_FFFF_FFFF) // 2^56 bits
            }
            val => Self::Unknown(val),
        }
    }

    pub fn new_random() -> Self {
        let val = GLOBAL_SECURE_AND_FAST_RANDOM.lock().next_u64();
        NetworkId::Random(val & 0xFF_FFFF_FFFF_FFFF)
    }

    pub fn as_u64(&self) -> u64 {
        match self {
            NetworkId::Reserved(id) => *id + PREFIX_RESERVED,
            NetworkId::AppDefault(id) => id.as_u64() + PREFIX_APP_ID,
            NetworkId::Random(id) => *id + PREFIX_RANDOM,
            NetworkId::Unknown(id) => *id,
        }
    }

    pub fn as_app_id(&self) -> Option<AppId> {
        match self {
            NetworkId::AppDefault(id) => Some(*id),
            _ => None,
        }
    }

    pub fn to_bytes(&self) -> [u8; 8] {
        self.as_u64().to_be_bytes()
    }

    pub fn from_ip(addr: &IpAddr, method: NetworkIdEncodingMethod) -> Option<(Self, u32)> {
        match addr {
            IpAddr::V4(_) => None,
            IpAddr::V6(addr) => Some(Self::from_ipv6(addr, method)),
        }
    }

    pub fn from_ipv6(addr: &Ipv6Addr, method: NetworkIdEncodingMethod) -> (Self, u32) {
        let addr = addr.octets();
        let e: &[u8] = match method {
            NetworkIdEncodingMethod::PrivateProjection => &addr[5..13],
            NetworkIdEncodingMethod::PublicProjection => &addr[8..16],
        };
        let mut encrypted: [u8; 8] = [e[0], e[1], e[2], e[3], e[4], e[5], e[6], e[7]];

        // Grab the seed from the upper bits of the IP address
        let iv = match method {
            NetworkIdEncodingMethod::PrivateProjection => {
                let iv = &addr[1..5];
                [iv[0], iv[1], iv[2], iv[3], iv[0], iv[1], iv[2], iv[3]]
            }
            NetworkIdEncodingMethod::PublicProjection => {
                let iv = &addr[0..8];
                [iv[0], iv[1], iv[2], iv[3], iv[4], iv[5], iv[6], iv[7]]
            }
        };

        // Use a cipher to decrypt the ID
        let key = [
            iv[0],
            NETWORK_ID_ENCRYPT_KEY[0],
            iv[1],
            NETWORK_ID_ENCRYPT_KEY[1],
            iv[2],
            NETWORK_ID_ENCRYPT_KEY[2],
            iv[3],
            NETWORK_ID_ENCRYPT_KEY[3],
            iv[4],
            NETWORK_ID_ENCRYPT_KEY[4],
            iv[5],
            NETWORK_ID_ENCRYPT_KEY[5],
            iv[6],
            NETWORK_ID_ENCRYPT_KEY[6],
            iv[7],
            NETWORK_ID_ENCRYPT_KEY[7],
        ];
        let cipher = sparx::sparx64::key_schedule_encrypt(&key);
        sparx::sparx64::decrypt_block(&mut encrypted, &cipher);

        // Split the result into the AppID and remainder
        let remainder = match method {
            NetworkIdEncodingMethod::PrivateProjection => {
                let remainder = [0, addr[13], addr[14], addr[15]];
                u32::from_be_bytes(remainder)
            }
            NetworkIdEncodingMethod::PublicProjection => 0,
        };
        let id = u64::from_be_bytes(encrypted);
        (Self::new(id), remainder)
    }

    pub fn into_ipv6(
        self,
        base: Ipv6Addr,
        remainder: u32,
        method: NetworkIdEncodingMethod,
    ) -> Ipv6Addr {
        let b = base.octets();

        // Turn the ID into a plain text
        let mut plain = self.as_u64().to_be_bytes();

        // Grab the seed from the upper bits of the IP address
        let iv = match method {
            NetworkIdEncodingMethod::PrivateProjection => {
                let iv = &b[1..5];
                [iv[0], iv[1], iv[2], iv[3], iv[0], iv[1], iv[2], iv[3]]
            }
            NetworkIdEncodingMethod::PublicProjection => {
                let iv = &b[0..8];
                [iv[0], iv[1], iv[2], iv[3], iv[4], iv[5], iv[6], iv[7]]
            }
        };

        // Use a cipher to decrypt the ID
        let key = [
            iv[0],
            NETWORK_ID_ENCRYPT_KEY[0],
            iv[1],
            NETWORK_ID_ENCRYPT_KEY[1],
            iv[2],
            NETWORK_ID_ENCRYPT_KEY[2],
            iv[3],
            NETWORK_ID_ENCRYPT_KEY[3],
            iv[4],
            NETWORK_ID_ENCRYPT_KEY[4],
            iv[5],
            NETWORK_ID_ENCRYPT_KEY[5],
            iv[6],
            NETWORK_ID_ENCRYPT_KEY[6],
            iv[7],
            NETWORK_ID_ENCRYPT_KEY[7],
        ];
        let cipher = sparx::sparx64::key_schedule_encrypt(&key);
        sparx::sparx64::encrypt_block(&mut plain, &cipher);
        let e = plain;

        // Build the address and return it
        let remainder = remainder.to_be_bytes();
        let addr = match method {
            NetworkIdEncodingMethod::PrivateProjection => [
                b[0],
                b[1],
                b[2],
                b[3],
                b[4],
                e[0],
                e[1],
                e[2],
                e[3],
                e[4],
                e[5],
                e[6],
                e[7],
                remainder[1],
                remainder[2],
                remainder[3],
            ],
            NetworkIdEncodingMethod::PublicProjection => [
                b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], e[0], e[1], e[2], e[3], e[4], e[5],
                e[6], e[7],
            ],
        };
        addr.into()
    }
}

impl Display for NetworkId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NetworkId::Reserved(id) => write!(f, "reserved({id})"),
            NetworkId::AppDefault(id) => write!(f, "app-id({id})"),
            NetworkId::Random(id) => write!(f, "random({id:0X?})"),
            NetworkId::Unknown(id) => write!(f, "unknown({id})"),
        }
    }
}

impl TryFrom<&str> for NetworkId {
    type Error = ();

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        FromStr::from_str(value)
    }
}

impl FromStr for NetworkId {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(if let Some(app_id) = s.strip_prefix("app-id(") {
            match app_id.strip_suffix(')') {
                Some(app_id) => NetworkId::AppDefault(AppId::from_str(app_id)?),
                None => NetworkId::AppDefault(AppId::from_str(app_id)?),
            }
        } else if let Some(id) = s.strip_prefix("random(") {
            match id.strip_suffix(')') {
                Some(id) => NetworkId::Random(u64::from_str(id).map_err(|_| ())?),
                None => NetworkId::Random(u64::from_str(id).map_err(|_| ())?),
            }
        } else if let Some(id) = s.strip_prefix("reserved(") {
            match id.strip_suffix(')') {
                Some(id) => NetworkId::Reserved(u64::from_str(id).map_err(|_| ())?),
                None => NetworkId::Reserved(u64::from_str(id).map_err(|_| ())?),
            }
        } else {
            return Err(());
        })
    }
}

impl From<u64> for NetworkId {
    fn from(value: u64) -> Self {
        Self::new(value)
    }
}

impl From<NetworkId> for u64 {
    fn from(val: NetworkId) -> Self {
        val.as_u64()
    }
}

impl From<AppId> for NetworkId {
    fn from(value: AppId) -> Self {
        Self::AppDefault(value)
    }
}

impl TryInto<AppId> for NetworkId {
    type Error = ();
    fn try_into(self) -> Result<AppId, Self::Error> {
        self.as_app_id().ok_or(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_network_id_with_randoms() {
        let addrs = [
            Ipv6Addr::LOCALHOST,
            Ipv6Addr::UNSPECIFIED,
            Ipv6Addr::new(0xff0e, 0, 0, 0, 0, 0, 0, 0),
            Ipv6Addr::new(0, 0, 0x1c9, 0, 0, 0xafc8, 0, 0x1),
            Ipv6Addr::new(0xff00, 0, 0, 0, 0, 0, 0, 0),
            Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 0),
        ];
        let methods = [
            NetworkIdEncodingMethod::PrivateProjection,
            NetworkIdEncodingMethod::PublicProjection,
        ];
        let network_ids = [
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
            NetworkId::new_random(),
        ];
        let remainders = [0u32, 1u32, 128u32, 255u32, 16_777_215u32];

        for method in methods {
            for network_id in network_ids {
                for r in remainders {
                    for base in addrs.iter() {
                        let addr = network_id.into_ipv6(*base, r, method);

                        let s1 = addr.segments();
                        let s2 = base.segments();

                        match method {
                            NetworkIdEncodingMethod::PrivateProjection => {
                                assert_eq!(s1[0], s2[0]);
                                assert_eq!(s1[1], s2[1]);
                                assert_eq!(s1[2].to_be_bytes()[0], s2[2].to_be_bytes()[0]);
                            }
                            NetworkIdEncodingMethod::PublicProjection => {
                                assert_eq!(s1[0], s2[0]);
                                assert_eq!(s1[1], s2[1]);
                                assert_eq!(s1[2], s2[2]);
                                assert_eq!(s1[3], s2[3]);
                            }
                        }

                        let (ret_network_id, ret_r) = NetworkId::from_ipv6(&addr, method);
                        assert_eq!(network_id, ret_network_id);

                        match method {
                            NetworkIdEncodingMethod::PrivateProjection => assert_eq!(r, ret_r),
                            NetworkIdEncodingMethod::PublicProjection => assert_eq!(0, ret_r),
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn test_network_id_with_app_id() {
        let app_ids = [0u64, 1u64, 1000u64, 100000000u64, 1000500200u64];
        for test_app_id in app_ids {
            let app_id = AppId::new(test_app_id);
            let network_id: NetworkId = app_id.into();
            let app_id: AppId = network_id.try_into().unwrap();
            let app_id = app_id.as_u64();
            assert_eq!(app_id, test_app_id)
        }
    }
}