use std::{
    error::Error,
    fs::File,
    collections::{HashMap, HashSet}, 
    net::IpAddr,
};

use pretty_table::prelude::*;
use ip_network::IpNetwork;
use trust_dns_resolver::{
    config::{ResolverConfig, ResolverOpts},
    TokioAsyncResolver,
};
use maxminddb::{geoip2::City, Reader as GeoIpReader};
use libp2p::{
    identify::Info, multiaddr::Protocol, Multiaddr, PeerId
};

pub struct Locator<'a> {
    maybe_filename: &'a Option<String>,
    peers: &'a HashMap<&'a PeerId, &'a Info>,
    db: maxminddb::Reader<&'static [u8]>,
}

impl<'a> Locator<'a> {
    const CITY_DATA: &'static  [u8] = include_bytes!("../artifacts/GeoLite2-City.mmdb");

    pub fn new(
        peers: &'a HashMap<&'a PeerId, &'a Info>, 
        maybe_filename: &'a Option<String>,
    ) -> Self {
        Self { 
            peers,
            maybe_filename,
            db: GeoIpReader::from_source(Self::CITY_DATA)
                .expect("City data is always valid; qed"),
        }
    }

    fn is_public_adress(&self, addr: &Multiaddr) -> bool {
        let ip = match addr.iter().next() {
            Some(Protocol::Ip4(ip)) => IpNetwork::from(ip),
            Some(Protocol::Ip6(ip)) => IpNetwork::from(ip),
            Some(Protocol::Dns(_)) | Some(Protocol::Dns4(_)) | Some(Protocol::Dns6(_)) => return true,
            _ => return false,
        };
        ip.is_global()
    }

    pub async fn locate_peers(&self) {
        let mut geo_peers: Vec<Vec<String>> = Vec::new();
        let used_addrs: HashSet<Multiaddr> = HashSet::new();

        let resolver = TokioAsyncResolver::tokio(
            ResolverConfig::default(),
            ResolverOpts::default(),
        );

        for (_, info) in self.peers {
            for addr in &info.listen_addrs {
                if !self.is_public_adress(addr) { continue; }
                let located = match addr.iter().next() {
                    Some(Protocol::Ip4(ip)) => self.locate(IpAddr::V4(ip)),
                    Some(Protocol::Ip6(ip)) => self.locate(IpAddr::V6(ip)),
                    Some(Protocol::Dns(dns)) |
                        Some(Protocol::Dns4(dns)) |
                        Some(Protocol::Dns6(dns)) => {
                            let Ok(lookup) = resolver.lookup_ip(dns.to_string()).await else {
                                continue;
                            };
                            lookup.iter().find_map(|ip| self.locate(ip))
                        }
                    _ => continue,
                };

                if let None = used_addrs.get(addr) {
                    geo_peers.push(vec![
                        addr.to_string(), 
                        located.unwrap_or_default(),
                    ]);
                }
            }
        }

        match &self.maybe_filename {
            Some(filename) => {
                if filename.ends_with(".csv") {
                    if let Err(err) = self.write_to_csv(&filename, &geo_peers) {
                        tracing::error!("Could not write to file {} because of: {err}", &filename);
                    }
                } else {
                    if let Err(err) = write_table_to_file(filename.clone(), geo_peers) {
                        tracing::error!("Could not write to file {} because of: {err}", &filename);
                    }
                }
                println!("Results saved to {}", &filename);
            },
            None => print_table!(geo_peers),
        };
    }

    fn write_to_csv(
        &self, 
        filename: &str,
        data: &Vec<Vec<String>>,
    ) -> Result<(), Box<dyn Error>> {
        let file = File::create(filename)?;
        let mut writer = csv::Writer::from_writer(file);

        for record in data {
            writer.write_record(record)?;
        }

        writer.flush()?;
        Ok(())
    }

    fn locate(&self, ip: IpAddr) -> Option<String> {
        let City { city, .. } = self.db.lookup(ip).ok()?;

        let city = city
            .as_ref()?
            .names
            .as_ref()?
            .get("en")?
            .to_string()
            .into_boxed_str();

        Some(city.into_string())
    }
}
