use clap::Parser;
use curl::easy::{Easy2, Handler, WriteError};
use rfc_8805::{RawRecord, Record, SerRecord};
use serde::Serialize;
use thiserror::Error;

use std::collections::HashSet;

#[derive(Debug, Error)]
enum Error {
    #[error("Unable to fetch geofeed url")]
    UnableToFetchFeed(#[from] curl::Error),
    #[error("Failed to parse geofeed")]
    UnableToParseFeed(#[from] csv::Error),
    #[error("Failed to unpack geofeed")]
    Deserialize(#[from] rfc_8805::Error),
}

struct Collector(Vec<u8>);

impl Handler for Collector {
    fn write(&mut self, data: &[u8]) -> Result<usize, WriteError> {
        self.0.extend_from_slice(data);
        Ok(data.len())
    }
}

fn fetch_geofeed(url: String) -> Result<String, Error> {
    let mut easy = Easy2::new(Collector(Vec::new()));
    easy.get(true)?;
    easy.url(&url)?;
    easy.perform()?;

    if easy.response_code()? != 200 {
        return Err(Error::UnableToFetchFeed(curl::Error::new(
            easy.response_code()?,
        )));
    }

    let contents = easy.get_ref();

    //remove any headers
    let raw_feed = String::from_utf8_lossy(&contents.0);
    let feed = raw_feed
        .trim()
        .lines()
        .filter(|line| !line.starts_with('#'))
        .collect::<Vec<&str>>();

    Ok(feed.join("\n").to_string())
}

fn parse_geofeed(input: String) -> Result<Vec<Record>, Error> {
    let mut geofeed_records = vec![];

    let mut csv = csv::ReaderBuilder::new()
        .has_headers(false)
        .from_reader(input.as_bytes());

    for result in csv.deserialize() {
        // We must tell Serde what type we want to deserialize into.
        let rec: RawRecord = result?;
        let record: Record = Record::try_from(rec)?;
        geofeed_records.push(record);
    }

    Ok(geofeed_records)
}

#[derive(Parser, Debug)]
#[command(version, about, long_about = None)]
struct Cli {
    /// Url to fetch geofeed from (ex. https://geoip.starlinkisp.net/feed.csv)
    #[arg(short, long)]
    geofeed_url: String,
    /// List of IPs to check is contained in the geofeed
    #[arg(short, long, num_args(1..))]
    ip: Vec<String>,
}

#[derive(Debug, Serialize)]
struct Output {
    ip: String,
    record: Option<SerRecord>,
}

fn main() -> Result<(), Error> {
    let cli = Cli::parse();
    let feed = fetch_geofeed(cli.geofeed_url)?;
    let feed_parsed = parse_geofeed(feed)?;
    let mut found: HashSet<String> = HashSet::new();

    for record in feed_parsed.iter() {
        for ip in cli.ip.iter() {
            if record.check_ip(ip)? {
                let output = Output {
                    ip: ip.to_string(),
                    record: Some(record.into()),
                };
                println!("{}", serde_json::to_string(&output).unwrap());
                found.insert(ip.to_string());
            }
        }
    }

    for ip in cli.ip.iter() {
        if !found.contains(ip) {
            let output = Output {
                ip: ip.to_string(),
                record: None,
            };
            println!("{}", serde_json::to_string(&output).unwrap());
        }
    }

    Ok(())
}
