use async_process::Command;
use std::{fmt, net::Ipv4Addr};

#[derive(Clone, Copy, Debug, serde::Deserialize, serde::Serialize)]
pub(crate) enum Proto {
    Tcp,
    Udp,
}

impl Proto {
    fn as_iptables_str(&self) -> &'static str {
        match self {
            Proto::Tcp => "tcp",
            Proto::Udp => "udp",
        }
    }
}

impl fmt::Display for Proto {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Proto::Tcp => write!(f, "TCP"),
            Proto::Udp => write!(f, "UDP"),
        }
    }
}

pub(crate) async fn input_accept(
    proto: Proto,
    external_interface: &str,
    external_ip: Ipv4Addr,
    external_port: u16,
    external_mask: u8,
) -> Result<(), anyhow::Error> {
    input(
        proto,
        external_interface,
        external_ip,
        external_port,
        external_mask,
        move |cmd| cmd.arg("-I"),
    )
    .await
}

pub(crate) async fn delete_input_accept(
    proto: Proto,
    external_interface: &str,
    external_ip: Ipv4Addr,
    external_port: u16,
    external_mask: u8,
) -> Result<(), anyhow::Error> {
    input(
        proto,
        external_interface,
        external_ip,
        external_port,
        external_mask,
        move |cmd| cmd.arg("-D"),
    )
    .await
}

async fn input(
    proto: Proto,
    external_interface: &str,
    external_ip: Ipv4Addr,
    external_port: u16,
    external_mask: u8,
    func: impl Fn(&mut Command) -> &mut Command,
) -> Result<(), anyhow::Error> {
    iptables_filter(move |cmd| {
        func(cmd).args(&[
            "INPUT",
            "-d",
            &format!("{}/{}", external_ip, external_mask),
            "-i",
            external_interface,
            "-p",
            proto.as_iptables_str(),
            "-m",
            "conntrack",
            "--ctstate",
            "NEW,RELATED,ESTABLISHED",
            "-m",
            proto.as_iptables_str(),
            "--dport",
            &external_port.to_string(),
            "-j",
            "ACCEPT",
        ])
    })
    .await
}

pub(crate) async fn forward_accept(
    external_interface: &str,
    internal_interface: &str,
    proto: Proto,
    external_port: u16,
) -> Result<(), anyhow::Error> {
    forward(
        external_interface,
        internal_interface,
        proto,
        external_port,
        move |cmd| cmd.arg("-I"),
    )
    .await
}

pub(crate) async fn delete_forward_accept(
    external_interface: &str,
    internal_interface: &str,
    proto: Proto,
    external_port: u16,
) -> Result<(), anyhow::Error> {
    forward(
        external_interface,
        internal_interface,
        proto,
        external_port,
        move |cmd| cmd.arg("-D"),
    )
    .await
}

async fn forward(
    external_interface: &str,
    internal_interface: &str,
    proto: Proto,
    external_port: u16,
    func: impl Fn(&mut Command) -> &mut Command,
) -> Result<(), anyhow::Error> {
    iptables_filter(move |cmd| {
        func(cmd).args(&[
            "FORWARD",
            "-i",
            external_interface,
            "-o",
            internal_interface,
            "-p",
            proto.as_iptables_str(),
            "--dport",
            &external_port.to_string(),
            "-m",
            "conntrack",
            "--ctstate",
            "NEW,ESTABLISHED,RELATED",
            "-j",
            "ACCEPT",
        ])
    })
    .await
}

pub(crate) async fn forward_postrouting(
    proto: Proto,
    internal_ip: Ipv4Addr,
    internal_mask: u8,
    external_ip: Ipv4Addr,
    external_port: u16,
    destination_ip: Ipv4Addr,
) -> Result<(), anyhow::Error> {
    forward_postrouting_snat(
        proto,
        internal_ip,
        internal_mask,
        external_ip,
        external_port,
        destination_ip,
        |cmd| cmd.arg("-I"),
    )
    .await
}

pub(crate) async fn delete_forward_postrouting(
    proto: Proto,
    internal_ip: Ipv4Addr,
    internal_mask: u8,
    external_ip: Ipv4Addr,
    external_port: u16,
    destination_ip: Ipv4Addr,
) -> Result<(), anyhow::Error> {
    forward_postrouting_snat(
        proto,
        internal_ip,
        internal_mask,
        external_ip,
        external_port,
        destination_ip,
        |cmd| cmd.arg("-D"),
    )
    .await
}

async fn forward_postrouting_snat(
    proto: Proto,
    internal_ip: Ipv4Addr,
    internal_mask: u8,
    external_ip: Ipv4Addr,
    external_port: u16,
    destination_ip: Ipv4Addr,
    func: impl Fn(&mut Command) -> &mut Command,
) -> Result<(), anyhow::Error> {
    iptables_nat(move |cmd| {
        func(cmd).args(&[
            "POSTROUTING",
            "-s",
            &format!("{}/{}", internal_ip, internal_mask),
            "-d",
            &destination_ip.to_string(),
            "-p",
            proto.as_iptables_str(),
            "-m",
            proto.as_iptables_str(),
            "--dport",
            &external_port.to_string(),
            "-m",
            "conntrack",
            "--ctstate",
            "NEW,RELATED,ESTABLISHED",
            "-j",
            "SNAT",
            "--to-source",
            &external_ip.to_string(),
        ])
    })
    .await
}

pub(crate) async fn forward_prerouting(
    proto: Proto,
    external_ip: Ipv4Addr,
    external_mask: u8,
    external_port: u16,
    destination_ip: Ipv4Addr,
    destination_port: u16,
) -> Result<(), anyhow::Error> {
    forward_prerouting_dnat(
        proto,
        external_ip,
        external_mask,
        external_port,
        destination_ip,
        destination_port,
        |cmd| cmd.arg("-I"),
    )
    .await
}

pub(crate) async fn delete_forward_prerouting(
    proto: Proto,
    external_ip: Ipv4Addr,
    external_mask: u8,
    external_port: u16,
    destination_ip: Ipv4Addr,
    destination_port: u16,
) -> Result<(), anyhow::Error> {
    forward_prerouting_dnat(
        proto,
        external_ip,
        external_mask,
        external_port,
        destination_ip,
        destination_port,
        |cmd| cmd.arg("-D"),
    )
    .await
}

async fn forward_prerouting_dnat(
    proto: Proto,
    external_ip: Ipv4Addr,
    external_mask: u8,
    external_port: u16,
    destination_ip: Ipv4Addr,
    destination_port: u16,
    func: impl Fn(&mut Command) -> &mut Command,
) -> Result<(), anyhow::Error> {
    iptables_nat(move |cmd| {
        func(cmd).args(&[
            "PREROUTING",
            "-p",
            proto.as_iptables_str(),
            "-d",
            &format!("{}/{}", external_ip, external_mask),
            "--dport",
            &external_port.to_string(),
            "-m",
            "conntrack",
            "--ctstate",
            "NEW,ESTABLISHED,RELATED",
            "-j",
            "DNAT",
            "--to",
            &format!("{}:{}", destination_ip, destination_port),
        ])
    })
    .await
}

/// iptables -I FORWARD -i %i -j ACCEPT
/// iptables -I FORWARD -o %i -j ACCEPT

async fn iptables_nat(func: impl Fn(&mut Command) -> &mut Command) -> Result<(), anyhow::Error> {
    iptables(move |cmd| func(cmd.args(&["-t", "nat"]))).await
}

async fn iptables_filter(func: impl Fn(&mut Command) -> &mut Command) -> Result<(), anyhow::Error> {
    iptables(move |cmd| func(cmd.args(&["-t", "filter"]))).await
}

async fn iptables(func: impl Fn(&mut Command) -> &mut Command) -> Result<(), anyhow::Error> {
    let mut command = Command::new("iptables");

    func(&mut command);

    let mut child = command.spawn()?;

    let status = child.status().await?;

    if !status.success() {
        return Err(anyhow::Error::msg(format!(
            "Command failed with status {}",
            status
        )));
    }

    Ok(())
}