use std::{
    ffi::{CStr, CString},
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    ptr::NonNull,
};

use nftnl::{
    Batch, Chain, ChainType, FinalizedBatch, Hook, MsgType, ProtoFamily, Rule, Table,
    expr::{Expression, Immediate, InterfaceName, Nat, NatType, Register},
    nft_expr, nft_nlmsg_maxsize,
    nftnl_sys::{
        NFTNL_SET_ELEM_FLAGS, NFTNL_SET_ELEM_KEY, NFTNL_SET_FLAGS, nftnl_set, nftnl_set_elem_add,
        nftnl_set_elem_alloc, nftnl_set_elem_set, nftnl_set_elem_set_u32, nftnl_set_set_u32,
    },
    set::{Set, SetKey},
};

use super::{Address, FilterRule, Policy, Port, Protocol, State, get_set_id};
use crate::enterprise::firewall::{FirewallError, SnatBinding, iprange::IpAddrRange, max_address};

const FILTER_TABLE: &str = "filter";
const NAT_TABLE: &str = "nat";
const DEFGUARD_TABLE: &str = "DEFGUARD-";
const POSTROUTING_CHAIN: &str = "POSTROUTING";
const FORWARD_CHAIN: &str = "FORWARD";
const ANON_SET_NAME: &CStr = c"__set%d";
const LOOPBACK_IFACE: &str = "lo";

const POSTROUTING_PRIORITY: i32 = 100;
const FORWARD_PRIORITY: i32 = 0;

struct InetService(u16);

impl SetKey for InetService {
    const LEN: u32 = 2;
    const TYPE: u32 = 13;

    fn data(&self) -> Box<[u8]> {
        Box::new(self.0.to_be_bytes())
    }
}

impl State {
    const fn to_expr_state(&self) -> nftnl::expr::ct::States {
        match self {
            Self::Established => nftnl::expr::ct::States::ESTABLISHED,
            Self::Invalid => nftnl::expr::ct::States::INVALID,
            Self::New => nftnl::expr::ct::States::NEW,
            Self::Related => nftnl::expr::ct::States::RELATED,
        }
    }
}

impl Protocol {
    pub(crate) fn as_port_payload_expr(&self) -> Result<&impl Expression, FirewallError> {
        match self {
            Self::Tcp => Ok(&nft_expr!(payload tcp dport)),
            Self::Udp => Ok(&nft_expr!(payload udp dport)),
            _ => Err(FirewallError::UnsupportedProtocol(*self as u8)),
        }
    }
}

impl From<Policy> for nftnl::Policy {
    fn from(policy: Policy) -> Self {
        match policy {
            Policy::Allow => Self::Accept,
            Policy::Deny => Self::Drop,
        }
    }
}

impl SetKey for Protocol {
    const LEN: u32 = 1;
    const TYPE: u32 = 12;

    fn data(&self) -> Box<[u8]> {
        Box::new([*self as u8])
    }
}

pub trait FirewallRule {
    fn to_chain_rule<'a>(
        &self,
        chain: &'a Chain,
        batch: &mut Batch,
    ) -> Result<Rule<'a>, FirewallError>;
}

fn add_address_to_set(set: NonNull<nftnl_set>, ip: &Address) -> Result<(), FirewallError> {
    match ip {
        Address::Network(network) => {
            let upper_bound = max_address(network);
            let net = network.network();
            match (net, upper_bound) {
                (IpAddr::V4(network), IpAddr::V4(upper_bound)) => {
                    add_to_set(set, &network, Some(&upper_bound))?;
                }
                (IpAddr::V6(network), IpAddr::V6(upper_bound)) => {
                    add_to_set(set, &network, Some(&upper_bound))?;
                }
                _ => {
                    return Err(FirewallError::InvalidConfiguration(format!(
                        "Expected both addresses to be of the same type, got {net:?} and \
                        {upper_bound:?}",
                    )));
                }
            }
        }
        Address::Range(addr_range) => match addr_range {
            IpAddrRange::V4(ipv4_range) => {
                add_to_set(set, ipv4_range.start(), Some(ipv4_range.end()))?;
            }
            IpAddrRange::V6(ipv6_range) => {
                add_to_set(set, ipv6_range.start(), Some(ipv6_range.end()))?;
            }
        },
    }

    Ok(())
}

fn add_port_to_set(set: NonNull<nftnl_set>, port: Port) -> Result<(), FirewallError> {
    match port {
        Port::Any => {
            // nothing to do
        }
        Port::Single(port) => {
            let inet_service = InetService(port);
            add_to_set(set, &inet_service, Some(&inet_service))?;
        }
        Port::Range(start, end) => {
            let start = InetService(start);
            let end = InetService(end);

            add_to_set(set, &start, Some(&end))?;
        }
    }

    Ok(())
}

fn add_protocol_to_set(set: NonNull<nftnl_set>, proto: Protocol) -> Result<(), FirewallError> {
    add_to_set(set, &proto, None)?;
    Ok(())
}

fn add_rule_comment(rule: &mut Rule, comment: &str) -> Result<(), FirewallError> {
    debug!("Adding comment to nftables expression: {comment:?}");
    // Since we are interoping with C, truncate the string to 255 *bytes* (not UTF-8 characters)
    // 256 is the maximum length of a comment string in nftables, leave 1 byte for the null terminator
    let maybe_truncated_str = if comment.len() > 255 {
        warn!("Comment string {comment} is too long, truncating to 255 bytes");
        &comment[..=255]
    } else {
        comment
    };
    let comment = &CString::new(maybe_truncated_str).map_err(|e| {
        FirewallError::NetlinkError(format!(
            "Failed to create CString from string {comment}. Error: {e:?}"
        ))
    })?;
    rule.set_comment(comment);
    debug!("Added comment to nftables expression: {comment:?}");
    Ok(())
}

impl FirewallRule for FilterRule<'_> {
    fn to_chain_rule<'a>(
        &self,
        chain: &'a Chain,
        batch: &mut Batch,
    ) -> Result<Rule<'a>, FirewallError> {
        let mut rule = Rule::new(chain);
        debug!("Converting {self:?} to nftables expression");
        // Debug purposes only
        let mut matches = Vec::new();

        if !self.dest_ports.is_empty() && self.protocols.len() > 1 {
            return Err(FirewallError::InvalidConfiguration(format!(
                "Cannot specify multiple protocols with destination ports, specified \
                protocols: {:?}, destination ports: {:?}, Defguard Rule ID: {}",
                self.protocols, self.dest_ports, self.defguard_rule_id
            )));
        }

        // TODO: Reduce code duplication here
        if !self.src_ips.is_empty() {
            if self.v4 {
                let set = new_anon_set::<Ipv4Addr>(chain.get_table(), ProtoFamily::Inet, true);
                batch.add(&set, MsgType::Add);

                for ip in self.src_ips {
                    add_address_to_set(set.as_ptr(), ip)?;
                }

                // ip saddr {x.x.x.x, x.x.x.x}
                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                rule.add_expr(&nft_expr!(meta nfproto));
                rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));
                rule.add_expr(&nft_expr!(payload ipv4 saddr));

                rule.add_expr(&nft_expr!(lookup & set));
            } else {
                let set = new_anon_set::<Ipv6Addr>(chain.get_table(), ProtoFamily::Inet, true);
                batch.add(&set, MsgType::Add);

                for ip in self.src_ips {
                    add_address_to_set(set.as_ptr(), ip)?;
                }

                // ip6 saddr {x.x.x.x, x.x.x.x}
                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                rule.add_expr(&nft_expr!(meta nfproto));
                rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
                rule.add_expr(&nft_expr!(payload ipv6 saddr));

                rule.add_expr(&nft_expr!(lookup & set));
            }
            debug!(
                "Added source IP addresses match to nftables expression: {:?}",
                self.src_ips
            );
            matches.push(format!("ANY SOURCE IPs: {:?}", self.src_ips));
        }

        // TODO: Reduce code duplication here
        if !self.dest_ips.is_empty() {
            if self.v4 {
                let set = new_anon_set::<Ipv4Addr>(chain.get_table(), ProtoFamily::Inet, true);
                batch.add(&set, MsgType::Add);

                for ip in self.dest_ips {
                    add_address_to_set(set.as_ptr(), ip)?;
                }

                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                // ip daddr {x.x.x.x, x.x.x.x}
                rule.add_expr(&nft_expr!(meta nfproto));
                rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));
                rule.add_expr(&nft_expr!(payload ipv4 daddr));

                rule.add_expr(&nft_expr!(lookup & set));
            } else {
                let set = new_anon_set::<Ipv6Addr>(chain.get_table(), ProtoFamily::Inet, true);
                batch.add(&set, MsgType::Add);

                for ip in self.dest_ips {
                    add_address_to_set(set.as_ptr(), ip)?;
                }

                // ip6 daddr {x.x.x.x, x.x.x.x}
                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                rule.add_expr(&nft_expr!(meta nfproto));
                rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
                rule.add_expr(&nft_expr!(payload ipv6 daddr));

                rule.add_expr(&nft_expr!(lookup & set));
            }
            debug!(
                "Added destination IP addresses match to nftables expression: {:?}",
                self.dest_ips
            );
            matches.push(format!("ANY DEST IPs: {:?}", self.dest_ips));
        }

        if !self.protocols.is_empty() {
            // > 0 Protocols
            // 0 Ports
            if self.protocols.len() > 1 {
                let set = new_anon_set::<Protocol>(chain.get_table(), ProtoFamily::Inet, false);
                batch.add(&set, MsgType::Add);

                for proto in &self.protocols {
                    add_protocol_to_set(set.as_ptr(), *proto)?;
                }

                // <protocol> dport {x, x-x}
                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                rule.add_expr(&nft_expr!(meta nfproto));

                if self.v4 {
                    rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));
                    rule.add_expr(&nft_expr!(payload ipv4 protocol));
                } else {
                    rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
                    rule.add_expr(&nft_expr!(payload ipv6 nextheader));
                }

                rule.add_expr(&nft_expr!(lookup & set));

                debug!("Added protocol match to rule: {:?}", self.protocols);
                matches.push(format!("ANY PROTOCOLS: {:?}", self.protocols));
            }
            // 1 Protocol
            // > 0 Ports
            else if !self.dest_ports.is_empty() {
                if let Some(protocol) = self.protocols.first()
                    && protocol.supports_ports()
                {
                    let set =
                        new_anon_set::<InetService>(chain.get_table(), ProtoFamily::Inet, true);
                    batch.add(&set, MsgType::Add);

                    for port in self.dest_ports {
                        add_port_to_set(set.as_ptr(), *port)?;
                    }

                    // <protocol> dport {x, x-x}
                    set.elems_iter().for_each(|elem| {
                        batch.add(&elem, MsgType::Add);
                    });

                    rule.add_expr(&nft_expr!(meta l4proto));
                    rule.add_expr(&nft_expr!(cmp == *protocol as u8));
                    rule.add_expr(protocol.as_port_payload_expr()?);
                    rule.add_expr(&nft_expr!(lookup & set));
                }

                debug!(
                    "Added single protocol ({:?}) match and destination ports match to nftables \
                    expression: {:?}",
                    self.protocols, self.dest_ports
                );
                matches.push(format!(
                    "PROTOCOL: {:?} AND ANY DEST PORTS: {:?}",
                    self.protocols, self.dest_ports
                ));
            }
            // 1 Protocol
            // 0 Ports
            else if let Some(protocol) = self.protocols.first() {
                // ip protocol <protocol>
                rule.add_expr(&nft_expr!(meta nfproto));

                if self.v4 {
                    rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));
                    rule.add_expr(&nft_expr!(payload ipv4 protocol));
                } else {
                    rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
                    rule.add_expr(&nft_expr!(payload ipv6 nextheader));
                }

                rule.add_expr(&nft_expr!(cmp == *protocol as u8));
                debug!("Added protocol match to rule: {protocol:?}");
                matches.push(format!("SINGLE PROTOCOL: {protocol:?}"));
            }
        }

        if let Some(iifname) = &self.iifname {
            // iifname <interface>
            rule.add_expr(&nft_expr!(meta iifname));
            let exact = InterfaceName::Exact(CString::new(iifname.as_str()).unwrap());
            if self.negated_iifname {
                rule.add_expr(&nft_expr!(cmp != exact));
            } else {
                rule.add_expr(&nft_expr!(cmp == exact));
            }
            debug!("Added input interface match to rule: {iifname:?}");
            matches.push(format!("INPUT INTERFACE: {iifname:?}"));
        }

        if let Some(oifname) = &self.oifname {
            // oifname <interface>
            rule.add_expr(&nft_expr!(meta oifname));
            let exact = InterfaceName::Exact(CString::new(oifname.as_str()).unwrap());
            if self.negated_oifname {
                rule.add_expr(&nft_expr!(cmp != exact));
            } else {
                rule.add_expr(&nft_expr!(cmp == exact));
            }
            debug!("Added output interface match to rule: {oifname:?}");
            matches.push(format!("OUTPUT INTERFACE: {oifname:?}"));
        }

        if !self.states.is_empty() {
            // ct state <state1>,<state2>
            let combined_states = self
                .states
                .iter()
                .fold(0u32, |acc, state| acc | state.to_expr_state().bits());
            rule.add_expr(&nft_expr!(ct state));
            rule.add_expr(&nft_expr!(bitwise mask combined_states, xor 0u32));
            rule.add_expr(&nft_expr!(cmp != 0u32));
            debug!(
                "Added connection tracking states match to nftables expression: {:?}",
                self.states
            );
            matches.push(format!("ANY CT STATES: {:?}", self.states));
        }

        if self.counter {
            // counter
            rule.add_expr(&nft_expr!(counter));
            debug!("Added counter expression to rule");
        }

        // accept/drop
        match self.action {
            Policy::Allow => {
                rule.add_expr(&nft_expr!(verdict accept));
            }
            Policy::Deny => {
                rule.add_expr(&nft_expr!(verdict drop));
            }
        }

        // comment <comment>
        if let Some(comment_string) = &self.comment {
            add_rule_comment(&mut rule, comment_string)?;
        } else {
            debug!("No comment provided for nftables expression");
        }

        let matches = matches.join(" AND ");
        debug!("Created nftables rule with matches: {matches:?}");

        Ok(rule)
    }
}

#[derive(Debug)]
struct MasqueradeRule {
    oifname: String,
    negated_oifname: bool,
    counter: bool,
}

impl FirewallRule for MasqueradeRule {
    fn to_chain_rule<'a>(
        &self,
        chain: &'a Chain,
        _batch: &mut Batch,
    ) -> Result<Rule<'a>, FirewallError> {
        let mut rule = Rule::new(chain);

        rule.add_expr(&nft_expr!(meta oifname));
        let exact = InterfaceName::Exact(CString::new(self.oifname.as_str()).unwrap());
        if self.negated_oifname {
            rule.add_expr(&nft_expr!(cmp != exact));
        } else {
            rule.add_expr(&nft_expr!(cmp == exact));
        }

        if self.counter {
            rule.add_expr(&nft_expr!(counter));
        }

        rule.add_expr(&nft_expr!(masquerade));

        Ok(rule)
    }
}

#[derive(Debug)]
struct SnatRule<'a> {
    src_ips: &'a [Address],
    public_ip: &'a IpAddr,
    oifname: String,
    negated_oifname: bool,
    counter: bool,
    ipv4: bool,
    comment: Option<String>,
}

impl FirewallRule for SnatRule<'_> {
    fn to_chain_rule<'a>(
        &self,
        chain: &'a Chain,
        batch: &mut Batch,
    ) -> Result<Rule<'a>, FirewallError> {
        let mut rule = Rule::new(chain);

        if !self.src_ips.is_empty() {
            if self.ipv4 {
                let set = new_anon_set::<Ipv4Addr>(chain.get_table(), ProtoFamily::Inet, true);
                batch.add(&set, MsgType::Add);

                for ip in self.src_ips {
                    add_address_to_set(set.as_ptr(), ip)?;
                }

                // ip saddr {x.x.x.x, x.x.x.x}
                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                rule.add_expr(&nft_expr!(meta nfproto));
                rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));
                rule.add_expr(&nft_expr!(payload ipv4 saddr));

                rule.add_expr(&nft_expr!(lookup & set));
            } else {
                let set = new_anon_set::<Ipv6Addr>(chain.get_table(), ProtoFamily::Inet, true);
                batch.add(&set, MsgType::Add);

                for ip in self.src_ips {
                    add_address_to_set(set.as_ptr(), ip)?;
                }

                // ip6 saddr {x.x.x.x, x.x.x.x}
                set.elems_iter().for_each(|elem| {
                    batch.add(&elem, MsgType::Add);
                });

                rule.add_expr(&nft_expr!(meta nfproto));
                rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
                rule.add_expr(&nft_expr!(payload ipv6 saddr));

                rule.add_expr(&nft_expr!(lookup & set));
            }
            debug!(
                "Added source IP addresses match to nftables expression: {:?}",
                self.src_ips
            );
        }

        rule.add_expr(&nft_expr!(meta oifname));
        let exact = InterfaceName::Exact(CString::new(self.oifname.as_str()).unwrap());
        if self.negated_oifname {
            rule.add_expr(&nft_expr!(cmp != exact));
        } else {
            rule.add_expr(&nft_expr!(cmp == exact));
        }

        if self.counter {
            rule.add_expr(&nft_expr!(counter));
        }

        // determine if public IP is IPv4 or IPv6 and store the address in a register
        let family = match self.public_ip {
            IpAddr::V4(ipv4_addr) => {
                rule.add_expr(&Immediate::new(*ipv4_addr, Register::Reg1));
                ProtoFamily::Ipv4
            }
            IpAddr::V6(ipv6_addr) => {
                rule.add_expr(&Immediate::new(*ipv6_addr, Register::Reg1));
                ProtoFamily::Ipv6
            }
        };
        let snat_expr = Nat {
            nat_type: NatType::SNat,
            family,
            ip_register: Register::Reg1,
            port_register: None,
        };

        rule.add_expr(&snat_expr);

        // comment <comment>
        if let Some(comment_string) = &self.comment {
            add_rule_comment(&mut rule, comment_string)?;
        } else {
            debug!("No comment provided for nftables expression");
        }

        Ok(rule)
    }
}

// Left in case if this is ever needed
// struct JumpRule<'a> {
//     src_chain: &'a Chain<'a>,
//     dest_chain: &'a Chain<'a>,
//     oifname: Option<String>,
//     iifname: Option<String>,
//     negated_oifname: bool,
//     negated_iifname: bool,
// }

// impl<'a> JumpRule<'a> {
//     fn to_chain_rule(&self) -> Result<Rule<'a>, FirewallError> {
//         let mut rule = Rule::new(self.src_chain);

//         if let Some(iifname) = &self.iifname {
//             rule.add_expr(&nft_expr!(meta iifname));
//             let exact = InterfaceName::Exact(CString::new(iifname.as_str()).unwrap());
//             if self.negated_iifname {
//                 rule.add_expr(&nft_expr!(cmp != exact));
//             } else {
//                 rule.add_expr(&nft_expr!(cmp == exact));
//             }
//         }

//         if let Some(oifname) = &self.oifname {
//             rule.add_expr(&nft_expr!(meta oifname));
//             let exact = InterfaceName::Exact(CString::new(oifname.as_str()).unwrap());
//             if self.negated_oifname {
//                 rule.add_expr(&nft_expr!(cmp != exact));
//             } else {
//                 rule.add_expr(&nft_expr!(cmp == exact));
//             }
//         }

//         // first, match

//         rule.add_expr(&nft_expr!(counter));
//         rule.add_expr(&nft_expr!(verdict jump self.dest_chain.get_name().into()));

//         Ok(rule)
//     }
// }

/// Sets up the default chains for the firewall
pub(super) fn init_firewall(
    initial_policy: Policy,
    defguard_fwd_chain_priority: Option<i32>,
    batch: &mut Batch,
    ifname: &str,
) {
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);

    batch.add(&table, MsgType::Add);
    batch.add(&table, MsgType::Del);
    batch.add(&table, MsgType::Add);

    let mut fw_chain = Chains::Forward.to_chain(&table);
    fw_chain.set_hook(
        Hook::Forward,
        defguard_fwd_chain_priority.unwrap_or(FORWARD_PRIORITY),
    );
    fw_chain.set_policy(initial_policy.into());
    fw_chain.set_type(ChainType::Filter);
    batch.add(&fw_chain, MsgType::Add);
}

pub(super) fn drop_table(batch: &mut Batch, ifname: &str) {
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
    batch.add(&table, MsgType::Add);
    batch.add(&table, MsgType::Del);
}

pub(super) fn drop_chain(chain: &Chains, batch: &mut Batch, ifname: &str) {
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
    let chain = chain.to_chain(&table);
    batch.add(&chain, MsgType::Add);
    batch.add(&chain, MsgType::Del);
}

/// Applies NAT rules on the specified interface for the outgoing packets
pub(super) fn set_nat_rules(
    batch: &mut Batch,
    ifname: &str,
    masquerade_enabled: bool,
    snat_bindings: &[SnatBinding],
) -> Result<(), FirewallError> {
    // cleanup existing POSTROUTING chain rules
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
    batch.add(&table, MsgType::Add);

    drop_chain(&Chains::Postrouting, batch, ifname);

    // initialize new POSTROUTING chain
    let mut nat_chain = Chains::Postrouting.to_chain(&table);
    nat_chain.set_hook(Hook::PostRouting, POSTROUTING_PRIORITY);
    nat_chain.set_policy(nftnl::Policy::Accept);
    nat_chain.set_type(ChainType::Nat);
    batch.add(&nat_chain, MsgType::Add);

    // add SNAT bindings
    for binding in snat_bindings {
        let snat_rule = SnatRule {
            oifname: LOOPBACK_IFACE.to_string(),
            negated_oifname: true,
            counter: true,
            src_ips: &binding.source_addrs,
            public_ip: &binding.public_ip,
            ipv4: binding.public_ip.is_ipv4(),
            comment: binding.comment.clone(),
        }
        .to_chain_rule(&nat_chain, batch)?;

        batch.add(&snat_rule, MsgType::Add);
    }

    if masquerade_enabled {
        let masquerade_rule = MasqueradeRule {
            oifname: LOOPBACK_IFACE.to_string(),
            negated_oifname: true,
            counter: true,
        }
        .to_chain_rule(&nat_chain, batch)?;

        batch.add(&masquerade_rule, MsgType::Add);
    }

    Ok(())
}

// pub(super) fn set_default_policy(
//     policy: Policy,
//     batch: &mut Batch,
//     ifname: &str,
// ) -> Result<(), FirewallError> {
//     let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
//     batch.add(&table, MsgType::Add);

//     let mut forward_chain = Chains::Forward.to_chain(&table);
//     forward_chain.set_policy(if policy == Policy::Allow {
//         Policy::Accept
//     } else {
//         Policy::Drop
//     });
//     batch.add(&forward_chain, MsgType::Add);

//     Ok(())
// }

pub(super) fn allow_established_traffic(
    batch: &mut Batch,
    ifname: &str,
) -> Result<(), FirewallError> {
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
    batch.add(&table, MsgType::Add);

    let forward_chain = Chains::Forward.to_chain(&table);
    batch.add(&forward_chain, MsgType::Add);

    let established_rule = FilterRule {
        states: vec![State::Established, State::Related],
        counter: true,
        action: Policy::Allow,
        comment: Some("Allow established and related traffic".to_string()),
        ..Default::default()
    }
    .to_chain_rule(&forward_chain, batch)?;
    batch.add(&established_rule, MsgType::Add);

    Ok(())
}

pub(super) fn ignore_unrelated_traffic(
    batch: &mut Batch,
    ifname: &str,
) -> Result<(), FirewallError> {
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
    batch.add(&table, MsgType::Add);

    let forward_chain = Chains::Forward.to_chain(&table);
    batch.add(&forward_chain, MsgType::Add);

    let ignore_rule = FilterRule {
        iifname: Some(ifname.to_string()),
        negated_iifname: true,
        action: Policy::Allow,
        counter: true,
        comment: Some("Ignore traffic not related to the VPN".to_string()),
        ..Default::default()
    }
    .to_chain_rule(&forward_chain, batch)?;
    batch.add(&ignore_rule, MsgType::Add);

    Ok(())
}

#[allow(dead_code)]
enum Tables {
    Filter(ProtoFamily),
    Nat(ProtoFamily),
    Defguard(ProtoFamily),
}

impl Tables {
    fn to_table(&self, ifname: &str) -> Table {
        match self {
            Self::Filter(family) => Table::new(
                CString::new(FILTER_TABLE)
                    .expect("Failed to create CString from FILTER_TABLE constant."),
                *family,
            ),
            Self::Nat(family) => Table::new(
                CString::new(NAT_TABLE).expect("Failed to create CString from NAT_TABLE constant."),
                *family,
            ),
            Self::Defguard(family) => Table::new(
                CString::new(DEFGUARD_TABLE.to_owned() + ifname)
                    .expect("Failed to create CString from DEFGUARD_TABLE constant."),
                *family,
            ),
        }
    }
}

pub(super) enum Chains {
    Forward,
    Postrouting,
}

impl Chains {
    fn to_chain<'a>(&self, table: &'a Table) -> Chain<'a> {
        match self {
            Self::Forward => Chain::new(
                CString::new(FORWARD_CHAIN)
                    .expect("Failed to create CString from FORWARD_CHAIN constant."),
                table,
            ),
            Self::Postrouting => Chain::new(
                CString::new(POSTROUTING_CHAIN)
                    .expect("Failed to create CString from POSTROUTING_CHAIN constant."),
                table,
            ),
        }
    }
}

pub(super) fn apply_filter_rules(
    rules: &[FilterRule],
    ifname: &str,
    socket: &mnl::Socket,
) -> Result<(), FirewallError> {
    let table = Tables::Defguard(ProtoFamily::Inet).to_table(ifname);
    let mut batch = Batch::new();
    batch.add(&table, MsgType::Add);

    let forward_chain = Chains::Forward.to_chain(&table);
    batch.add(&forward_chain, MsgType::Add);

    for rule in rules {
        let chain_rule = rule.to_chain_rule(&forward_chain, &mut batch)?;
        batch.add(&chain_rule, MsgType::Add);
    }

    send_batch(&batch.finalize(), socket)
}

pub(crate) fn send_batch(
    batch: &FinalizedBatch,
    socket: &mnl::Socket,
) -> Result<(), FirewallError> {
    socket.send_all(batch).map_err(|err| {
        FirewallError::NetlinkError(format!("Failed to send batch through socket: {err:?}"))
    })?;

    let portid = socket.portid();
    let mut buffer = vec![0; nft_nlmsg_maxsize() as usize];

    let mut expected_seqs = batch.sequence_numbers();
    while !expected_seqs.is_empty() {
        for message in socket.recv(&mut buffer).map_err(|err| {
            FirewallError::NetlinkError(format!("Failed reading message from socket: {err:?}"))
        })? {
            let Ok(message) = message else {
                warn!("Invalid netlink message");
                continue;
            };
            let Some(seq) = expected_seqs.next() else {
                warn!("Unexpected ACK in netlink messages");
                continue;
            };
            match mnl::cb_run(message, seq, portid) {
                Ok(mnl::CbResult::Stop) => {
                    debug!("Received stop signal from netlink callback");
                    break;
                }
                Ok(mnl::CbResult::Ok) => {
                    debug!("Received OK signal from netlink callback");
                }
                Err(err) => {
                    return Err(FirewallError::NetlinkError(format!(
                        "Failed to receive netlink callback: {err:?}"
                    )));
                }
            }
        }
    }

    Ok(())
}

fn new_anon_set<T>(table: &Table, family: ProtoFamily, interval_set: bool) -> Set<'_, T>
where
    T: SetKey,
{
    let set = Set::<T>::new(ANON_SET_NAME, get_set_id(), table, family);

    if interval_set {
        unsafe {
            nftnl_set_set_u32(
                set.as_ptr().as_ptr(),
                NFTNL_SET_FLAGS as u16,
                (libc::NFT_SET_ANONYMOUS | libc::NFT_SET_CONSTANT | libc::NFT_SET_INTERVAL) as u32,
            );
        }
    }

    set
}

/// Adds key to a set. If the `range_end` option is specified, it will assume the lower and upper
/// bounds of a range need to be added.
fn add_to_set<K>(
    set: NonNull<nftnl_set>,
    key: &K,
    range_end: Option<&K>,
) -> Result<(), FirewallError>
where
    K: SetKey,
{
    let key_data = key.data();
    let key_data_len = key_data.len() as u32;
    unsafe {
        let elem = nftnl_set_elem_alloc();
        if elem.is_null() {
            return Err(FirewallError::OutOfMemory(
                "Failed to allocate memory for set element".to_string(),
            ));
        }
        nftnl_set_elem_set(
            elem,
            NFTNL_SET_ELEM_KEY as u16,
            key_data.as_ptr().cast(),
            key_data_len,
        );
        nftnl_set_elem_add(set.as_ptr(), elem);

        if let Some(end) = range_end {
            let mut end_data = end.data();

            // This is a workaround to make the upper bound inclusive.
            // Perhaps there is a better way to do this.
            increment_bytes(&mut end_data);
            let end_data_len = (end_data.len()) as u32;

            let elem = nftnl_set_elem_alloc();
            if elem.is_null() {
                return Err(FirewallError::OutOfMemory(
                    "Failed to allocate memory for set element".to_string(),
                ));
            }
            nftnl_set_elem_set(
                elem,
                NFTNL_SET_ELEM_KEY as u16,
                end_data.as_ptr().cast(),
                end_data_len,
            );
            nftnl_set_elem_set_u32(
                elem,
                NFTNL_SET_ELEM_FLAGS as u16,
                libc::NFT_SET_ELEM_INTERVAL_END as u32,
            );
            nftnl_set_elem_add(set.as_ptr(), elem);
        }
    }

    Ok(())
}

fn increment_bytes(bytes: &mut [u8]) {
    for i in (0..bytes.len()).rev() {
        if bytes[i] < 255 {
            bytes[i] += 1;
            return;
        }
        bytes[i] = 0;
    }

    // the bytes have overflown, but that's okay for our purposes
}

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

    #[test]
    fn test_increment_ipv4_basic() {
        let mut ip = [192, 168, 1, 1];
        increment_bytes(&mut ip);
        assert_eq!(ip, [192, 168, 1, 2]);
    }

    #[test]
    fn test_increment_ipv4_overflow_last_octet() {
        let mut ip = [192, 168, 1, 255];
        increment_bytes(&mut ip);
        assert_eq!(ip, [192, 168, 2, 0]);
    }

    #[test]
    fn test_increment_ipv4_overflow_multiple_octets() {
        let mut ip = [192, 168, 255, 255];
        increment_bytes(&mut ip);
        assert_eq!(ip, [192, 169, 0, 0]);
    }

    #[test]
    fn test_increment_ipv4_max_address() {
        let mut ip = [255, 255, 255, 255];
        increment_bytes(&mut ip);
        assert_eq!(ip, [0, 0, 0, 0]);
    }

    #[test]
    fn test_increment_ipv4_zero_address() {
        let mut ip = [0, 0, 0, 0];
        increment_bytes(&mut ip);
        assert_eq!(ip, [0, 0, 0, 1]);
    }

    #[test]
    fn test_increment_ipv6_basic() {
        let mut ip = [0, 0, 0, 0, 0, 0, 0, 0];
        increment_bytes(&mut ip);
        assert_eq!(ip, [0, 0, 0, 0, 0, 0, 0, 1]);
    }

    #[test]
    fn test_increment_ipv6_overflow_last_octet() {
        let mut ip = [0, 0, 0, 0, 0, 0, 0, 255];
        increment_bytes(&mut ip);
        assert_eq!(ip, [0, 0, 0, 0, 0, 0, 1, 0]);
    }

    #[test]
    fn test_increment_ipv6_overflow_multiple_octets() {
        let mut ip = [0, 0, 0, 0, 0, 0, 255, 255];
        increment_bytes(&mut ip);
        assert_eq!(ip, [0, 0, 0, 0, 0, 1, 0, 0]);
    }

    #[test]
    fn test_increment_ipv6_max_address() {
        let mut ip = [255, 255, 255, 255, 255, 255, 255, 255];
        increment_bytes(&mut ip);
        assert_eq!(ip, [0, 0, 0, 0, 0, 0, 0, 0]);
    }
}
