use std::env;
use std::io;
use std::net::IpAddr;
use std::process;
use std::str::FromStr;
use toluol::DnsType;

#[derive(Debug, PartialEq)]
pub enum ConnectionType {
    UDP,
    TCP,
    TLS
}

#[derive(Debug)]
pub struct Args {
    pub nameserver: String,
    pub hostname: Option<String>,
    pub url: String,
    pub qtype: DnsType,
    pub use_dnssec: bool,
    pub connection_type: ConnectionType
}

pub fn parse_args() -> io::Result<Args> {
    let mut args: Vec<String> = env::args().collect();
    args.remove(0); // remove executable name

    let mut nameserver = "1.1.1.1".to_string();
    let mut hostname = None;
    let mut url = "example.com".to_string();
    let mut qtype = DnsType::A;
    let mut use_dnssec = false;
    let mut connection_type = ConnectionType::UDP;
    let mut reverse = false;

    for arg in args {
        if let Some(ns) = arg.strip_prefix('@') { // nameserver
            nameserver = ns.to_string();
        } else if let Some(flag) = arg.strip_prefix('+') { // flags
            match flag {
                "do"  => { use_dnssec = true; },
                "tcp" => { connection_type = ConnectionType::TCP; },
                "tls" => { connection_type = ConnectionType::TLS; },
                x => { eprintln!("Invalid flag: {}.", x); process::exit(1); }
            }
        } else if let Some(option) = arg.strip_prefix('-') { // options
            if option.contains('=') {
                let split: Vec<&str> = option.split('=').collect();
                if split.len() > 2 {
                    eprintln!("Invalid option: {} (wrong length).", option); process::exit(1);
                }
                match split[0] {
                    "-hostname" => { hostname = Some(split[1].to_string()); },
                    x => { eprintln!("Invalid option: {}.", x); process::exit(1); }
                }
            } else {
                match option {
                    "h" | "-help" => { print_help(); process::exit(0); },
                    "x" => { reverse = true; },
                    x => { eprintln!("Invalid option: {}.", x); process::exit(1); }
                }
            }
        } else {
            match DnsType::from_str(&arg.to_uppercase()) {
                Ok(t) => { qtype = t; },
                Err(_) => { url = arg; }
            }
        }
    }

    if reverse {
        match IpAddr::from_str(url.as_str()) {
            Err(_) => {
                eprintln!("Invalid IP address: {}.", url);
                process::exit(1);
            },
            Ok(IpAddr::V4(addr)) => {
                let octets = addr.octets();
                url = format!("{}.{}.{}.{}.in-addr.arpa", octets[3], octets[2], octets[1], octets[0]);
                qtype = DnsType::PTR;
            },
            Ok(IpAddr::V6(addr)) => {
                url = String::with_capacity(72);
                for s in addr.segments().iter().rev() {
                    for c in format!("{:04x}", s).chars().rev() {
                        url.push(c);
                        url.push('.');
                    }
                }
                url.push_str("ip6.arpa");
                qtype = DnsType::PTR;
            }
        }
    }

    if connection_type == ConnectionType::TLS {
        if hostname.is_none() {
            eprintln!("Connection type set to TLS but hostname not set.");
            process::exit(1);
        } else if IpAddr::from_str(nameserver.as_str()).is_err() {
            eprintln!("Invalid nameserver: {}. Must be an IP address when using TLS.", nameserver);
            process::exit(1);
        }
    }

    Ok(Args {
        nameserver,
        hostname,
        url,
        qtype,
        use_dnssec,
        connection_type
    })
}

fn print_help() {
    println!("Usage: \ttoluol [@nameserver] [domain] [q-type] [options] [flags]");
    println!("Where: \tq-type is a valid DNS QTYPE (A, AAAA, TXT, MX, SOA, ...)");
    println!("       \toptions is one or more of the following:");
    println!("       \t    -h | --help            (prints this help message)");
    println!("       \t    -x                     (shortcut for reverse lookup)");
    println!("       \t    --hostname=<hostname>  (hostname of the nameserver)");
    println!("       \tflags is one or more of the following:");
    println!("       \t    +do   (indicates DNSSEC records shall be fetched)");
    println!("       \t    +tcp  (indicates TCP instead of UDP shall be used)");
    println!("       \t    +tls  (indicates DNS over TLS shall be used)");
    println!();
    println!("NOTE: order of the arguments does not matter.");
    println!("If no arguments are specified, the default behaviour is `toluol @1.1.1.1 A example.com`.")
}
