#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate rust_util;

mod config;
mod x509;
// mod simple_thread_pool;

use std::env;
use rust_util::XResult;
use acme_lib::Directory;
use acme_lib::{create_p384_key, create_p256_key, create_rsa_key};
use acme_lib::persist::FilePersist;
use clap::{App, Arg};
use std::fs;
use std::str::FromStr;
use std::sync::RwLock;
use std::collections::BTreeMap;
use tide::Request;
use std::process::exit;
use std::time::{Duration, SystemTime};
use async_std::task;
use async_std::channel;
use async_std::channel::Sender;
use config::AcmeMode;
use crate::config::{CertConfig, CERT_NAME, KEY_NAME};
use crate::x509::{X509PublicKeyAlgo, X509EcPublicKeyAlgo};
use std::path::PathBuf;

const NAME: &str = env!("CARGO_PKG_NAME");
const VERSION: &str = env!("CARGO_PKG_VERSION");
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
const DESCRIPTION: &str = env!("CARGO_PKG_DESCRIPTION");

lazy_static! {
    static ref TOKEN_MAP: RwLock<BTreeMap<String, String>> = RwLock::new(BTreeMap::new());
}

#[derive(Debug, Default)]
struct AcmeRequest<'a> {
    contract_email: &'a str,
    primary_name: &'a str,
    alt_names: &'a [&'a str],
    algo: X509PublicKeyAlgo,
    mode: AcmeMode,
    account_dir: &'a str,
    timeout: u64,
    key_file: Option<String>,
    cert_file: Option<String>,
}

#[async_std::main]
async fn main() -> tide::Result<()> {
    let matches = App::new(NAME)
        .version(VERSION)
        .about(DESCRIPTION)
        .author(AUTHORS)
        .arg(Arg::with_name("version").short("V").long("version").help("Print version"))
        .arg(Arg::with_name("verbose").short("v").long("verbose").help("Verbose"))
        .arg(Arg::with_name("type").short("t").long("type").default_value("http").takes_value(true).help("Type http or dns"))
        .arg(Arg::with_name("port").short("p").long("port").default_value("80").takes_value(true).help("Http port"))
        .arg(Arg::with_name("domain").short("d").long("domain").multiple(true).takes_value(true).help("Domains"))
        .arg(Arg::with_name("email").long("email").takes_value(true).help("Contract email"))
        .arg(Arg::with_name("algo").short("a").long("algo").takes_value(true).default_value("ec384").help("Pki algo"))
        .arg(Arg::with_name("timeout").long("timeout").takes_value(true).default_value("5000").help("Timeout (ms)"))
        .arg(Arg::with_name("mode").short("m").long("mode").takes_value(true).default_value("prod").help("Mode"))
        .arg(Arg::with_name("dir").long("dir").takes_value(true).default_value("acme_dir").help("Account key dir"))
        .arg(Arg::with_name("config").short("c").long("config").takes_value(true).help("Cert config"))
        .arg(Arg::with_name("check").long("check").help("Check cert config"))
        .arg(Arg::with_name("hide-logo").long("hide-logo").help("Hide logo"))
        .get_matches();

    if matches.is_present("verbose") {
        env::set_var("LOGGER_LEVEL", "*");
    }

    if matches.is_present("version") {
        println!("{}", include_str!("logo.txt"));
        information!("{} v{}", NAME, VERSION);
        exit(1);
    }

    if !matches.is_present("hide-logo") {
        println!("{}", include_str!("logo.txt"));
    }

    debugging!("Clap matches: {:?}", matches);

    let account_dir = matches.value_of("dir").unwrap_or("acme_dir");

    let mut account_email = PathBuf::from(account_dir);
    account_email.push("account_email.conf");

    let email = if account_email.exists() {
        match fs::read_to_string(&account_email) {
            Err(e) => {
                failure!("Read from file: {:?}, failed: {}", account_email, e);
                exit(1);
            }
            Ok(email) => {
                if let Some(email_from_args) = matches.value_of("email") {
                    if &email != email_from_args {
                        warning!("Get email from account config: {}", email);
                    }
                }
                email.trim().to_string()
            }
        }
    } else {
        let email = matches.value_of("email").unwrap_or_else(|| {
            failure!("Email is not assigned.");
            exit(1);
        });

        information!("Write email to account config: {:?}", account_email);
        fs::write(account_email, email);

        email.to_string()
    };

    match matches.value_of("type") {
        Some("http") => {}
        _ => {
            failure!("Type is not assigned or must be http.");
            exit(1);
        }
    }
    let port: u16 = match matches.value_of("port") {
        Some(p) => p.parse().unwrap_or_else(|e| {
            failure!("Parse port: {}, failed: {}", p, e);
            exit(1);
        }),
        None => {
            failure!("Port is not assigned.");
            exit(1);
        }
    };
    let timeout: u64 = match matches.value_of("timeout") {
        Some(p) => p.parse().unwrap_or_else(|e| {
            failure!("Parse timeout: {}, failed: {}", p, e);
            exit(1);
        }),
        None => {
            failure!("Timeout is not assigned.");
            exit(1);
        }
    };
    let algo = match matches.value_of("algo") {
        Some(a) => X509PublicKeyAlgo::from_str(a).unwrap_or_else(|e| {
            failure!("{}", e);
            exit(1);
        }),
        _ => {
            failure!("Algo is not assigned, should be: ec256, ec384, rsa2048, rsa3073 or rsa4096.");
            exit(1);
        }
    };
    let mode = match matches.value_of("mode") {
        Some(m) => AcmeMode::parse(m).unwrap_or_else(|e| {
            failure!("{}", e);
            exit(1);
        }),
        _ => {
            failure!("AcmeMode is not assigned, should be: prod or test");
            exit(1);
        }
    };

    let check = matches.is_present("check");

    if !check {
        let (s, r) = channel::bounded(1);
        startup_http_server(s, port);
        r.recv().await.ok();
        task::sleep(Duration::from_millis(500)).await;
    }

    let cert_config = matches.value_of("config");
    match cert_config {
        None => { // cert config is not assigned
            let domains_val = matches.values_of("domain").unwrap_or_else(|| {
                failure!("Domains is not assigned.");
                exit(1);
            });
            let domains: Vec<&str> = domains_val.collect::<Vec<_>>();
            let primary_name = domains[0];
            let alt_names: Vec<&str> = domains.into_iter().skip(1).collect();
            information!("Domains, main: {}, alt: {:?}", primary_name, alt_names);
            let acme_request = AcmeRequest {
                contract_email: &email,
                primary_name,
                alt_names: &alt_names,
                algo,
                mode,
                account_dir,
                timeout,
                ..Default::default()
            };
            if let Err(e) = request_acme_certificate(acme_request) {
                failure!("Request certificate by acme failed: {}", e);
                exit(1);
            }
        }
        Some(cert_config) => { // cert config is assigned
            let cert_config = {
                CertConfig::load(cert_config).unwrap_or_else(|e| {
                    failure!("Load cert config: {}, failed: {}", cert_config, e);
                    exit(1);
                })
            };
            if check {
                check_cert_config(&cert_config);
                return Ok(());
            }
            let filtered_cert_config = cert_config.filter_cert_config_items(30);
            for item in &filtered_cert_config.cert_items {
                if item.common_name.as_ref().map(|n| n.contains('*')).unwrap_or(false)
                    || item.dns_names.as_ref().map(|dns_names| dns_names.iter().any(|n| n.contains('*'))).unwrap_or(false) {
                    warning!("Currently not support wide card domain name");
                    continue;
                }
                if let (Some(common_name), Some(dns_names)) = (&item.common_name, &item.dns_names) {
                    information!("Domains, main: {}, alt: {:?}", common_name, dns_names);
                    let alt_names: Vec<&str> = dns_names.iter().map(|n| n.as_str()).collect();
                    let acme_request = AcmeRequest {
                        contract_email: &email,
                        primary_name: common_name,
                        alt_names: &alt_names,
                        algo,
                        mode,
                        account_dir,
                        timeout,
                        cert_file: Some(format!("{}/{}", item.path, CERT_NAME)),
                        key_file: Some(format!("{}/{}", item.path, KEY_NAME)),
                    };
                    if let Err(e) = request_acme_certificate(acme_request) {
                        failure!("Request certificate: {}, by acme failed: {}", item.path, e);
                    }
                }
            }
        }
    }

    Ok(())
}

fn check_cert_config(cert_config: &CertConfig) {
    let secs_from_unix_epoch = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs() as i64;
    let item_count = cert_config.cert_items.len();
    for (i, item) in cert_config.cert_items.iter().enumerate() {
        information!("Checking: {}, item {} of {}", item.path, i, item_count);
        let cert_fn = format!("{}/{}", item.path, CERT_NAME);
        let pem = match fs::read_to_string(&cert_fn) {
            Ok(pem) => pem,
            Err(e) => {
                warning!("Read file: {}, failed: {}", cert_fn, e);
                continue;
            }
        };
        let x509_certificate = match x509::parse_x509(&cert_fn, &pem) {
            Ok(cert) => cert,
            Err(e) => {
                failure!("Parse x509 file: {}, failed: {}", cert_fn, e);
                continue;
            }
        };
        success!("Found certificate: common name: {}, dns names: {:?}, public key algo: {:?}, valid days: {}",
            x509_certificate.common_name,
            x509_certificate.alt_names,
            x509_certificate.public_key_algo,
            (x509_certificate.certificate_not_after - secs_from_unix_epoch) / (24 * 3600)
        );
    }
}

fn request_acme_certificate(acme_request: AcmeRequest) -> XResult<()> {
    information!("Acme mode: {:?}", acme_request.mode);
    let url = acme_request.mode.directory_url();
    information!("Acme dir: {}", acme_request.account_dir);
    fs::create_dir_all(acme_request.account_dir).ok();
    let persist = FilePersist::new(acme_request.account_dir);
    let dir = opt_result!(Directory::from_url(persist, url), "Create directory from url failed: {}");
    let acc = opt_result!(dir.account(acme_request.contract_email), "Directory set account failed: {}");
    let mut ord_new = opt_result!( acc.new_order(acme_request.primary_name, acme_request.alt_names), "Create order failed: {}");

    let mut order_csr_index = 0;
    let ord_csr = loop {
        if let Some(ord_csr) = ord_new.confirm_validations() {
            debugging!("Valid acme certificate http challenge success");
            break ord_csr;
        }

        information!("Loop for acme challenge auth, #{}", order_csr_index);
        order_csr_index += 1;

        debugging!("Start acme certificate http challenge");
        let auths = opt_result!(ord_new.authorizations(), "Order auth failed: {}");
        for auth in &auths {
            let chall = auth.http_challenge();
            let token = chall.http_token();
            let proof = chall.http_proof();

            {
                information!("Add acme http challenge: {} -> {}",token, proof);
                TOKEN_MAP.write().unwrap().insert(token.to_string(), proof);
            }
            debugging!("Valid acme certificate http challenge");
            opt_result!(chall.validate(acme_request.timeout), "Validate http challenge failed: {}");
        }

        debugging!("Refresh acme certificate order");
        opt_result!(ord_new.refresh(), "Refresh order failed: {}");
    };

    information!("Generate private key, key type: {:?}", acme_request.algo);
    let pkey_pri = match acme_request.algo {
        X509PublicKeyAlgo::EcKey(X509EcPublicKeyAlgo::Secp256r1) => create_p256_key(),
        X509PublicKeyAlgo::EcKey(X509EcPublicKeyAlgo::Secp384r1) => create_p384_key(),
        X509PublicKeyAlgo::EcKey(X509EcPublicKeyAlgo::Secp521r1) => return simple_error!("Algo ec521 is not supported"),
        X509PublicKeyAlgo::Rsa(bits) => create_rsa_key(bits),
    };

    debugging!("Invoking csr finalize pkey");
    let ord_cert = opt_result!( ord_csr.finalize_pkey(pkey_pri, acme_request.timeout), "Submit CSR failed: {}");
    debugging!("Downloading and save cert");
    let cert = opt_result!( ord_cert.download_and_save_cert(), "Download and save certificate failed: {}");

    if let (Some(cert_file), Some(key_file)) = (&acme_request.cert_file, &acme_request.key_file) {
        debugging!("Certificate key: {}",  cert.private_key());
        debugging!("Certificate pem: {}",  cert.certificate());
        information!("Write file: {}", cert_file);
        if let Err(e) = fs::write(cert_file, cert.certificate()) {
            failure!("Write file: {}, failed: {}", cert_file, e);
        }
        information!("Write file: {}", key_file);
        if let Err(e) = fs::write(key_file, cert.private_key()) {
            failure!("Write file: {}, failed: {}", key_file, e);
        }
        success!("Write files success: {} and {}", cert_file, key_file);
    } else {
        information!("Certificate key: {}",  cert.private_key());
        information!("Certificate pem: {}",  cert.certificate());
    }

    Ok(())
}

fn startup_http_server(s: Sender<i32>, port: u16) {
    task::spawn(async move {
        information!("Listen at 0.0.0.0:{}", port);
        let mut app = tide::new();
        app.at("/.well-known/acme-challenge/:token").get(|req: Request<()>| async move {
            let token = match req.param("token") {
                Ok(token) => token,
                Err(e) => {
                    warning!("Cannot get token from url, query: {:?}, error: {}", req.url().query(), e);
                    return Ok("400 - bad request".to_string());
                }
            };
            let peer = req.peer_addr().unwrap_or("none");
            let auth_token = { TOKEN_MAP.read().unwrap().get(token).cloned() };
            match auth_token {
                Some(auth_token) => {
                    information!("Request acme challenge: {} -> {}, peer: {:?}", token, auth_token, peer);
                    Ok(auth_token)
                }
                None => {
                    warning!("Request acme challenge not found: {}, peer: {:?}", token, peer);
                    Ok("404 - not found".to_string())
                }
            }
        });
        s.send(1).await.ok();
        if let Err(e) = app.listen(&format!("0.0.0.0:{}", port)).await {
            failure!("Failed to listen 0.0.0.0:{}, program will exit, error: {}", port, e);
            exit(1);
        }
    });
}