use std::io::{BufReader, Error, ErrorKind, Write};
use rcgen::{Certificate, CertificateParams, DnType, KeyPair};
use rustls_pemfile::{Item, read_all, read_one};
use tokio_rustls::rustls::{ClientConfig, RootCertStore};
use tokio_rustls::rustls::pki_types::{CertificateDer, PrivatePkcs8KeyDer};
use x509_parser::nom::AsBytes;
use pem::{self, Pem};

pub struct RawCertPair {
    pub cert: Vec<u8>,
    pub prkey: Vec<u8>,
}

// TODO: Eventually consider shifting from CertificateDer to x509_parser::X509Certificate
//  for the extra information it can provide runtime
pub struct LeafCertPair<'a> {
    cert: CertificateDer<'a>,
    ca_chain: Vec<CertificateDer<'a>>,
    prkey: PrivatePkcs8KeyDer<'a>,
}

impl Clone for LeafCertPair<'_> {
    fn clone(&self) -> Self {
        Self {
            cert: self.cert.clone(),
            ca_chain: self.ca_chain.clone(),
            prkey: self.prkey.clone_key(),
        }
    }
}

impl LeafCertPair<'_> {
    pub fn parse<'a>(cert: Vec<u8>, ca_chain: Vec<Vec<u8>>, prkey: Vec<u8>) -> LeafCertPair<'a> {
        let cert = CertificateDer::from(cert);
        let ca_chain = ca_chain.into_iter().map(CertificateDer::from).collect();
        let prkey = PrivatePkcs8KeyDer::from(prkey);
        LeafCertPair {
            cert,
            ca_chain,
            prkey,
        }
    }

    pub fn save_into_file<P: AsRef<std::path::Path>>(&self, filename: P) -> std::io::Result<()> {
        let mut file = std::fs::File::create(filename)?;
        let mut pems = vec![
            Pem::new("CERTIFICATE", self.cert.as_bytes())
        ];
        for c in self.ca_chain.iter() {
            pems.push(Pem::new("CERTIFICATE", c.as_bytes()));
        }
        pems.push(Pem::new("PRIVATE KEY", self.prkey.secret_pkcs8_der()));
        file.write_all(pem::encode_many(&pems).as_bytes())?;
        Ok(())
    }

    pub fn load_from_file<'a, P: AsRef<std::path::Path>>(filename: P) -> std::io::Result<LeafCertPair<'a>> {
        let file = std::fs::File::open(filename).unwrap();
        let mut buf = BufReader::new(file);
        if let Item::X509Certificate(cert) = read_one(&mut buf).unwrap().unwrap() {
            let parsed_cert = x509_parser::parse_x509_certificate(&cert).unwrap().1;
            if parsed_cert.is_ca() {
                return Err(Error::new(ErrorKind::InvalidInput, "main cert is ca"));
            }
            let mut ca_chain: Vec<CertificateDer> = Vec::new();
            for item in read_all(&mut buf) {
                match item {
                    Ok(Item::X509Certificate(c)) => {
                        let parsed_cert = x509_parser::parse_x509_certificate(&c).unwrap().1;
                        if !parsed_cert.is_ca() {
                            return Err(Error::new(ErrorKind::InvalidInput, "chain cert is not ca"));
                        }
                        ca_chain.push(c);
                    },
                    Ok(Item::Pkcs8Key(prkey)) => {
                        return Ok(LeafCertPair { cert, ca_chain, prkey });
                    }
                    _ => {
                        return Err(Error::new(ErrorKind::InvalidInput, "invalid format"));
                    }
                }
            }
            Err(Error::new(ErrorKind::InvalidInput, "pkcs8key not found"))
        } else {
            Err(Error::new(ErrorKind::InvalidInput, "no main x509 cert"))
        }
    }

    pub fn cert(&self) -> &CertificateDer {
        &self.cert
    }

    pub fn ca_chain(&self) -> &Vec<CertificateDer> {
        &self.ca_chain
    }

    pub fn prkey(&self) -> &PrivatePkcs8KeyDer {
        &self.prkey
    }

    pub fn clone_key(&self) -> PrivatePkcs8KeyDer<'static> {
        self.prkey.clone_key()
    }

    pub fn fullchain<'a>(&self) -> Vec<CertificateDer<'a>> {
        let mut chain: Vec<CertificateDer> = Vec::with_capacity(self.ca_chain.len() + 1);
        chain.push(self.cert.clone().into_owned());
        let mut ca_chain = self.ca_chain.clone().into_iter().map(|c| c.into_owned()).collect();
        chain.append(&mut ca_chain);
        chain
    }

    pub fn to_raw(&self) -> RawCertPair {
        let cert = self.cert.to_vec();
        let prkey = self.prkey.secret_pkcs8_der().to_vec();
        RawCertPair { cert, prkey }
    }

    pub fn to_tlsclientconfig(&self, broker_root_certs: &BrokerRootCerts) -> ClientConfig {
        let broker_root_cert_store = broker_root_certs.to_rootcertstore();
        let cert_chain = self.fullchain();
        ClientConfig::builder()
            .with_root_certificates(broker_root_cert_store)
            .with_client_auth_cert(cert_chain, self.prkey.clone_key().into()).expect("Invalid Cert chain")
    }
}

pub struct CACertPair<'a> {
    cert: CertificateDer<'a>,
    ca_chain: Vec<CertificateDer<'a>>,
    prkey: PrivatePkcs8KeyDer<'a>,
}

impl CACertPair<'_> {
    pub fn load_from_file<'a, P: AsRef<std::path::Path>>(filename: P) -> std::io::Result<CACertPair<'a>> {
        let file = std::fs::File::open(filename).unwrap();
        let mut buf = BufReader::new(file);
        if let Item::X509Certificate(cert) = read_one(&mut buf).unwrap().unwrap() {
            let parsed_cert = x509_parser::parse_x509_certificate(&cert).unwrap().1;
            if !parsed_cert.is_ca() {
                return Err(Error::new(ErrorKind::InvalidInput, "cert is not ca"));
            }
            // TODO: Implement ca_chain reading (for now it's unused)
            if let Item::Pkcs8Key(prkey) = read_one(&mut buf).unwrap().unwrap() {
                Ok(CACertPair { cert, ca_chain: vec![], prkey })
            } else {
                Err(Error::new(ErrorKind::InvalidInput, "no ca pkcs8key"))
            }
        } else {
            Err(Error::new(ErrorKind::InvalidInput, "no ca x509 cert"))
        }
    }

    fn rcgen_keypair(&self) -> KeyPair {
        KeyPair::from_der(self.prkey.secret_pkcs8_der()).unwrap()
    }

    fn rcgen_certificate(&self) -> Certificate {
        Certificate::from_params(CertificateParams::from_ca_cert_der(
            &self.cert,
            self.rcgen_keypair()
        ).unwrap()).unwrap()
    }

    pub fn sign_new_cert(&self, params: CertificateParams) -> LeafCertPair {
        let root_cert = self.rcgen_certificate();
        let certificate = Certificate::from_params(params).unwrap();
        let b_cert = certificate.serialize_der_with_signer(&root_cert).unwrap();
        let b_prkey = certificate.serialize_private_key_der();
        let cert = CertificateDer::from(b_cert);
        let prkey = PrivatePkcs8KeyDer::from(b_prkey);
        let mut ca_chain = Vec::with_capacity(self.ca_chain.len() + 1);
        ca_chain.push(self.cert.clone());
        ca_chain.extend(self.ca_chain.iter().cloned());
        LeafCertPair { cert, ca_chain, prkey }
    }

    pub fn cert(&self) -> &CertificateDer {
        &self.cert
    }

    pub fn clone_key(&self) -> PrivatePkcs8KeyDer {
        self.prkey.clone_key()
    }

    pub fn to_raw(&self) -> RawCertPair {
        let cert = self.cert.to_vec();
        let prkey = self.prkey.secret_pkcs8_der().to_vec();
        RawCertPair { cert, prkey }
    }
}

pub fn server_leaf_certparams(name: &str) -> CertificateParams {
    let mut params = CertificateParams::new(vec!["entity.other.host".into(), format!("bonk.server.{name}")]);
    params.distinguished_name.push(DnType::CommonName, name);
    params.use_authority_key_identifier_extension = true;
    params.key_usages.push(rcgen::KeyUsagePurpose::DigitalSignature);
    params.extended_key_usages.push(rcgen::ExtendedKeyUsagePurpose::ClientAuth);
    params
}

#[derive(Clone)]
pub struct BrokerRootCerts<'a> {
    root_cert: CertificateDer<'a>
}

impl BrokerRootCerts<'_> {
    pub fn load_from_file<'a, P: AsRef<std::path::Path>>(filename: P) -> std::io::Result<BrokerRootCerts<'a>> {
        let file = std::fs::File::open(filename).unwrap();
        let mut buf = BufReader::new(file);
        if let Item::X509Certificate(root_cert) = read_one(&mut buf).unwrap().unwrap() {
            Ok(BrokerRootCerts { root_cert })
        } else {
            Err(Error::new(ErrorKind::InvalidInput, "no broker root x509 cert"))
        }
    }

    pub fn to_rootcertstore(&self) -> RootCertStore {
        let mut broker_root_cert_store = RootCertStore::empty();
        broker_root_cert_store.add(self.root_cert.clone()).expect("Invalid Broker Root");
        broker_root_cert_store
    }

    pub fn certs(&self) -> Vec<&CertificateDer> {
        vec![&self.root_cert]
    }
}