pub mod firewall;
pub mod routes;

use crate::{model::Trojan, runtime::Runtime, xray};
use anyhow::Result;
use async_trait::async_trait;
use std::{path::Path, time::Duration};

pub struct LinuxRuntime {
    routing: String,
    applied: Option<Trojan>,
    armed: bool,
}
impl Default for LinuxRuntime {
    fn default() -> Self {
        Self {
            routing: "unknown".into(),
            applied: None,
            armed: false,
        }
    }
}

#[async_trait]
impl Runtime for LinuxRuntime {
    async fn arm(&mut self) -> Result<()> {
        xray::preflight().await?;
        firewall::preflight()?;
        routes::preflight().await?;
        anyhow::ensure!(
            self.applied.is_some() || routes::tun_index().await?.is_none(),
            "roentgen0 already exists outside the active owned configuration; resolve the interface conflict before enabling"
        );
        let (uid, _) = xray::identity()?;
        firewall::apply(true, uid)?;
        // Persist the boot guard input before recording enabled intent in SQLite.
        write_guard(true)?;
        self.routing = "blocked".into();
        self.armed = true;
        Ok(())
    }
    async fn block(&mut self) -> Result<()> {
        firewall::apply(true, xray::identity()?.0)?;
        write_guard(true)?;
        self.routing = "blocked".into();
        self.applied = None;
        xray::systemctl(&["stop", "roentgen-xray.service"]).await?;
        routes::remove().await?;
        Ok(())
    }
    async fn apply(&mut self, connection: Option<&Trojan>, _: u64) -> Result<()> {
        if connection.is_none() {
            return self.block().await;
        }
        anyhow::ensure!(
            self.armed,
            "Activation requires successful compatibility checks and an armed guard"
        );
        if connection.is_some() && self.applied.as_ref() == connection {
            if let Some(index) = routes::tun_index().await? {
                routes::capture(index).await?;
                firewall::forward(xray::identity()?.0)?;
                self.routing = "tun".into();
                return Ok(());
            }
        }
        self.routing = "blocked".into();
        self.applied = None;
        xray::systemctl(&["stop", "roentgen-xray.service"]).await?;
        routes::remove().await?;
        if let Some(connection) = connection {
            anyhow::ensure!(
                routes::tun_index().await?.is_none(),
                "roentgen0 remained after stopping the owned Xray unit; refusing to take over another interface"
            );
            xray::preflight().await?;
            firewall::preflight()?;
            routes::preflight().await?;
            xray::write_config(
                Path::new("/run/roentgen/xray.json"),
                &xray::configuration(connection, true),
            )?;
            xray::systemctl(&["start", "roentgen-xray.service"]).await?;
            for _ in 0..100 {
                if let Some(index) = routes::tun_index().await? {
                    routes::capture(index).await?;
                    firewall::forward(xray::identity()?.0)?;
                    self.routing = "tun".into();
                    self.applied = Some(connection.clone());
                    return Ok(());
                }
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
            anyhow::bail!("Xray did not create roentgen0; traffic remains blocked");
        }
        Ok(())
    }
    async fn maintain(&mut self, connection: Option<&Trojan>, generation: u64) -> Result<()> {
        if connection.is_none() {
            return self.block().await;
        }
        if !self.armed {
            self.arm().await?;
        }
        self.apply(connection, generation).await
    }
    async fn probe(&mut self, connection: &Trojan, duration: u64) -> Result<xray::Probe> {
        xray::Probe::start(connection, duration).await
    }
    async fn disable(&mut self) -> Result<()> {
        xray::systemctl(&["stop", "roentgen-xray.service"]).await?;
        routes::remove().await?;
        self.applied = None;
        write_guard(false)?;
        firewall::apply(false, xray::identity()?.0)?;
        self.routing = "direct".into();
        self.armed = false;
        Ok(())
    }
    fn routing(&self) -> &str {
        &self.routing
    }
}

pub fn write_guard(enabled: bool) -> Result<()> {
    crate::storage::atomic_private_write(
        Path::new("/etc/roentgen/guard.toml"),
        format!("version = 1\nenabled = {enabled}\n").as_bytes(),
    )
}

pub fn restore_guard() -> Result<()> {
    #[derive(serde::Deserialize)]
    #[serde(deny_unknown_fields)]
    struct Guard {
        version: u32,
        enabled: bool,
    }
    let uid = xray::identity()?.0;
    let input = std::fs::read_to_string("/etc/roentgen/guard.toml");
    let value = input
        .ok()
        .and_then(|input| toml::from_str::<Guard>(&input).ok())
        .filter(|value| value.version == 1);
    match value {
        Some(value) => firewall::apply(value.enabled, uid),
        None => {
            firewall::apply(true, uid)?;
            anyhow::bail!(
                "Corrupt or missing guard state; protection was armed. Use explicit recover --disable"
            );
        }
    }
}