use std::collections::HashSet;

use anyhow::Result;
use iptables;
use log::*;
use tokio::{
    select,
    sync::watch,
    time::{self, Duration},
};

use crate::{fw, messages};

pub struct FirewallActor {
    pub ipt_v4: iptables::IPTables,
    pub ipt_v6: iptables::IPTables,
    rx_ports: watch::Receiver<messages::PublicExposedPorts>,
    last_ports: messages::PublicExposedPorts,
    refresh: Duration,
}

impl FirewallActor {
    pub async fn new(
        refresh: Duration,
        rxp: &watch::Receiver<messages::PublicExposedPorts>,
    ) -> Result<Self> {
        let ctx = Self {
            ipt_v4: iptables::new(false)?,
            ipt_v6: iptables::new(true)?,
            rx_ports: rxp.clone(),
            last_ports: messages::PublicExposedPorts::new(),
            refresh,
        };

        fw::setup(&ctx.ipt_v4)?;
        fw::setup(&ctx.ipt_v6)?;

        return Ok(ctx);
    }

    pub async fn listen(&mut self) -> Result<()> {
        let mut interval = time::interval(self.refresh);
        loop {
            // 1. Wait for an event
            let new_ports = select! {
              _ = self.rx_ports.changed() => Some(self.rx_ports.borrow().clone()),
              _ = interval.tick() => None,
              else => return Ok(()) // Sender dropped, terminate loop.
            };

            // 2. Update last ports if needed
            if let Some(p) = new_ports {
                self.last_ports = p;
            }

            // 3. Update firewall rules
            match self.do_fw_update().await {
                Ok(()) => debug!("Successfully updated firewall rules"),
                Err(e) => error!("An error occured while updating firewall rules. {}", e),
            }
        }
    }

    pub async fn do_fw_update(&self) -> Result<()> {
        for ipt in [&self.ipt_v4, &self.ipt_v6] {
            let curr_opened_ports = fw::get_opened_ports(ipt)?;

            let diff_tcp = self
                .last_ports
                .tcp_ports
                .difference(&curr_opened_ports.tcp_ports)
                .copied()
                .collect::<HashSet<u16>>();
            let diff_udp = self
                .last_ports
                .udp_ports
                .difference(&curr_opened_ports.udp_ports)
                .copied()
                .collect::<HashSet<u16>>();

            let ports_to_open = messages::PublicExposedPorts {
                tcp_ports: diff_tcp,
                udp_ports: diff_udp,
            };

            fw::open_ports(ipt, ports_to_open)?;
        }

        return Ok(());
    }
}