use std::cmp::{Ordering, min};
use std::collections::{BTreeMap, BTreeSet};
use std::net::IpAddr;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use futures::FutureExt;
use imbl_value::InternedString;
use libc::time_t;
use openssl::asn1::{Asn1Integer, Asn1Time, Asn1TimeRef};
use openssl::bn::{BigNum, MsbOption};
use openssl::ec::{EcGroup, EcKey};
use openssl::error::ErrorStack;
use openssl::hash::MessageDigest;
use openssl::pkey::{PKey, Private};
use openssl::x509::extension::{
    AuthorityKeyIdentifier, BasicConstraints, KeyUsage, SubjectAlternativeName,
    SubjectKeyIdentifier,
};
use openssl::x509::{X509, X509Builder, X509NameBuilder, X509Ref};
use openssl::*;
use patch_db::HasModel;
use serde::{Deserialize, Serialize};
use tokio_rustls::rustls::ServerConfig;
use tokio_rustls::rustls::crypto::CryptoProvider;
use tokio_rustls::rustls::pki_types::{PrivateKeyDer, PrivatePkcs8KeyDer};
use tokio_rustls::rustls::server::ClientHello;
use tracing::instrument;
use visit_rs::Visit;

use crate::SOURCE_DATE;
use crate::account::AccountInfo;
use crate::db::model::Database;
use crate::db::{DbAccess, DbAccessMut};
use crate::hostname::ServerHostname;
use crate::init::check_time_is_synchronized;
use crate::net::gateway::GatewayInfo;
use crate::net::tls::{TlsHandler, TlsHandlerAction};
use crate::net::web_server::{Accept, ExtractVisitor, TcpMetadata, extract};
use crate::prelude::*;
use crate::util::serde::Pem;

pub fn should_use_cert(cert: &X509Ref) -> Result<bool, ErrorStack> {
    Ok(cert
        .not_before()
        .compare(Asn1Time::days_from_now(0)?.as_ref())?
        == Ordering::Less
        && cert
            .not_after()
            .compare(Asn1Time::days_from_now(30)?.as_ref())?
            == Ordering::Greater)
}

#[derive(Debug, Deserialize, Serialize, HasModel)]
#[model = "Model<Self>"]
#[serde(rename_all = "camelCase")]
pub struct CertStore {
    pub root_key: Pem<PKey<Private>>,
    pub root_cert: Pem<X509>,
    pub int_key: Pem<PKey<Private>>,
    pub int_cert: Pem<X509>,
    pub leaves: BTreeMap<JsonKey<BTreeSet<InternedString>>, CertData>,
}
impl CertStore {
    pub fn new(account: &AccountInfo) -> Result<Self, Error> {
        let int_key = gen_nistp256()?;
        let int_cert = make_int_cert((&account.root_ca_key, &account.root_ca_cert), &int_key)?;
        Ok(Self {
            root_key: Pem::new(account.root_ca_key.clone()),
            root_cert: Pem::new(account.root_ca_cert.clone()),
            int_key: Pem::new(int_key),
            int_cert: Pem::new(int_cert),
            leaves: BTreeMap::new(),
        })
    }
}
impl Model<CertStore> {
    /// This function will grant any cert for any domain. It is up to the *caller* to enusure that the calling service has permission to sign a cert for the requested domain
    pub fn cert_for(
        &mut self,
        hostnames: &BTreeSet<InternedString>,
    ) -> Result<FullchainCertData, Error> {
        let keys = if let Some(cert_data) = self
            .as_leaves()
            .as_idx(JsonKey::new_ref(hostnames))
            .map(|m| m.de())
            .transpose()?
        {
            if should_use_cert(&cert_data.certs.ed25519)?
                && should_use_cert(&cert_data.certs.nistp256)?
            {
                return Ok(FullchainCertData {
                    root: self.as_root_cert().de()?.0,
                    int: self.as_int_cert().de()?.0,
                    leaf: cert_data,
                });
            }
            cert_data.keys
        } else {
            PKeyPair {
                ed25519: PKey::generate_ed25519()?,
                nistp256: gen_nistp256()?,
            }
        };
        let int_key = self.as_int_key().de()?.0;
        let int_cert = self.as_int_cert().de()?.0;
        let cert_data = CertData {
            certs: CertPair {
                ed25519: make_leaf_cert(
                    (&int_key, &int_cert),
                    (&keys.ed25519, &SANInfo::new(hostnames)),
                )?,
                nistp256: make_leaf_cert(
                    (&int_key, &int_cert),
                    (&keys.nistp256, &SANInfo::new(hostnames)),
                )?,
            },
            keys,
        };
        self.as_leaves_mut()
            .insert(JsonKey::new_ref(hostnames), &cert_data)?;
        Ok(FullchainCertData {
            root: self.as_root_cert().de()?.0,
            int: self.as_int_cert().de()?.0,
            leaf: cert_data,
        })
    }
}
impl DbAccess<CertStore> for Database {
    fn access<'a>(db: &'a Model<Self>) -> &'a Model<CertStore> {
        db.as_private().as_key_store().as_local_certs()
    }
}
impl DbAccessMut<CertStore> for Database {
    fn access_mut<'a>(db: &'a mut Model<Self>) -> &'a mut Model<CertStore> {
        db.as_private_mut().as_key_store_mut().as_local_certs_mut()
    }
}

#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub struct CertData {
    pub keys: PKeyPair,
    pub certs: CertPair,
}
impl CertData {
    pub fn expiration(&self) -> Result<SystemTime, Error> {
        self.certs.expiration()
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FullchainCertData {
    pub root: X509,
    pub int: X509,
    pub leaf: CertData,
}
impl FullchainCertData {
    pub fn fullchain_ed25519(&self) -> Vec<&X509> {
        vec![&self.leaf.certs.ed25519, &self.int, &self.root]
    }
    pub fn fullchain_nistp256(&self) -> Vec<&X509> {
        vec![&self.leaf.certs.nistp256, &self.int, &self.root]
    }
    pub fn expiration(&self) -> Result<SystemTime, Error> {
        [
            asn1_time_to_system_time(self.root.not_after())?,
            asn1_time_to_system_time(self.int.not_after())?,
            self.leaf.expiration()?,
        ]
        .into_iter()
        .min()
        .ok_or_else(|| Error::new(eyre!("{}", t!("net.ssl.unreachable")), ErrorKind::Unknown))
    }
}

static CERTIFICATE_VERSION: i32 = 2; // X509 version 3 is actually encoded as '2' in the cert because fuck you.

fn unix_time(time: SystemTime) -> time_t {
    time.duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs() as time_t)
        .or_else(|_| UNIX_EPOCH.elapsed().map(|d| -(d.as_secs() as time_t)))
        .unwrap_or_default()
}

lazy_static::lazy_static! {
    static ref ASN1_UNIX_EPOCH: Asn1Time = Asn1Time::from_unix(0).unwrap();
}

fn asn1_time_to_system_time(time: &Asn1TimeRef) -> Result<SystemTime, Error> {
    let diff = time.diff(&**ASN1_UNIX_EPOCH)?;
    let mut res = UNIX_EPOCH;
    if diff.days >= 0 {
        res += Duration::from_secs(diff.days as u64 * 86400);
    } else {
        res -= Duration::from_secs((-1 * diff.days) as u64 * 86400);
    }
    if diff.secs >= 0 {
        res += Duration::from_secs(diff.secs as u64);
    } else {
        res -= Duration::from_secs((-1 * diff.secs) as u64);
    }
    Ok(res)
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PKeyPair {
    #[serde(with = "crate::util::serde::pem")]
    pub ed25519: PKey<Private>,
    #[serde(with = "crate::util::serde::pem")]
    pub nistp256: PKey<Private>,
}
impl PartialEq for PKeyPair {
    fn eq(&self, other: &Self) -> bool {
        self.ed25519.public_eq(&other.ed25519) && self.nistp256.public_eq(&other.nistp256)
    }
}
impl Eq for PKeyPair {}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Deserialize, Serialize)]
pub struct CertPair {
    #[serde(with = "crate::util::serde::pem")]
    pub ed25519: X509,
    #[serde(with = "crate::util::serde::pem")]
    pub nistp256: X509,
}
impl CertPair {
    pub fn expiration(&self) -> Result<SystemTime, Error> {
        Ok(min(
            asn1_time_to_system_time(self.ed25519.not_after())?,
            asn1_time_to_system_time(self.nistp256.not_after())?,
        ))
    }
}

pub async fn root_ca_start_time() -> SystemTime {
    if check_time_is_synchronized()
        .await
        .log_err()
        .unwrap_or(false)
    {
        SystemTime::now()
    } else {
        *SOURCE_DATE
    }
}

const EC_CURVE_NAME: nid::Nid = nid::Nid::X9_62_PRIME256V1;
lazy_static::lazy_static! {
    static ref EC_GROUP: EcGroup = EcGroup::from_curve_name(EC_CURVE_NAME).unwrap();
}

pub async fn export_key(key: &PKey<Private>, target: &Path) -> Result<(), Error> {
    tokio::fs::write(target, key.private_key_to_pem_pkcs8()?)
        .map(|res| res.with_ctx(|_| (ErrorKind::Filesystem, target.display().to_string())))
        .await?;
    Ok(())
}
pub async fn export_cert(chain: &[&X509], target: &Path) -> Result<(), Error> {
    tokio::fs::write(
        target,
        chain
            .into_iter()
            .flat_map(|c| c.to_pem().unwrap())
            .collect::<Vec<u8>>(),
    )
    .await?;
    Ok(())
}

#[instrument(skip_all)]
fn rand_serial() -> Result<Asn1Integer, Error> {
    let mut bn = BigNum::new()?;
    bn.rand(64, MsbOption::MAYBE_ZERO, false)?;
    let asn1 = Asn1Integer::from_bn(&bn)?;
    Ok(asn1)
}
#[instrument(skip_all)]
pub fn gen_nistp256() -> Result<PKey<Private>, Error> {
    Ok(PKey::from_ec_key(EcKey::generate(EC_GROUP.as_ref())?)?)
}

#[instrument(skip_all)]
pub fn make_root_cert(
    root_key: &PKey<Private>,
    hostname: &ServerHostname,
    start_time: SystemTime,
) -> Result<X509, Error> {
    let mut builder = X509Builder::new()?;
    builder.set_version(CERTIFICATE_VERSION)?;

    let unix_start_time = unix_time(start_time);

    let embargo = Asn1Time::from_unix(unix_start_time - 86400)?;
    builder.set_not_before(&embargo)?;

    let expiration = Asn1Time::from_unix(unix_start_time + (10 * 364 * 86400))?;
    builder.set_not_after(&expiration)?;

    builder.set_serial_number(&*rand_serial()?)?;

    let mut subject_name_builder = X509NameBuilder::new()?;
    subject_name_builder
        .append_entry_by_text("CN", &format!("{} Local Root CA", hostname.as_ref()))?;
    subject_name_builder.append_entry_by_text("O", "Start9")?;
    subject_name_builder.append_entry_by_text("OU", "StartOS")?;
    let subject_name = subject_name_builder.build();
    builder.set_subject_name(&subject_name)?;

    builder.set_issuer_name(&subject_name)?;

    builder.set_pubkey(&root_key)?;

    // Extensions
    let cfg = conf::Conf::new(conf::ConfMethod::default())?;
    let ctx = builder.x509v3_context(None, Some(&cfg));
    // subjectKeyIdentifier = hash
    let subject_key_identifier = SubjectKeyIdentifier::new().build(&ctx)?;
    // authorityKeyIdentifier = keyid,issuer:always
    let authority_key_identifier = AuthorityKeyIdentifier::new()
        .keyid(false)
        .issuer(true)
        .build(&ctx)?;
    // basicConstraints = critical, CA:true, pathlen:0
    let basic_constraints = BasicConstraints::new().critical().ca().build()?;
    // keyUsage = critical, digitalSignature, cRLSign, keyCertSign
    let key_usage = KeyUsage::new()
        .critical()
        .digital_signature()
        .crl_sign()
        .key_cert_sign()
        .build()?;
    builder.append_extension(subject_key_identifier)?;
    builder.append_extension(authority_key_identifier)?;
    builder.append_extension(basic_constraints)?;
    builder.append_extension(key_usage)?;
    builder.sign(&root_key, MessageDigest::sha256())?;
    let cert = builder.build();
    Ok(cert)
}
#[instrument(skip_all)]
pub fn make_int_cert(
    signer: (&PKey<Private>, &X509),
    applicant: &PKey<Private>,
) -> Result<X509, Error> {
    let mut builder = X509Builder::new()?;
    builder.set_version(CERTIFICATE_VERSION)?;

    builder.set_not_before(signer.1.not_before())?;

    builder.set_not_after(signer.1.not_after())?;

    builder.set_serial_number(&*rand_serial()?)?;

    let mut subject_name_builder = X509NameBuilder::new()?;
    subject_name_builder.append_entry_by_text("CN", "StartOS Local Intermediate CA")?;
    subject_name_builder.append_entry_by_text("O", "Start9")?;
    subject_name_builder.append_entry_by_text("OU", "StartOS")?;
    let subject_name = subject_name_builder.build();
    builder.set_subject_name(&subject_name)?;

    builder.set_issuer_name(signer.1.subject_name())?;

    builder.set_pubkey(&applicant)?;

    let cfg = conf::Conf::new(conf::ConfMethod::default())?;
    let ctx = builder.x509v3_context(Some(&signer.1), Some(&cfg));

    // subjectKeyIdentifier = hash
    let subject_key_identifier = SubjectKeyIdentifier::new().build(&ctx)?;
    // authorityKeyIdentifier = keyid:always,issuer:always
    let authority_key_identifier = AuthorityKeyIdentifier::new()
        .keyid(true)
        .issuer(true)
        .build(&ctx)?;
    // basicConstraints = critical, CA:true, pathlen:0
    let basic_constraints = BasicConstraints::new().critical().ca().pathlen(0).build()?;
    // keyUsage = critical, digitalSignature, cRLSign, keyCertSign
    let key_usage = KeyUsage::new()
        .critical()
        .digital_signature()
        .crl_sign()
        .key_cert_sign()
        .build()?;
    builder.append_extension(subject_key_identifier)?;
    builder.append_extension(authority_key_identifier)?;
    builder.append_extension(basic_constraints)?;
    builder.append_extension(key_usage)?;
    builder.sign(&signer.0, MessageDigest::sha256())?;
    let cert = builder.build();
    Ok(cert)
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
pub enum MaybeWildcard {
    WithWildcard(String),
    WithoutWildcard(InternedString),
}
impl MaybeWildcard {
    pub fn as_str(&self) -> &str {
        match self {
            MaybeWildcard::WithWildcard(s) => s.as_str(),
            MaybeWildcard::WithoutWildcard(s) => &**s,
        }
    }
}
impl std::fmt::Display for MaybeWildcard {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MaybeWildcard::WithWildcard(dns) => write!(f, "DNS:{dns},DNS:*.{dns}"),
            MaybeWildcard::WithoutWildcard(dns) => write!(f, "DNS:{dns}"),
        }
    }
}

#[derive(Debug)]
pub struct SANInfo {
    pub dns: BTreeSet<MaybeWildcard>,
    pub ips: BTreeSet<IpAddr>,
}
impl SANInfo {
    pub fn new(hostnames: &BTreeSet<InternedString>) -> Self {
        let mut dns = BTreeSet::new();
        let mut ips = BTreeSet::new();
        for hostname in hostnames {
            if let Ok(ip) = hostname.parse::<IpAddr>() {
                ips.insert(ip);
            } else {
                dns.insert(MaybeWildcard::WithoutWildcard(hostname.clone())); // TODO: wildcards?
            }
        }
        Self { dns, ips }
    }
    pub fn x509_extension(&self) -> SubjectAlternativeName {
        let mut san = SubjectAlternativeName::new();
        for h in &self.dns {
            san.dns(&h.as_str());
        }
        for ip in &self.ips {
            san.ip(&ip.to_string());
        }
        san
    }
}
impl std::fmt::Display for SANInfo {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut written = false;
        for dns in &self.dns {
            if written {
                write!(f, ",")?;
            }
            written = true;
            write!(f, "{dns}")?;
        }
        for ip in &self.ips {
            if written {
                write!(f, ",")?;
            }
            written = true;
            write!(f, "IP:{ip}")?;
        }
        Ok(())
    }
}

#[instrument(skip_all)]
pub fn make_leaf_cert(
    signer: (&PKey<Private>, &X509),
    applicant: (&PKey<Private>, &SANInfo),
) -> Result<X509, Error> {
    let mut builder = X509Builder::new()?;
    builder.set_version(CERTIFICATE_VERSION)?;

    let embargo = Asn1Time::from_unix(unix_time(SystemTime::now()) - 86400)?;
    builder.set_not_before(&embargo)?;

    // Google Apple and Mozilla reject certificate horizons longer than 398 days
    // https://techbeacon.com/security/google-apple-mozilla-enforce-1-year-max-security-certifications
    let expiration = Asn1Time::days_from_now(397)?;
    builder.set_not_after(&expiration)?;

    builder.set_serial_number(&*rand_serial()?)?;

    let mut subject_name_builder = X509NameBuilder::new()?;
    subject_name_builder.append_entry_by_text(
        "CN",
        applicant
            .1
            .dns
            .first()
            .map(MaybeWildcard::as_str)
            .unwrap_or("localhost"),
    )?;
    subject_name_builder.append_entry_by_text("O", "Start9")?;
    subject_name_builder.append_entry_by_text("OU", "StartOS")?;
    let subject_name = subject_name_builder.build();
    builder.set_subject_name(&subject_name)?;

    builder.set_issuer_name(signer.1.subject_name())?;

    builder.set_pubkey(&applicant.0)?;

    // Extensions
    let cfg = conf::Conf::new(conf::ConfMethod::default())?;
    let ctx = builder.x509v3_context(Some(&signer.1), Some(&cfg));

    let subject_key_identifier = SubjectKeyIdentifier::new().build(&ctx)?;
    let authority_key_identifier = AuthorityKeyIdentifier::new()
        .keyid(true)
        .issuer(false)
        .build(&ctx)?;
    let subject_alt_name = applicant.1.x509_extension().build(&ctx)?;
    let basic_constraints = BasicConstraints::new().build()?;
    let key_usage = KeyUsage::new()
        .critical()
        .digital_signature()
        .key_encipherment()
        .build()?;

    builder.append_extension(subject_key_identifier)?;
    builder.append_extension(authority_key_identifier)?;
    builder.append_extension(subject_alt_name)?;
    builder.append_extension(basic_constraints)?;
    builder.append_extension(key_usage)?;

    builder.sign(&signer.0, MessageDigest::sha256())?;

    let cert = builder.build();
    Ok(cert)
}

#[instrument(skip_all)]
pub fn make_self_signed(applicant: (&PKey<Private>, &SANInfo)) -> Result<X509, Error> {
    let mut builder = X509Builder::new()?;
    builder.set_version(CERTIFICATE_VERSION)?;

    let embargo = Asn1Time::from_unix(unix_time(SystemTime::now()) - 86400)?;
    builder.set_not_before(&embargo)?;

    // Google Apple and Mozilla reject certificate horizons longer than 398 days
    // https://techbeacon.com/security/google-apple-mozilla-enforce-1-year-max-security-certifications
    let expiration = Asn1Time::days_from_now(397)?;
    builder.set_not_after(&expiration)?;

    builder.set_serial_number(&*rand_serial()?)?;

    let mut subject_name_builder = X509NameBuilder::new()?;
    subject_name_builder.append_entry_by_text(
        "CN",
        applicant
            .1
            .dns
            .first()
            .map(MaybeWildcard::as_str)
            .unwrap_or("localhost"),
    )?;
    subject_name_builder.append_entry_by_text("O", "Start9")?;
    subject_name_builder.append_entry_by_text("OU", "StartOS")?;
    let subject_name = subject_name_builder.build();
    builder.set_subject_name(&subject_name)?;

    builder.set_issuer_name(&subject_name)?;

    builder.set_pubkey(&applicant.0)?;

    // Extensions
    let cfg = conf::Conf::new(conf::ConfMethod::default())?;
    let ctx = builder.x509v3_context(None, Some(&cfg));

    let subject_key_identifier = SubjectKeyIdentifier::new().build(&ctx)?;
    let authority_key_identifier = AuthorityKeyIdentifier::new()
        .keyid(false)
        .issuer(true)
        .build(&ctx)?;
    let subject_alt_name = applicant.1.x509_extension().build(&ctx)?;
    let basic_constraints = BasicConstraints::new().build()?;
    let key_usage = KeyUsage::new()
        .critical()
        .digital_signature()
        .key_encipherment()
        .build()?;

    builder.append_extension(subject_key_identifier)?;
    builder.append_extension(authority_key_identifier)?;
    builder.append_extension(subject_alt_name)?;
    builder.append_extension(basic_constraints)?;
    builder.append_extension(key_usage)?;

    builder.sign(&applicant.0, MessageDigest::sha256())?;

    let cert = builder.build();
    Ok(cert)
}

pub struct RootCaTlsHandler<M: HasModel> {
    pub db: TypedPatchDb<M>,
    pub crypto_provider: Arc<CryptoProvider>,
}
impl<M: HasModel> Clone for RootCaTlsHandler<M> {
    fn clone(&self) -> Self {
        Self {
            db: self.db.clone(),
            crypto_provider: self.crypto_provider.clone(),
        }
    }
}

impl<'a, A, M> TlsHandler<'a, A> for RootCaTlsHandler<M>
where
    A: Accept + 'a,
    <A as Accept>::Metadata: Visit<ExtractVisitor<TcpMetadata>>
        + Visit<ExtractVisitor<GatewayInfo>>
        + Clone
        + Send
        + Sync
        + 'static,
    M: HasModel<Model = Model<M>> + DbAccessMut<CertStore> + Send + Sync,
{
    async fn get_config(
        &mut self,
        hello: &ClientHello<'_>,
        metadata: &<A as Accept>::Metadata,
    ) -> Option<TlsHandlerAction> {
        let hostnames: BTreeSet<InternedString> = hello
            .server_name()
            .map(InternedString::from)
            .into_iter()
            .chain(
                extract::<TcpMetadata, _>(metadata)
                    .map(|m| m.local_addr.ip())
                    .as_ref()
                    .map(InternedString::from_display),
            )
            .chain(
                extract::<GatewayInfo, _>(metadata)
                    .and_then(|i| i.info.ip_info)
                    .and_then(|i| i.wan_ip)
                    .as_ref()
                    .map(InternedString::from_display),
            )
            .collect();
        let cert = self
            .db
            .mutate(|db| M::access_mut(db).cert_for(&hostnames))
            .await
            .result
            .log_err()?;
        let cfg = ServerConfig::builder_with_provider(self.crypto_provider.clone())
            .with_safe_default_protocol_versions()
            .log_err()?
            .with_no_client_auth();
        if hello
            .signature_schemes()
            .contains(&tokio_rustls::rustls::SignatureScheme::ED25519)
        {
            cfg.with_single_cert(
                cert.fullchain_ed25519()
                    .into_iter()
                    .map(|c| {
                        Ok(tokio_rustls::rustls::pki_types::CertificateDer::from(
                            c.to_der()?,
                        ))
                    })
                    .collect::<Result<_, Error>>()
                    .log_err()?,
                PrivateKeyDer::from(PrivatePkcs8KeyDer::from(
                    cert.leaf.keys.ed25519.private_key_to_pkcs8().log_err()?,
                )),
            )
        } else {
            cfg.with_single_cert(
                cert.fullchain_nistp256()
                    .into_iter()
                    .map(|c| {
                        Ok(tokio_rustls::rustls::pki_types::CertificateDer::from(
                            c.to_der()?,
                        ))
                    })
                    .collect::<Result<_, Error>>()
                    .log_err()?,
                PrivateKeyDer::from(PrivatePkcs8KeyDer::from(
                    cert.leaf.keys.nistp256.private_key_to_pkcs8().log_err()?,
                )),
            )
        }
        .log_err()
        .map(TlsHandlerAction::Tls)
    }
}