use anyhow::{Context, Result};
use comfy_table::{presets::NOTHING, Table};
use ipnet::Ipv4Net;
use serde::Serialize;
use std::{
    collections::BTreeSet,
    fmt::Debug,
    io::{Read, Write},
    net::{Ipv4Addr, TcpStream},
    ops::Deref,
    thread::sleep,
    time::Duration,
};
use tracing::{debug, instrument, trace, warn};

#[derive(Clone, Debug, Serialize)]
pub struct Output<T> {
    pub input: Option<T>,
    pub valid: bool,
    pub canonical: Ipv4Net,
    pub contained_by: Option<Ipv4Net>,
    pub network: Ipv4Addr,
    pub address: Ipv4Addr,
    pub netmask: Ipv4Addr,
    pub hosts: usize,
    pub whois: Vec<String>,
}

pub trait Tabular {
    fn table(&self, headers: bool) -> Table;
}

impl<T: Debug> Tabular for Vec<Output<T>> {
    fn table(&self, headers: bool) -> Table {
        let mut table = Table::new();
        table.load_preset(NOTHING);
        if headers {
            table.add_row(vec![
                "POS",
                "INPUT",
                "INVALID",
                "CANONICAL",
                "CONTAINED",
                "NETWORK",
                "ADDRESS",
                "NETMASK",
                "HOSTS",
                "WHOIS",
            ]);
        }

        let mut host_count = 0;
        for (n, output) in self.iter().enumerate() {
            host_count += output.hosts;

            table.add_row(vec![
                (n + 1).to_string(),
                if let Some(input) = &output.input {
                    format!("{input:?}")
                } else {
                    String::new()
                },
                if output.valid {
                    ""
                } else {
                    "!"
                }
                .to_string(),
                output.canonical.to_string(),
                output
                    .contained_by
                    .map_or_else(String::new, |s| ToString::to_string(&s)),
                output.network.to_string(),
                output.address.to_string(),
                output.netmask.to_string(),
                output.hosts.to_string(),
                output.whois.join(" "),
            ]);
        }

        table
    }
}

pub trait Json {
    /// # Errors
    ///
    /// Returns an error if anything goes wrong during serialization.
    fn to_json(&self) -> Result<String>;
}

impl<T: Serialize> Json for Vec<Output<T>> {
    fn to_json(&self) -> Result<String> {
        Ok(serde_json::to_string_pretty(self)?)
    }
}

#[derive(Clone, Debug, Default)]
pub struct IpSet<T> {
    ips: BTreeSet<(Option<T>, Ipv4Net)>,
}

impl<T: AsRef<str> + Ord> IpSet<T> {
    /// Attempt to add a CIDR block from a string. This will first try to parse
    /// it as an `Ipv4Net`, and then fall back to trying it as an `Ipv4Addr`,
    /// which can be converted to a /32.
    ///
    /// # Errors
    ///
    /// Returns an error if the IP fails to parse.
    pub fn insert(&mut self, s: T) -> Result<&mut Self> {
        let ip = if let Ok(ip) = s.as_ref().parse() {
            ip
        } else {
            debug!("parsing as a CIDR block failed; trying as an IP");
            let ip: Ipv4Addr =
                s.as_ref().parse().context("parsing as Ipv4Addr")?;

            ip.into()
        };

        self.ips.insert((Some(s), ip));

        Ok(self)
    }
}

impl From<Vec<Ipv4Net>> for IpSet<String> {
    fn from(value: Vec<Ipv4Net>) -> Self {
        Self {
            ips: value
                .into_iter()
                .map(|i| (Some(i.to_string()), i))
                .collect(),
        }
    }
}

impl<T> Deref for IpSet<T> {
    type Target = BTreeSet<(Option<T>, Ipv4Net)>;

    fn deref(&self) -> &Self::Target {
        &self.ips
    }
}

impl<T: std::fmt::Debug + PartialEq<String> + Clone + ToString> IpSet<T> {
    fn ips(&self) -> Vec<Ipv4Net> {
        self.iter().map(|ip| ip.1).collect()
    }

    /// # Errors
    ///
    /// Returns an error if anything goes wrong talking to whois.
    #[instrument(ret, err)]
    pub fn check_ips(&self, whois: bool) -> Result<Vec<Output<T>>> {
        // Aggregate the IPs so we can look for supersets of ranges.
        let aggregated_ips = Ipv4Net::aggregate(&self.ips());

        // If the lengths between passed IPs and aggregated IPs don't match, it
        // means we've got redundant subnets. Warn the user.
        if aggregated_ips.len() != self.len() {
            warn!(
                aggregated = aggregated_ips.len(),
                ips = self.len(),
                "one or more subnets are not aggregated"
            );
        }

        let mut outputs = vec![];
        for (raw, ip) in self.iter() {
            let canonical = ip.trunc();

            // Figure out whether this is a subnet that's part of a greater range in
            // the set.
            let contained_by =
                aggregated_ips.iter().find(|i| *i != ip && i.contains(ip));

            let output = Output {
                input: raw.clone(),
                valid: raw
                    .as_ref()
                    .map_or(true, |s| *s == canonical.to_string()),
                canonical,
                contained_by: contained_by.copied(),
                network: ip.network(),
                address: ip.addr(),
                netmask: ip.netmask(),
                hosts: ip.hosts().count(),
                whois: if whois {
                    lookup(*ip)?
                } else {
                    vec![]
                },
            };
            outputs.push(output);

            sleep(Duration::from_millis(100));
        }

        Ok(outputs)
    }
}

fn lookup(ip: Ipv4Net) -> Result<Vec<String>> {
    // Look up the whois info.
    let mut stream = TcpStream::connect("whois.arin.net:43")?;
    stream.set_read_timeout(Some(Duration::from_secs(5)))?;
    stream.set_write_timeout(Some(Duration::from_secs(5)))?;
    let ip_addr = ip.addr().to_string();
    stream.write_all(&[b"n + ", ip_addr.as_bytes(), b"\n"].concat())?;
    trace!(?ip_addr, "wrote request");
    let mut buf = String::new();
    let bytes = stream.read_to_string(&mut buf)?;
    trace!(?bytes, ?buf, "read response");

    // Parse the response into key/val pairs.
    Ok(buf
        .lines()
        .filter_map(|l| {
            if l.starts_with('#') {
                return None;
            }

            let (k, v) = l.split_once(':')?;

            match k.trim() {
                "NetRange" | "CIDR" | "Organization" | "City" | "StateProv" => {
                    Some(v.trim().to_string())
                },
                _ => None,
            }
        })
        .collect())
}

#[cfg(test)]
mod tests {
    use super::*;
}