use anyhow::{Context, Result, ensure};
use nftnl::{Batch, Chain, ChainType, Hook, MsgType, Policy, ProtoFamily, Rule, Table, nft_expr};
use std::{
    collections::HashSet,
    ffi::CString,
    io,
    net::{Ipv4Addr, Ipv6Addr},
    os::fd::AsRawFd,
};

pub const CAPTURE_MARK: u32 = 0x0200_0000;
const INCOMING_MARK: u32 = 0x0100_0000;
const TABLE: &std::ffi::CStr = c"roentgen";

fn socket() -> io::Result<mnl::Socket> {
    let socket = mnl::Socket::new(mnl::Bus::Netfilter)?;
    let timeout = libc::timeval {
        tv_sec: 5,
        tv_usec: 0,
    };
    // SAFETY: the live socket fd and timeval pointer/length match SO_RCVTIMEO.
    let result = unsafe {
        libc::setsockopt(
            socket.as_raw_fd(),
            libc::SOL_SOCKET,
            libc::SO_RCVTIMEO,
            (&timeout as *const libc::timeval).cast(),
            std::mem::size_of_val(&timeout) as _,
        )
    };
    if result != 0 {
        return Err(io::Error::last_os_error());
    }
    Ok(socket)
}

fn tables() -> Result<HashSet<CString>> {
    let socket = socket()?;
    socket.send(&nftnl::table::get_tables_nlmsg(1))?;
    let mut tables = HashSet::new();
    let mut buffer = vec![0u8; nftnl::nft_nlmsg_maxsize() as usize];
    loop {
        for message in socket.recv(&mut buffer)? {
            if matches!(
                mnl::cb_run2(
                    message?,
                    1,
                    socket.portid(),
                    nftnl::table::get_tables_cb,
                    &mut tables
                )?,
                mnl::CbResult::Stop
            ) {
                return Ok(tables);
            }
        }
    }
}

fn send(batch: nftnl::FinalizedBatch) -> Result<()> {
    let socket = socket()?;
    socket.send_all(&batch)?;
    let mut buffer = vec![0u8; nftnl::nft_nlmsg_maxsize() as usize];
    let mut sequences = batch.sequence_numbers();
    while !sequences.is_empty() {
        for message in socket.recv(&mut buffer)? {
            let sequence = sequences
                .next()
                .context("Unexpected firewall acknowledgement")?;
            mnl::cb_run(message?, sequence, socket.portid())?;
        }
    }
    Ok(())
}

pub fn preflight() -> Result<()> {
    let tables = tables()?;
    ensure!(
        tables.iter().all(|name| name.as_c_str() == TABLE),
        "Other nftables tables exist; this firewall arrangement is not supported. Roentgen did not alter them"
    );
    Ok(())
}

/// Replacing the table is one kernel transaction. Connection marks survive table replacement.
pub fn apply(enabled: bool, xray_uid: u32) -> Result<()> {
    apply_policy(enabled, xray_uid, false)
}
pub fn forward(xray_uid: u32) -> Result<()> {
    apply_policy(true, xray_uid, true)
}
fn apply_policy(enabled: bool, xray_uid: u32, allow_tun: bool) -> Result<()> {
    let existing = tables()?;
    let table = Table::new(TABLE, ProtoFamily::Inet);
    let mut batch = Batch::new();
    if existing.contains(TABLE) {
        batch.add(&table, MsgType::Del);
    }
    batch.add(&table, MsgType::Add);
    let mut incoming = Chain::new(c"incoming", &table);
    incoming.set_type(ChainType::Filter);
    incoming.set_hook(Hook::In, -150);
    incoming.set_policy(Policy::Accept);
    batch.add(&incoming, MsgType::Add);
    let mut classify = Rule::new(&incoming);
    classify.add_expr(&nft_expr!(meta iifname));
    classify.add_expr(&nft_expr!(cmp != &b"lo\0"[..]));
    classify.add_expr(&nft_expr!(ct state));
    classify.add_expr(&nft_expr!(bitwise mask 8u32, xor 0u32));
    classify.add_expr(&nft_expr!(cmp == 8u32));
    classify.add_expr(&nft_expr!(ct mark));
    classify.add_expr(&nft_expr!(bitwise mask !INCOMING_MARK, xor INCOMING_MARK));
    classify.add_expr(&nft_expr!(ct mark set));
    batch.add(&classify, MsgType::Add);
    if enabled {
        let mut routing = Chain::new(c"route", &table);
        routing.set_type(ChainType::Route);
        routing.set_hook(Hook::Out, -150);
        routing.set_policy(Policy::Accept);
        batch.add(&routing, MsgType::Add);
        exceptions(&mut batch, &routing, xray_uid);
        let mut mark = Rule::new(&routing);
        mark.add_expr(&nft_expr!(meta mark));
        mark.add_expr(&nft_expr!(bitwise mask !CAPTURE_MARK, xor CAPTURE_MARK));
        mark.add_expr(&nft_expr!(meta mark set));
        batch.add(&mark, MsgType::Add);
        let mut guard = Chain::new(c"guard", &table);
        guard.set_type(ChainType::Filter);
        guard.set_hook(Hook::Out, 100);
        guard.set_policy(Policy::Drop);
        batch.add(&guard, MsgType::Add);
        exceptions(&mut batch, &guard, xray_uid);
        if allow_tun {
            let mut tun = Rule::new(&guard);
            tun.add_expr(&nft_expr!(meta oifname));
            tun.add_expr(&nft_expr!(cmp == &b"roentgen0\0"[..]));
            tun.add_expr(&nft_expr!(verdict accept));
            batch.add(&tun, MsgType::Add);
        }
    }
    send(batch.finalize()).context("Could not apply Roentgen firewall transaction")
}

fn exceptions(batch: &mut Batch, chain: &Chain<'_>, xray_uid: u32) {
    let mut loopback = Rule::new(chain);
    loopback.add_expr(&nft_expr!(meta oifname));
    loopback.add_expr(&nft_expr!(cmp == &b"lo\0"[..]));
    loopback.add_expr(&nft_expr!(verdict accept));
    batch.add(&loopback, MsgType::Add);
    let mut core = Rule::new(chain);
    core.add_expr(&nft_expr!(meta skuid));
    core.add_expr(&nft_expr!(cmp == xray_uid));
    core.add_expr(&nft_expr!(verdict accept));
    batch.add(&core, MsgType::Add);
    let mut replies = Rule::new(chain);
    replies.add_expr(&nft_expr!(ct mark));
    replies.add_expr(&nft_expr!(bitwise mask INCOMING_MARK, xor 0u32));
    replies.add_expr(&nft_expr!(cmp == INCOMING_MARK));
    replies.add_expr(&nft_expr!(verdict accept));
    batch.add(&replies, MsgType::Add);
    for (network, mask) in [
        (Ipv4Addr::new(10, 0, 0, 0), Ipv4Addr::new(255, 0, 0, 0)),
        (Ipv4Addr::new(172, 16, 0, 0), Ipv4Addr::new(255, 240, 0, 0)),
        (Ipv4Addr::new(192, 168, 0, 0), Ipv4Addr::new(255, 255, 0, 0)),
    ] {
        let mut rule = Rule::new(chain);
        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!(bitwise mask mask, xor 0u32));
        rule.add_expr(&nft_expr!(cmp == network));
        rule.add_expr(&nft_expr!(verdict accept));
        batch.add(&rule, MsgType::Add);
    }
    for (network, mask) in [(
        Ipv6Addr::new(0xfc00, 0, 0, 0, 0, 0, 0, 0),
        Ipv6Addr::new(0xfe00, 0, 0, 0, 0, 0, 0, 0),
    )] {
        let mut rule = Rule::new(chain);
        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!(bitwise mask mask, xor Ipv6Addr::UNSPECIFIED));
        rule.add_expr(&nft_expr!(cmp == network));
        rule.add_expr(&nft_expr!(verdict accept));
        batch.add(&rule, MsgType::Add);
    }
    // DHCPv4 broadcast is link control, not a general broadcast exception.
    let mut dhcp = Rule::new(chain);
    dhcp.add_expr(&nft_expr!(meta nfproto));
    dhcp.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV4 as u8));
    dhcp.add_expr(&nft_expr!(payload ipv4 daddr));
    dhcp.add_expr(&nft_expr!(cmp == Ipv4Addr::BROADCAST));
    dhcp.add_expr(&nft_expr!(meta l4proto));
    dhcp.add_expr(&nft_expr!(cmp == libc::IPPROTO_UDP as u8));
    dhcp.add_expr(&nft_expr!(payload udp sport));
    dhcp.add_expr(&nft_expr!(cmp == 68u16.to_be()));
    dhcp.add_expr(&nft_expr!(payload udp dport));
    dhcp.add_expr(&nft_expr!(cmp == 67u16.to_be()));
    dhcp.add_expr(&nft_expr!(verdict accept));
    batch.add(&dhcp, MsgType::Add);
    let mut dhcp6 = Rule::new(chain);
    dhcp6.add_expr(&nft_expr!(meta nfproto));
    dhcp6.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
    dhcp6.add_expr(&nft_expr!(payload ipv6 daddr));
    dhcp6.add_expr(&nft_expr!(
        cmp == Ipv6Addr::new(0xff02, 0, 0, 0, 0, 0, 1, 2)
    ));
    dhcp6.add_expr(&nft_expr!(meta l4proto));
    dhcp6.add_expr(&nft_expr!(cmp == libc::IPPROTO_UDP as u8));
    dhcp6.add_expr(&nft_expr!(payload udp sport));
    dhcp6.add_expr(&nft_expr!(cmp == 546u16.to_be()));
    dhcp6.add_expr(&nft_expr!(payload udp dport));
    dhcp6.add_expr(&nft_expr!(cmp == 547u16.to_be()));
    dhcp6.add_expr(&nft_expr!(verdict accept));
    batch.add(&dhcp6, MsgType::Add);
    // Limit multicast ICMPv6 to router/neighbor discovery with hop limit 255.
    for kind in [133u8, 134, 135, 136] {
        let mut rule = Rule::new(chain);
        rule.add_expr(&nft_expr!(meta nfproto));
        rule.add_expr(&nft_expr!(cmp == libc::NFPROTO_IPV6 as u8));
        rule.add_expr(&nft_expr!(meta l4proto));
        rule.add_expr(&nft_expr!(cmp == libc::IPPROTO_ICMPV6 as u8));
        rule.add_expr(&nft_expr!(payload ipv6 hoplimit));
        rule.add_expr(&nft_expr!(cmp == 255u8));
        rule.add_expr(&nftnl::expr::Payload::Transport(
            nftnl::expr::TransportHeaderField::Icmpv6(nftnl::expr::Icmpv6HeaderField::Type),
        ));
        rule.add_expr(&nft_expr!(cmp == kind));
        rule.add_expr(&nft_expr!(verdict accept));
        batch.add(&rule, MsgType::Add);
    }
}

pub fn remove() -> Result<()> {
    if tables()?.contains(TABLE) {
        let mut batch = Batch::new();
        batch.add(&Table::new(TABLE, ProtoFamily::Inet), MsgType::Del);
        send(batch.finalize())?;
    }
    Ok(())
}