use super::firewall::CAPTURE_MARK;
use anyhow::{Result, ensure};
use futures_util::TryStreamExt;
use rtnetlink::{
    Handle, IpVersion, RouteMessageBuilder,
    packet_route::{
        route::{RouteAttribute, RouteProtocol},
        rule::{RuleAction, RuleAttribute},
    },
};
use std::net::{Ipv4Addr, Ipv6Addr};

const TABLE_ID: u32 = 20260;
const PRIORITY: u32 = 10260;
const PROTOCOL: u8 = 186;

fn connection() -> Result<Handle> {
    let (connection, handle, _) = rtnetlink::new_connection()?;
    tokio::spawn(connection);
    Ok(handle)
}

pub async fn preflight() -> Result<()> {
    let handle = connection()?;
    for version in [IpVersion::V4, IpVersion::V6] {
        let mut rules = handle.rule().get(version).execute();
        while let Some(rule) = rules.try_next().await? {
            let priority = rule
                .attributes
                .iter()
                .find_map(|a| {
                    if let RuleAttribute::Priority(v) = a {
                        Some(*v)
                    } else {
                        None
                    }
                })
                .unwrap_or(0);
            ensure!(
                matches!(priority, 0 | 32766 | 32767) || owned_rule(&rule),
                "Unsupported policy-routing rule at priority {}; remove conflicting VPN or policy routing before enabling",
                priority
            );
        }
    }
    let mut routes = handle
        .route()
        .get(rtnetlink::RouteMessageBuilder::<std::net::IpAddr>::new().build())
        .execute();
    while let Some(route) = routes.try_next().await? {
        if table(&route) == TABLE_ID {
            ensure!(
                route.header.protocol == RouteProtocol::Other(PROTOCOL),
                "Roentgen route table is already in use"
            );
        }
    }
    Ok(())
}

fn table(route: &rtnetlink::packet_route::route::RouteMessage) -> u32 {
    route
        .attributes
        .iter()
        .find_map(|a| {
            if let RouteAttribute::Table(v) = a {
                Some(*v)
            } else {
                None
            }
        })
        .unwrap_or(u32::from(route.header.table))
}
fn owned_rule(rule: &rtnetlink::packet_route::rule::RuleMessage) -> bool {
    rule.attributes.contains(&RuleAttribute::Priority(PRIORITY))
        && rule
            .attributes
            .contains(&RuleAttribute::FwMark(CAPTURE_MARK))
        && rule
            .attributes
            .contains(&RuleAttribute::FwMask(CAPTURE_MARK))
        && rule.attributes.contains(&RuleAttribute::Table(TABLE_ID))
        && rule.header.action == RuleAction::ToTable
}

pub async fn capture(interface: u32) -> Result<()> {
    let handle = connection()?;
    for route in [
        RouteMessageBuilder::<Ipv4Addr>::new()
            .table_id(TABLE_ID)
            .output_interface(interface)
            .protocol(RouteProtocol::Other(PROTOCOL))
            .build(),
        RouteMessageBuilder::<Ipv6Addr>::new()
            .table_id(TABLE_ID)
            .output_interface(interface)
            .protocol(RouteProtocol::Other(PROTOCOL))
            .build(),
    ] {
        handle.route().add(route).replace().execute().await?;
    }
    for version in [IpVersion::V4, IpVersion::V6] {
        let existing: Vec<_> = handle
            .rule()
            .get(version.clone())
            .execute()
            .try_collect()
            .await?;
        if existing.iter().any(owned_rule) {
            continue;
        }
        match version {
            IpVersion::V4 => {
                let mut request = handle
                    .rule()
                    .add()
                    .v4()
                    .priority(PRIORITY)
                    .table_id(TABLE_ID)
                    .fw_mark(CAPTURE_MARK)
                    .action(RuleAction::ToTable);
                request
                    .message_mut()
                    .attributes
                    .push(RuleAttribute::FwMask(CAPTURE_MARK));
                request.execute().await?;
            }
            IpVersion::V6 => {
                let mut request = handle
                    .rule()
                    .add()
                    .v6()
                    .priority(PRIORITY)
                    .table_id(TABLE_ID)
                    .fw_mark(CAPTURE_MARK)
                    .action(RuleAction::ToTable);
                request
                    .message_mut()
                    .attributes
                    .push(RuleAttribute::FwMask(CAPTURE_MARK));
                request.execute().await?;
            }
        }
    }
    Ok(())
}

pub async fn remove() -> Result<()> {
    let handle = connection()?;
    for version in [IpVersion::V4, IpVersion::V6] {
        let rules: Vec<_> = handle
            .rule()
            .get(version.clone())
            .execute()
            .try_collect()
            .await?;
        for rule in rules.into_iter().filter(owned_rule) {
            handle.rule().del(rule).execute().await?;
        }
    }
    let routes: Vec<_> = handle
        .route()
        .get(RouteMessageBuilder::<std::net::IpAddr>::new().build())
        .execute()
        .try_collect()
        .await?;
    for route in routes
        .into_iter()
        .filter(|r| table(r) == TABLE_ID && r.header.protocol == RouteProtocol::Other(PROTOCOL))
    {
        handle.route().del(route).execute().await?;
    }
    Ok(())
}

pub async fn tun_index() -> Result<Option<u32>> {
    let handle = connection()?;
    Ok(handle
        .link()
        .get()
        .match_name("roentgen0")
        .execute()
        .try_next()
        .await?
        .map(|link| link.header.index))
}