use hmac::{Hmac, Mac};
use prost::Message;
use sha2::Sha256;
#[cfg(feature = "api")]
use std::collections::BTreeMap;
use std::collections::BTreeSet;

type PeerId = u32;

include!(concat!(env!("OUT_DIR"), "/peer_rpc.rs"));
#[cfg(feature = "json-rpc")]
include!(concat!(env!("OUT_DIR"), "/peer_rpc.serde.rs"));

impl PeerGroupInfo {
    pub fn generate_with_proof(group_name: String, group_secret: String, peer_id: PeerId) -> Self {
        let mut mac = Hmac::<Sha256>::new_from_slice(group_secret.as_bytes())
            .expect("HMAC can take key of any size");

        let mut data_to_sign = group_name.as_bytes().to_vec();
        data_to_sign.push(0x00); // Add a delimiter byte
        data_to_sign.extend_from_slice(&peer_id.to_be_bytes());

        mac.update(&data_to_sign);

        let proof = mac.finalize().into_bytes().to_vec();

        PeerGroupInfo {
            group_name,
            group_proof: proof,
        }
    }

    pub fn verify(&self, group_secret: &str, peer_id: PeerId) -> bool {
        let mut verifier = Hmac::<Sha256>::new_from_slice(group_secret.as_bytes())
            .expect("HMAC can take key of any size");

        let mut data_to_sign = self.group_name.as_bytes().to_vec();
        data_to_sign.push(0x00); // Add a delimiter byte
        data_to_sign.extend_from_slice(&peer_id.to_be_bytes());

        verifier.update(&data_to_sign);

        verifier.verify_slice(&self.group_proof).is_ok()
    }
}

impl TrustedCredentialPubkeyProof {
    pub fn generate_credential_hmac_from_bytes(
        credential_bytes: &[u8],
        network_secret: &str,
    ) -> Vec<u8> {
        let mut mac = Hmac::<Sha256>::new_from_slice(network_secret.as_bytes())
            .expect("HMAC can take key of any size");
        mac.update(b"easytier credential proof");
        mac.update(credential_bytes);
        mac.finalize().into_bytes().to_vec()
    }

    pub fn generate_credential_hmac(
        credential: &TrustedCredentialPubkey,
        network_secret: &str,
    ) -> Vec<u8> {
        Self::generate_credential_hmac_from_bytes(&credential.encode_to_vec(), network_secret)
    }

    pub fn new_signed(credential: TrustedCredentialPubkey, network_secret: &str) -> Self {
        let credential_hmac = Self::generate_credential_hmac(&credential, network_secret);
        Self {
            credential: Some(credential),
            credential_hmac,
        }
    }

    pub fn verify_credential_hmac(&self, network_secret: &str) -> bool {
        let Some(credential) = self.credential.as_ref() else {
            return false;
        };
        self.verify_credential_hmac_with_bytes(&credential.encode_to_vec(), network_secret)
    }

    pub fn verify_credential_hmac_with_bytes(
        &self,
        credential_bytes: &[u8],
        network_secret: &str,
    ) -> bool {
        if self.credential_hmac.is_empty() {
            return false;
        }

        let mut mac = Hmac::<Sha256>::new_from_slice(network_secret.as_bytes())
            .expect("HMAC can take key of any size");
        mac.update(b"easytier credential proof");
        mac.update(credential_bytes);
        mac.verify_slice(&self.credential_hmac).is_ok()
    }
}

impl From<RouteConnBitmap> for sync_route_info_request::ConnInfo {
    fn from(val: RouteConnBitmap) -> Self {
        Self::ConnBitmap(val)
    }
}

impl From<RouteConnPeerList> for sync_route_info_request::ConnInfo {
    fn from(val: RouteConnPeerList) -> Self {
        Self::ConnPeerList(val)
    }
}

#[cfg(feature = "api")]
impl From<Vec<crate::api::instance::PeerInfo>> for PeerInfoForGlobalMap {
    fn from(peers: Vec<crate::api::instance::PeerInfo>) -> Self {
        let mut peer_map = BTreeMap::new();
        for peer in peers {
            let Some(min_lat) = peer
                .conns
                .iter()
                .map(|conn| conn.stats.as_ref().unwrap().latency_us)
                .min()
            else {
                continue;
            };

            let dp_info = DirectConnectedPeerInfo {
                latency_ms: std::cmp::max(1, (min_lat as u32 / 1000) as i32),
            };

            peer_map.insert(peer.peer_id, dp_info);
        }
        PeerInfoForGlobalMap {
            direct_peers: peer_map,
        }
    }
}

impl From<RoutePeerInfo> for crate::core_peer::peer::Route {
    fn from(val: RoutePeerInfo) -> Self {
        let network_length = if val.network_length == 0 {
            24
        } else {
            val.network_length
        };

        crate::core_peer::peer::Route {
            peer_id: val.peer_id,
            ipv4_addr: val.ipv4_addr.map(|ipv4_addr| crate::common::Ipv4Inet {
                address: Some(ipv4_addr),
                network_length,
            }),
            next_hop_peer_id: 0,
            cost: 0,
            path_latency: 0,
            proxy_cidrs: val.proxy_cidrs.clone(),
            hostname: val.hostname.unwrap_or_default(),
            stun_info: {
                let mut stun_info = crate::common::StunInfo::default();
                if let Ok(udp_nat_type) = crate::common::NatType::try_from(val.udp_nat_type) {
                    stun_info.set_udp_nat_type(udp_nat_type);
                }
                if let Ok(tcp_nat_type) = crate::common::NatType::try_from(val.tcp_nat_type) {
                    stun_info.set_tcp_nat_type(tcp_nat_type);
                }
                Some(stun_info)
            },
            inst_id: val.inst_id.map(|x| x.to_string()).unwrap_or_default(),
            version: val.easytier_version,
            feature_flag: val.feature_flag,

            next_hop_peer_id_latency_first: None,
            cost_latency_first: None,
            path_latency_latency_first: None,

            ipv6_addr: val.ipv6_addr,
            public_ipv6_addr: val.ipv6_public_addr_lease,
            ipv6_public_addr_prefix: val.ipv6_public_addr_prefix,
        }
    }
}

#[cfg(feature = "api")]
impl From<RoutePeerInfo> for crate::api::instance::Route {
    fn from(val: RoutePeerInfo) -> Self {
        let network_length = if val.network_length == 0 {
            24
        } else {
            val.network_length
        };

        crate::api::instance::Route {
            peer_id: val.peer_id,
            ipv4_addr: val.ipv4_addr.map(|ipv4_addr| crate::common::Ipv4Inet {
                address: Some(ipv4_addr),
                network_length,
            }),
            next_hop_peer_id: 0,
            cost: 0,
            path_latency: 0,
            proxy_cidrs: val.proxy_cidrs.clone(),
            hostname: val.hostname.unwrap_or_default(),
            stun_info: {
                let mut stun_info = crate::common::StunInfo::default();
                if let Ok(udp_nat_type) = crate::common::NatType::try_from(val.udp_nat_type) {
                    stun_info.set_udp_nat_type(udp_nat_type);
                }
                if let Ok(tcp_nat_type) = crate::common::NatType::try_from(val.tcp_nat_type) {
                    stun_info.set_tcp_nat_type(tcp_nat_type);
                }
                Some(stun_info)
            },
            inst_id: val.inst_id.map(|x| x.to_string()).unwrap_or_default(),
            version: val.easytier_version,
            feature_flag: val.feature_flag,

            next_hop_peer_id_latency_first: None,
            cost_latency_first: None,
            path_latency_latency_first: None,

            ipv6_addr: val.ipv6_addr,
            public_ipv6_addr: val.ipv6_public_addr_lease,
            ipv6_public_addr_prefix: val.ipv6_public_addr_prefix,
        }
    }
}

impl RouteConnBitmap {
    pub fn get_bit(&self, idx: usize) -> bool {
        let byte_idx = idx / 8;
        let bit_idx = idx % 8;
        let byte = self.bitmap[byte_idx];
        (byte >> bit_idx) & 1 == 1
    }

    pub fn get_connected_peers(&self, peer_idx: usize) -> BTreeSet<PeerId> {
        let mut connected_peers = BTreeSet::new();
        for (idx, peer_id_version) in self.peer_ids.iter().enumerate() {
            if self.get_bit(peer_idx * self.peer_ids.len() + idx) {
                connected_peers.insert(peer_id_version.peer_id);
            }
        }
        connected_peers
    }
}

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

    #[test]
    fn test_peer_group_info_new() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name.clone(), group_secret, peer_id);

        assert_eq!(peer_group_info.group_name, group_name);
        assert!(!peer_group_info.group_proof.is_empty());
        // HMAC-SHA256 produces a 32-byte output
        assert_eq!(peer_group_info.group_proof.len(), 32);
    }

    #[test]
    fn test_peer_group_info_verify_valid() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        // Verification should succeed using the same secret and peer_id
        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    fn test_peer_group_info_verify_invalid_secret() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let peer_group_info = PeerGroupInfo::generate_with_proof(group_name, group_secret, peer_id);

        // Verification should fail with a wrong secret
        assert!(!peer_group_info.verify("wrong_secret", peer_id));
    }

    #[test]
    fn test_peer_group_info_verify_invalid_peer_id() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        // Verification should fail with a wrong peer_id
        assert!(!peer_group_info.verify(&group_secret, 999u32));
    }

    #[test]
    fn test_peer_group_info_different_groups_different_proofs() {
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let group1 =
            PeerGroupInfo::generate_with_proof("group1".to_string(), group_secret.clone(), peer_id);
        let group2 =
            PeerGroupInfo::generate_with_proof("group2".to_string(), group_secret, peer_id);

        // Different group names should produce different proofs
        assert_ne!(group1.group_proof, group2.group_proof);
    }

    #[test]
    fn test_peer_group_info_same_params_same_proof() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let group1 =
            PeerGroupInfo::generate_with_proof(group_name.clone(), group_secret.clone(), peer_id);
        let group2 = PeerGroupInfo::generate_with_proof(group_name, group_secret, peer_id);

        // Same parameters should produce the same proof
        assert_eq!(group1.group_proof, group2.group_proof);
    }

    #[test]
    fn test_peer_group_info_empty_group_name() {
        let group_name = "".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name.clone(), group_secret.clone(), peer_id);

        assert_eq!(peer_group_info.group_name, group_name);
        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    fn test_peer_group_info_empty_secret() {
        let group_name = "test_group".to_string();
        let group_secret = "".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    fn test_peer_group_info_unicode_group_name() {
        let group_name = "测试组🚀".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name.clone(), group_secret.clone(), peer_id);

        assert_eq!(peer_group_info.group_name, group_name);
        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    fn test_peer_group_info_unicode_secret() {
        let group_name = "test_group".to_string();
        let group_secret = "密码123🔐".to_string();
        let peer_id = 42u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    fn test_peer_group_info_zero_peer_id() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 0u32;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    fn test_peer_group_info_max_peer_id() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = u32::MAX;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        assert!(peer_group_info.verify(&group_secret, peer_id));
    }

    #[test]
    #[ignore]
    fn perf_test_generate_with_proof() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;
        let iterations = 100000;

        let start = std::time::Instant::now();
        for _ in 0..iterations {
            let _ = PeerGroupInfo::generate_with_proof(
                group_name.clone(),
                group_secret.clone(),
                peer_id,
            );
        }
        let duration = start.elapsed();

        println!(
            "generate_with_proof took {:?} for {} iterations",
            duration, iterations
        );
        println!("Avg time per iteration: {:?}", duration / iterations as u32);
    }

    #[test]
    #[ignore]
    fn perf_test_verify() {
        let group_name = "test_group".to_string();
        let group_secret = "secret123".to_string();
        let peer_id = 42u32;
        let iterations = 100000;

        let peer_group_info =
            PeerGroupInfo::generate_with_proof(group_name, group_secret.clone(), peer_id);

        let start = std::time::Instant::now();
        for _ in 0..iterations {
            assert!(peer_group_info.verify(&group_secret, peer_id));
        }
        let duration = start.elapsed();

        println!("verify took {:?} for {} iterations", duration, iterations);
        println!("Avg time per iteration: {:?}", duration / iterations as u32);
    }

    #[test]
    fn test_trusted_credential_pubkey_hmac_valid() {
        let credential = TrustedCredentialPubkey {
            pubkey: vec![7u8; 32],
            groups: vec!["ops".to_string(), "guest".to_string()],
            allow_relay: true,
            expiry_unix: 123456,
            allowed_proxy_cidrs: vec!["10.0.0.0/24".to_string()],
            reusable: Some(true),
        };
        let tc = TrustedCredentialPubkeyProof::new_signed(credential, "sec-1");

        assert!(tc.verify_credential_hmac("sec-1"));
        assert!(!tc.verify_credential_hmac("sec-2"));
    }

    #[test]
    fn test_trusted_credential_pubkey_hmac_tampered() {
        let credential = TrustedCredentialPubkey {
            pubkey: vec![8u8; 32],
            groups: vec!["g1".to_string()],
            allow_relay: false,
            expiry_unix: 1,
            allowed_proxy_cidrs: vec![],
            reusable: Some(true),
        };
        let tc = TrustedCredentialPubkeyProof::new_signed(credential, "sec-1");

        let mut tampered = tc.clone();
        tampered.credential.as_mut().unwrap().allow_relay = true;
        assert!(!tampered.verify_credential_hmac("sec-1"));
    }

    #[test]
    fn test_trusted_credential_pubkey_hmac_with_raw_bytes() {
        let credential = TrustedCredentialPubkey {
            pubkey: vec![9u8; 32],
            groups: vec!["raw".to_string()],
            allow_relay: true,
            expiry_unix: 123456,
            allowed_proxy_cidrs: vec![],
            reusable: Some(true),
        };

        let mut raw_credential_bytes = credential.encode_to_vec();
        prost::encoding::encode_key(
            9999,
            prost::encoding::WireType::Varint,
            &mut raw_credential_bytes,
        );
        prost::encoding::encode_varint(42, &mut raw_credential_bytes);

        let proof = TrustedCredentialPubkeyProof {
            credential: Some(credential),
            credential_hmac: TrustedCredentialPubkeyProof::generate_credential_hmac_from_bytes(
                &raw_credential_bytes,
                "sec-1",
            ),
        };

        assert!(proof.verify_credential_hmac_with_bytes(&raw_credential_bytes, "sec-1"));
        assert!(!proof.verify_credential_hmac("sec-1"));
    }
}
