use crate::constants;
use crate::libvirt_rest_structures::net::NetworkName;
use crate::nat::nat_definition::{Nat, NatSourceIP, NetNat};
use crate::utils::net_utils;
use clap::Parser;
use std::collections::HashMap;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::process::{Command, ExitStatus};

#[derive(thiserror::Error, Debug)]
enum NatConfModeError {
    #[error("UpdateFirewall failed!")]
    UpdateFirewall,
}

/// VirtWeb NAT configuration mode. This executable should never be executed manually
#[derive(Parser, Debug, Clone)]
#[clap(author, version, about, long_about = None)]
struct NatArgs {
    /// Storage directory
    #[clap(short, long)]
    storage: String,

    /// Network name
    #[clap(short, long)]
    network_name: String,

    /// Operation
    #[clap(short, long)]
    operation: String,

    /// Sub operation
    #[clap(long)]
    sub_operation: String,
}

impl NatArgs {
    pub fn network_file(&self) -> PathBuf {
        let network_name = NetworkName(self.network_name.to_string());
        Path::new(&self.storage)
            .join(constants::STORAGE_NAT_DIR)
            .join(network_name.nat_file_name())
    }
}

/// NAT sub main function
pub async fn sub_main() -> anyhow::Result<()> {
    let args = NatArgs::parse();

    if !args.network_file().exists() {
        log::warn!("Cannot do anything for the network, because the NAT configuration file does not exixsts!");
        return Ok(());
    }

    let conf_json = std::fs::read_to_string(args.network_file())?;
    let conf: NetNat = serde_json::from_str(&conf_json)?;

    let nic_ips = net_utils::net_list_and_ips()?;

    match (args.operation.as_str(), args.sub_operation.as_str()) {
        ("started", "begin") => {
            log::info!("Enable port forwarding for network");
            trigger_nat_forwarding(true, &conf, &nic_ips).await?
        }
        ("stopped", "end") => {
            log::info!("Disable port forwarding for network");
            trigger_nat_forwarding(false, &conf, &nic_ips).await?
        }
        _ => log::debug!(
            "Operation {} - {} not supported",
            args.operation,
            args.sub_operation
        ),
    }

    Ok(())
}

async fn trigger_nat_forwarding(
    enable: bool,
    conf: &NetNat,
    nic_ips: &HashMap<String, Vec<IpAddr>>,
) -> anyhow::Result<()> {
    if let Some(ipv4) = &conf.ipv4 {
        trigger_nat_forwarding_nat_ipv(
            enable,
            &conf.interface,
            &ipv4.iter().map(|i| i.generalize()).collect::<Vec<_>>(),
            nic_ips,
        )
        .await?;
    }

    if let Some(ipv6) = &conf.ipv6 {
        trigger_nat_forwarding_nat_ipv(
            enable,
            &conf.interface,
            &ipv6.iter().map(|i| i.generalize()).collect::<Vec<_>>(),
            nic_ips,
        )
        .await?;
    }

    Ok(())
}

async fn trigger_nat_forwarding_nat_ipv(
    enable: bool,
    net_interface: &str,
    rules: &[Nat<IpAddr>],
    nic_ips: &HashMap<String, Vec<IpAddr>>,
) -> anyhow::Result<()> {
    for r in rules {
        let host_ips = match &r.host_ip {
            NatSourceIP::Interface { name } => nic_ips.get(name).cloned().unwrap_or_default(),
            NatSourceIP::Ip { ip } => vec![*ip],
        };

        for host_ip in host_ips {
            let mut guest_port = r.guest_port;
            for host_port in r.host_port.as_seq() {
                if r.protocol.has_tcp() {
                    toggle_port_forwarding(
                        enable,
                        false,
                        host_ip,
                        host_port,
                        net_interface,
                        r.guest_ip,
                        guest_port,
                    )?
                }

                if r.protocol.has_udp() {
                    toggle_port_forwarding(
                        enable,
                        true,
                        host_ip,
                        host_port,
                        net_interface,
                        r.guest_ip,
                        guest_port,
                    )?
                }

                guest_port += 1;
            }
        }
    }

    Ok(())
}

fn check_cmd(s: ExitStatus) -> anyhow::Result<()> {
    if !s.success() {
        log::error!("Failed to update firewall rules!");
        return Err(NatConfModeError::UpdateFirewall.into());
    }
    Ok(())
}

fn toggle_port_forwarding(
    enable: bool,
    is_udp: bool,
    host_ip: IpAddr,
    host_port: u16,
    net_interface: &str,
    guest_ip: IpAddr,
    guest_port: u16,
) -> anyhow::Result<()> {
    if host_ip.is_ipv4() != guest_ip.is_ipv4() {
        log::trace!("Skipping invalid combination {host_ip} -> {guest_ip}");
        return Ok(());
    }

    let program = match host_ip.is_ipv4() {
        true => "/sbin/iptables",
        false => "/sbin/ip6tables",
    };

    let protocol = match is_udp {
        true => "udp",
        false => "tcp",
    };

    log::info!("Forward (add={enable}) incoming {protocol} connections for {host_ip}:{host_port} to {guest_ip}:{guest_port} int {net_interface}");

    // Rule 1
    let cmd = Command::new(program)
        .arg(match enable {
            true => "-I",
            false => "-D",
        })
        .arg("FORWARD")
        .arg("-o")
        .arg(net_interface)
        .arg("-p")
        .arg(protocol)
        .arg("-d")
        .arg(guest_ip.to_string())
        .arg("--dport")
        .arg(guest_port.to_string())
        .arg("-j")
        .arg("ACCEPT")
        .status()?;
    check_cmd(cmd)?;

    // Rule 2
    let cmd = Command::new(program)
        .arg("-t")
        .arg("nat")
        .arg(match enable {
            true => "-I",
            false => "-D",
        })
        .arg("PREROUTING")
        .arg("-p")
        .arg(protocol)
        .arg("-d")
        .arg(host_ip.to_string())
        .arg("--dport")
        .arg(host_port.to_string())
        .arg("-j")
        .arg("DNAT")
        .arg("--to")
        .arg(format!("{guest_ip}:{guest_port}"))
        .status()?;
    check_cmd(cmd)?;

    Ok(())
}