use anyhow::{Context, Result};
use futures::stream::TryStreamExt;
use netlink_packet_route::route::RouteProtocol;
use rtnetlink::{new_connection, Handle, IpVersion};
use std::net::{IpAddr, Ipv4Addr};
use tracing::{info, warn};

const VPN_ROUTE_METRIC: u32 = 50;

/// Manages routing table for VPN connection using netlink.
/// Tracks every route it installs and removes exactly those on teardown.
/// Pre-existing routes are never touched.
pub struct RouteManager {
    handle: Handle,
    added_routes: Vec<(Ipv4Addr, u8)>,
    tun_interface: String,
    tun_index: Option<u32>,
    dns_manager: crate::dns::DnsManager,
    // Saved during setup for re-adding the server host route after resume.
    server_ip: Option<Ipv4Addr>,
    wifi_gateway: Option<Ipv4Addr>,
    // Saved for re-adding the VPN default route after reconnect.
    vpn_gateway: Option<Ipv4Addr>,
}

impl RouteManager {
    pub async fn new(tun_interface: String) -> Result<Self> {
        let (connection, handle, _) =
            new_connection().context("Failed to create netlink connection")?;

        tokio::spawn(connection);

        Ok(Self {
            handle,
            added_routes: Vec::new(),
            tun_interface: tun_interface.clone(),
            tun_index: None,
            dns_manager: crate::dns::DnsManager::new(tun_interface, None, IpAddr::V4(Ipv4Addr::UNSPECIFIED)),
            server_ip: None,
            wifi_gateway: None,
            vpn_gateway: None,
        })
    }

    async fn get_interface_index(&self, name: &str) -> Result<u32> {
        let mut links = self
            .handle
            .link()
            .get()
            .match_name(name.to_string())
            .execute();

        if let Some(link) = links.try_next().await? {
            Ok(link.header.index)
        } else {
            anyhow::bail!("Interface {} not found", name)
        }
    }

    async fn get_default_gateway(&self) -> Result<Ipv4Addr> {
        let mut routes = self.handle.route().get(IpVersion::V4).execute();

        while let Some(route) = routes.try_next().await? {
            if route.header.destination_prefix_length == 0 {
                for nla in route.attributes.iter() {
                    if let netlink_packet_route::route::RouteAttribute::Gateway(addr) = nla {
                        use netlink_packet_route::route::RouteAddress;
                        if let RouteAddress::Inet(ipv4_addr) = addr {
                            return Ok(*ipv4_addr);
                        }
                    }
                }
            }
        }

        anyhow::bail!("No default gateway found")
    }

    /// Install routes for the VPN connection.
    ///
    /// Uses metric 50 (lower than typical DHCP default of 100) so VPN routes
    /// win without deleting the original default route.
    pub async fn setup_vpn_routes(
        &mut self,
        server_addr: &str,
        vpn_gateway: IpAddr,
        routes: &[String],
    ) -> Result<()> {
        info!("Setting up VPN routes...");

        self.tun_index = Some(
            self.get_interface_index(&self.tun_interface)
                .await
                .context("Failed to get TUN interface index")?,
        );
        info!(
            "TUN interface {} has index {}",
            self.tun_interface,
            self.tun_index.unwrap()
        );

        let current_gateway = self
            .get_default_gateway()
            .await
            .context("Failed to get current gateway")?;
        info!("Current gateway: {}", current_gateway);

        // Resolve server address to IP
        let server_ip = if let Ok(ip) = server_addr.parse::<IpAddr>() {
            ip
        } else {
            let host = server_addr.split(':').next().unwrap_or(server_addr);
            let addrs = tokio::net::lookup_host(format!("{}:0", host))
                .await
                .context("Failed to resolve server hostname")?;
            addrs
                .map(|addr| addr.ip())
                .next()
                .context("No IP address found for server")?
        };
        info!("Server IP: {}", server_ip);

        // Host route for VPN server through current gateway so the tunnel itself
        // doesn't loop back through itself.
        if let IpAddr::V4(server_v4) = server_ip {
            self.server_ip = Some(server_v4);
            self.wifi_gateway = Some(current_gateway);

            let result = self
                .handle
                .route()
                .add()
                .v4()
                .destination_prefix(server_v4, 32)
                .gateway(current_gateway)
                .priority(VPN_ROUTE_METRIC)
                .execute()
                .await;

            match result {
                Ok(_) => self.added_routes.push((server_v4, 32)),
                Err(e) => {
                    let err_msg = format!("{:?}", e);
                    if err_msg.contains("File exists")
                        || err_msg.contains("EEXIST")
                        || err_msg.contains("code: Some(-17)")
                    {
                        warn!("Server route already exists, continuing...");
                    } else {
                        return Err(anyhow::anyhow!("Failed to add server route: {:#}", e));
                    }
                }
            }
        } else {
            warn!("IPv6 server addresses not yet supported for routing");
        }

        let IpAddr::V4(gateway_v4) = vpn_gateway else {
            return Err(anyhow::anyhow!("IPv6 VPN gateways not supported"));
        };
        self.vpn_gateway = Some(gateway_v4);

        let is_full_tunnel = routes.iter().any(|r| r == "0.0.0.0/0");

        if is_full_tunnel {
            info!("Full-tunnel mode: adding default route via VPN...");
            self.handle
                .route()
                .add()
                .v4()
                .destination_prefix(Ipv4Addr::new(0, 0, 0, 0), 0)
                .gateway(gateway_v4)
                .output_interface(self.tun_index.unwrap())
                .priority(VPN_ROUTE_METRIC)
                .protocol(RouteProtocol::Boot)
                .execute()
                .await
                .map_err(|e| anyhow::anyhow!("Failed to add VPN default route: {:?}", e))?;

            self.added_routes.push((Ipv4Addr::new(0, 0, 0, 0), 0));
        } else {
            info!(
                "Split-tunnel mode: adding {} routes via VPN...",
                routes.len()
            );
            for route_str in routes {
                let Some((dest_str, prefix_str)) = route_str.split_once('/') else {
                    warn!("Invalid route format: {}", route_str);
                    continue;
                };
                let dest: Ipv4Addr = match dest_str.parse() {
                    Ok(ip) => ip,
                    Err(_) => {
                        warn!("Invalid route destination: {}", dest_str);
                        continue;
                    }
                };
                let prefix_len: u8 = match prefix_str.parse() {
                    Ok(p) => p,
                    Err(_) => {
                        warn!("Invalid prefix length: {}", prefix_str);
                        continue;
                    }
                };

                let result = self
                    .handle
                    .route()
                    .add()
                    .v4()
                    .destination_prefix(dest, prefix_len)
                    .gateway(gateway_v4)
                    .output_interface(self.tun_index.unwrap())
                    .priority(VPN_ROUTE_METRIC)
                    .execute()
                    .await;

                match result {
                    Ok(_) => {
                        self.added_routes.push((dest, prefix_len));
                        info!("Added route: {}/{}", dest, prefix_len);
                    }
                    Err(e) => {
                        let err_msg = format!("{:?}", e);
                        if err_msg.contains("File exists")
                            || err_msg.contains("EEXIST")
                            || err_msg.contains("code: Some(-17)")
                        {
                            warn!("Route {} already exists, continuing...", route_str);
                        } else {
                            warn!("Failed to add route {}: {:?}", route_str, e);
                        }
                    }
                }
            }
        }

        info!("✅ VPN routes configured successfully");
        Ok(())
    }

    /// Setup DNS configuration
    pub async fn setup_dns(&mut self, dns_servers: &[IpAddr], assigned_ip: IpAddr) -> Result<()> {
        if dns_servers.is_empty() {
            return Ok(());
        }

        self.dns_manager =
            crate::dns::DnsManager::new(self.tun_interface.clone(), self.tun_index, assigned_ip);

        self.dns_manager.set_dns(dns_servers).await
    }

    /// Restore DNS configuration
    pub async fn restore_dns(&self) -> Result<()> {
        self.dns_manager.restore().await
    }

    /// Re-add the server host route if it has been removed by the kernel.
    ///
    /// On Linux, suspending the machine brings wlan0 down, which causes the kernel
    /// to purge all routes installed on that interface — including our host route for
    /// the VPN server IP via the WiFi gateway.  Without that route, reconnect SYNs
    /// follow the VPN default route into tun0, get dropped, and time out.  This
    /// method must be called before each reconnect attempt.
    pub async fn ensure_server_route(&mut self) -> Result<()> {
        let (server_ip, gateway) = match (self.server_ip, self.wifi_gateway) {
            (Some(s), Some(g)) => (s, g),
            _ => return Ok(()),
        };

        let result = self
            .handle
            .route()
            .add()
            .v4()
            .destination_prefix(server_ip, 32)
            .gateway(gateway)
            .priority(VPN_ROUTE_METRIC)
            .execute()
            .await;

        match result {
            Ok(_) => {
                info!(
                    "Re-added server host route {}/32 via {} (was removed on suspend)",
                    server_ip, gateway
                );
            }
            Err(e) => {
                let msg = format!("{:?}", e);
                if msg.contains("File exists") || msg.contains("EEXIST") || msg.contains("code: Some(-17)") {
                    // Route is already present — nothing to do.
                } else {
                    return Err(anyhow::anyhow!("Failed to re-add server route: {:#}", e));
                }
            }
        }

        Ok(())
    }

    /// Remove the VPN default route (0.0.0.0/0) during a reconnect window so
    /// traffic falls through to the WiFi default route (metric 600).
    /// Call `resume_default_route()` after the new session is established.
    pub async fn suspend_default_route(&mut self) -> Result<()> {
        if !self.added_routes.contains(&(Ipv4Addr::new(0, 0, 0, 0), 0)) {
            return Ok(());
        }

        let mut routes = self.handle.route().get(IpVersion::V4).execute();
        let mut to_delete = None;
        while let Some(route) = routes.try_next().await? {
            if route.header.destination_prefix_length == 0 {
                let is_ours = route.attributes.iter().any(|nla| {
                    matches!(
                        nla,
                        netlink_packet_route::route::RouteAttribute::Priority(p)
                            if *p == VPN_ROUTE_METRIC
                    )
                });
                if is_ours {
                    to_delete = Some(route);
                    break;
                }
            }
        }

        if let Some(route) = to_delete {
            self.handle.route().del(route).execute().await.map_err(|e| {
                anyhow::anyhow!("Failed to suspend VPN default route: {:?}", e)
            })?;
            info!("VPN default route suspended — traffic falls to WiFi during reconnect");
        }
        Ok(())
    }

    /// Re-add the VPN default route after a successful reconnect.
    pub async fn resume_default_route(&mut self) -> Result<()> {
        let gateway = match self.vpn_gateway {
            Some(gw) => gw,
            None => return Ok(()),
        };
        if !self.added_routes.contains(&(Ipv4Addr::new(0, 0, 0, 0), 0)) {
            return Ok(());
        }

        let result = self
            .handle
            .route()
            .add()
            .v4()
            .destination_prefix(Ipv4Addr::new(0, 0, 0, 0), 0)
            .gateway(gateway)
            .output_interface(self.tun_index.unwrap_or(0))
            .priority(VPN_ROUTE_METRIC)
            .protocol(RouteProtocol::Boot)
            .execute()
            .await;

        match result {
            Ok(_) => info!("VPN default route restored after reconnect"),
            Err(e) => {
                let msg = format!("{:?}", e);
                if msg.contains("File exists") || msg.contains("EEXIST") || msg.contains("code: Some(-17)") {
                    // already present
                } else {
                    return Err(anyhow::anyhow!("Failed to resume VPN default route: {:#}", e));
                }
            }
        }
        Ok(())
    }

    /// Remove every route that was added by `setup_vpn_routes`.
    /// Pre-existing routes are untouched.
    pub async fn restore_routes(&mut self) -> Result<()> {
        info!(
            "Removing {} tracked VPN routes...",
            self.added_routes.len()
        );

        if self.added_routes.is_empty() {
            return Ok(());
        }

        let mut routes_to_delete = Vec::new();
        let mut routes = self.handle.route().get(IpVersion::V4).execute();

        while let Some(route) = routes.try_next().await? {
            let prefix_len = route.header.destination_prefix_length;

            let dest_ip: Option<Ipv4Addr> = if prefix_len == 0 {
                Some(Ipv4Addr::new(0, 0, 0, 0))
            } else {
                route.attributes.iter().find_map(|nla| {
                    if let netlink_packet_route::route::RouteAttribute::Destination(addr) = nla {
                        use netlink_packet_route::route::RouteAddress;
                        if let RouteAddress::Inet(ipv4) = addr {
                            return Some(*ipv4);
                        }
                    }
                    None
                })
            };

            if let Some(ip) = dest_ip {
                if self.added_routes.contains(&(ip, prefix_len)) {
                    let is_ours = route.attributes.iter().any(|nla| {
                        matches!(
                            nla,
                            netlink_packet_route::route::RouteAttribute::Priority(p)
                                if *p == VPN_ROUTE_METRIC
                        )
                    });
                    if is_ours {
                        routes_to_delete.push(route);
                    }
                }
            }
        }

        for route in routes_to_delete {
            if let Err(e) = self.handle.route().del(route).execute().await {
                warn!("Failed to delete tracked route: {:?}", e);
            }
        }

        self.added_routes.clear();
        info!("✅ VPN routes removed");
        Ok(())
    }
}

impl Drop for RouteManager {
    fn drop(&mut self) {
        if !self.added_routes.is_empty() {
            warn!(
                "RouteManager dropped with {} unrestored routes — call restore_routes() before dropping",
                self.added_routes.len()
            );
        }
    }
}