// Copyright (c) 2025, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file
use std::net::{IpAddr, SocketAddr};

use ipnet::IpNet;

use crate::model::network_address::Address::{IPv4, IPv6};
use crate::IpFilterError;

/// Represents an IP address.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Address<Addr> {
    /// Represents an IPv4 address.
    IPv4(Addr),
    /// Represents an IPv6 address.
    IPv6(Addr),
    /// Represents an unknown address.
    Unknown,
}

impl Address<NetworkAddress> {
    /// Checks if the current address contains the other address.
    pub fn contains(&self, other: &Address<NetworkAddress>) -> bool {
        *self == *other || self.match_subnet(other)
    }

    fn match_subnet(&self, other: &Address<NetworkAddress>) -> bool {
        match (self, other) {
            (IPv4(addr1), IPv4(addr2)) => addr1.contains(addr2),
            (IPv6(addr1), IPv6(addr2)) => addr1.contains(addr2),
            _ => false,
        }
    }
}

/// Trait representing a valid IP address.
pub trait ValidIpAddress {
    fn is_ipv6(&self) -> bool;
    fn is_ipv4(&self) -> bool;
}

/// Represents a network address.
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
pub struct NetworkAddress {
    /// The IP address.
    address: IpNet,
}

impl From<IpAddr> for NetworkAddress {
    fn from(address: IpAddr) -> Self {
        NetworkAddress {
            address: IpNet::from(address),
        }
    }
}

impl From<IpNet> for NetworkAddress {
    fn from(network: IpNet) -> Self {
        NetworkAddress { address: network }
    }
}

impl ValidIpAddress for NetworkAddress {
    fn is_ipv6(&self) -> bool {
        self.address.addr().is_ipv6()
    }

    fn is_ipv4(&self) -> bool {
        self.address.addr().is_ipv4()
    }
}

impl NetworkAddress {
    /// Checks if the current address contains the other address.
    pub fn contains(&self, other: &NetworkAddress) -> bool {
        self.address.contains(&other.address)
    }

    /// Parses a string into a network address.
    pub fn parse(address: String) -> Result<NetworkAddress, IpFilterError> {
        let address_str = address.as_str();
        Self::socket_parser(address_str)
            .or_else(|_| Self::address_parser(address_str))
            .or_else(|_| Self::cidr_parser(address_str))
    }

    fn socket_parser(address: &str) -> Result<NetworkAddress, IpFilterError> {
        address
            .parse::<SocketAddr>()
            .map(|socket| socket.ip())
            .map(NetworkAddress::from)
            .map_err(|_| IpFilterError::InvalidIp(address.to_string()))
    }

    fn address_parser(address: &str) -> Result<NetworkAddress, IpFilterError> {
        address
            .parse::<IpAddr>()
            .map(NetworkAddress::from)
            .map_err(|_| IpFilterError::InvalidIp(address.to_string()))
    }

    fn cidr_parser(network: &str) -> Result<NetworkAddress, IpFilterError> {
        network
            .parse::<IpNet>()
            .map(NetworkAddress::from)
            .or_else(|_| {
                let complete_network_addr = complete_network_address(network.to_string());
                complete_network_addr
                    .parse::<IpNet>()
                    .map(NetworkAddress::from)
                    .map_err(|_| IpFilterError::InvalidIp(network.to_string()))
            })
    }
}

fn complete_network_address(addr: String) -> String {
    if is_ipv4(&addr) {
        complete_ipv4(&addr)
    } else if is_ipv6(&addr) {
        complete_ipv6(&addr)
    } else {
        addr
    }
}

fn complete_ipv4(addr: &str) -> String {
    complete_missing_groups(addr, 4, ".", "0", 8)
}

fn complete_ipv6(addr: &str) -> String {
    complete_missing_groups(addr, 8, ":", "0", 16)
}

// For example: we would like an address like 192.0.0 to be completed as 192.0.0.0/24. Our groups are
// octets (like 192, 0, and 0) each with a bit size of 8 (hence octets). We would like the new octet
// to be added to be 0 so we set the empty group parameter to "0" ( as a &str). We separate octets with
// "." in ipv4 so we set the separator to be "." and we have 4 octets in an ipv4 so we set the group
// count to be 4
fn complete_missing_groups(
    addr: &str,
    group_count: i32,
    separator: &str,
    empty_group: &str,
    group_bit_size: i32,
) -> String {
    let ip_group_count: i32 = group_count;
    let pattern = separator.to_string();
    let mut group: Vec<&str> = addr.split(&pattern).collect();
    let missing_group_count: i32 = ip_group_count - group.len() as i32;
    for _ in 0..missing_group_count {
        group.push(empty_group)
    }
    let mask = (group_count - missing_group_count) * group_bit_size;
    let full_addr = group.join(separator);
    if missing_group_count == 0 {
        full_addr
    } else {
        format!("{full_addr}/{mask}")
    }
}

fn is_ipv4(addr: &str) -> bool {
    addr.contains('.')
}

fn is_ipv6(addr: &str) -> bool {
    addr.contains(':')
}

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

    #[test]
    fn parse_valid_ipv4_address() {
        let result = NetworkAddress::parse("192.168.1.1".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_valid_ipv6_address() {
        let result = NetworkAddress::parse("2001:db8::1".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_valid_socket_address_ipv4() {
        let result = NetworkAddress::parse("192.168.1.1:8080".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_valid_socket_address_ipv6() {
        let result = NetworkAddress::parse("[2001:db8::1]:8080".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_valid_cidr_ipv4() {
        let result = NetworkAddress::parse("192.168.0.0/24".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_valid_cidr_ipv6() {
        let result = NetworkAddress::parse("2001:db8::/32".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_incomplete_ipv4_completes_with_mask() {
        let result = NetworkAddress::parse("192.168.0".to_string());
        assert!(result.is_ok());
    }

    #[test]
    fn parse_invalid_address_returns_error() {
        let result = NetworkAddress::parse("invalid".to_string());
        assert!(matches!(result, Err(IpFilterError::InvalidIp(_))));
    }

    #[test]
    fn parse_empty_string_returns_error() {
        let result = NetworkAddress::parse("".to_string());
        assert!(matches!(result, Err(IpFilterError::InvalidIp(_))));
    }

    #[test]
    fn contains_same_address_returns_true() {
        let addr1 = NetworkAddress::parse("192.168.1.1".to_string()).unwrap();
        let addr2 = NetworkAddress::parse("192.168.1.1".to_string()).unwrap();
        assert!(addr1.contains(&addr2));
    }

    #[test]
    fn cidr_contains_address_in_range() {
        let cidr = NetworkAddress::parse("192.168.0.0/24".to_string()).unwrap();
        let addr = NetworkAddress::parse("192.168.0.100".to_string()).unwrap();
        assert!(cidr.contains(&addr));
    }

    #[test]
    fn cidr_does_not_contain_address_outside_range() {
        let cidr = NetworkAddress::parse("192.168.0.0/24".to_string()).unwrap();
        let addr = NetworkAddress::parse("192.168.1.1".to_string()).unwrap();
        assert!(!cidr.contains(&addr));
    }

    #[test]
    fn address_enum_contains_same_ipv4() {
        let addr1 = Address::IPv4(NetworkAddress::parse("192.168.1.1".to_string()).unwrap());
        let addr2 = Address::IPv4(NetworkAddress::parse("192.168.1.1".to_string()).unwrap());
        assert!(addr1.contains(&addr2));
    }

    #[test]
    fn address_enum_contains_same_ipv6() {
        let addr1 = Address::IPv6(NetworkAddress::parse("::1".to_string()).unwrap());
        let addr2 = Address::IPv6(NetworkAddress::parse("::1".to_string()).unwrap());
        assert!(addr1.contains(&addr2));
    }

    #[test]
    fn address_enum_ipv4_does_not_contain_ipv6() {
        let ipv4 = Address::IPv4(NetworkAddress::parse("192.168.1.1".to_string()).unwrap());
        let ipv6 = Address::IPv6(NetworkAddress::parse("::1".to_string()).unwrap());
        assert!(!ipv4.contains(&ipv6));
    }

    #[test]
    fn address_enum_ipv6_does_not_contain_ipv4() {
        let ipv6 = Address::IPv6(NetworkAddress::parse("::1".to_string()).unwrap());
        let ipv4 = Address::IPv4(NetworkAddress::parse("192.168.1.1".to_string()).unwrap());
        assert!(!ipv6.contains(&ipv4));
    }

    #[test]
    fn address_enum_cidr_contains_address_in_range() {
        let cidr = Address::IPv4(NetworkAddress::parse("192.168.0.0/24".to_string()).unwrap());
        let addr = Address::IPv4(NetworkAddress::parse("192.168.0.100".to_string()).unwrap());
        assert!(cidr.contains(&addr));
    }

    #[test]
    fn address_enum_unknown_does_not_contain_ipv4() {
        let unknown = Address::<NetworkAddress>::Unknown;
        let ipv4 = Address::IPv4(NetworkAddress::parse("192.168.1.1".to_string()).unwrap());
        assert!(!unknown.contains(&ipv4));
    }

    // ValidIpAddress trait tests
    #[test]
    fn network_address_is_ipv4() {
        let addr = NetworkAddress::parse("192.168.1.1".to_string()).unwrap();
        assert!(addr.is_ipv4());
        assert!(!addr.is_ipv6());
    }

    #[test]
    fn network_address_is_ipv6() {
        let addr = NetworkAddress::parse("::1".to_string()).unwrap();
        assert!(addr.is_ipv6());
        assert!(!addr.is_ipv4());
    }

    #[test]
    fn network_address_from_ipaddr_ipv4() {
        let ip: std::net::IpAddr = "192.168.1.1".parse().unwrap();
        let addr = NetworkAddress::from(ip);
        assert!(addr.is_ipv4());
    }

    #[test]
    fn network_address_from_ipaddr_ipv6() {
        let ip: std::net::IpAddr = "::1".parse().unwrap();
        let addr = NetworkAddress::from(ip);
        assert!(addr.is_ipv6());
    }

    #[test]
    fn network_address_from_ipnet() {
        let net: ipnet::IpNet = "192.168.0.0/24".parse().unwrap();
        let addr = NetworkAddress::from(net);
        assert!(addr.is_ipv4());
    }
}
