use log::debug;
use rsa::{PublicKey, RSAPublicKey};

use super::*;

impl KeyExt for RSAPublicKey {
    fn to_pem_pkcs8(&self) -> Result<String, KeyError> {
        let bytes = write_pkcs1(self);
        let oid = yasna::models::ObjectIdentifier::from_slice(&RSA_OID);

        let contents = yasna::construct_der(|writer| {
            writer.write_sequence(|writer| {
                writer.next().write_sequence(|writer| {
                    writer.next().write_oid(&oid);
                    writer.next().write_null();
                });
                writer
                    .next()
                    .write_bitvec(&bit_vec::BitVec::from_bytes(&bytes));
            });
        });

        let p = pem::Pem {
            tag: "PUBLIC KEY".to_owned(),
            contents,
        };

        Ok(pem::encode(&p))
    }

    fn from_pem_pkcs8(pem: &str) -> Result<Self, KeyError> {
        let data = pem::parse(pem).map_err(|_| KeyError::Pem)?;

        if data.tag != "PUBLIC KEY" {
            return Err(KeyError::Kind.into());
        }

        let expected_oid = yasna::models::ObjectIdentifier::from_slice(&RSA_OID);

        let pkey = yasna::parse_der(&data.contents, |reader| {
            let o = reader.read_sequence(|reader| {
                debug!("Parse OID");
                let oid = reader.next().read_sequence(|reader| {
                    // TODO: parse more in here
                    let oid = reader.next().read_oid()?;
                    reader.next().read_null()?;
                    Ok(oid)
                })?;

                if oid != expected_oid {
                    debug!("OID was unexpected, {}", oid);
                    return Err(yasna::ASN1Error::new(yasna::ASN1ErrorKind::Invalid));
                }

                debug!("Parse bytes");
                let bitvec = reader.next().read_bitvec()?;

                parse_pkcs1(&bitvec.to_bytes())
            })?;
            debug!("Parsed sequence");
            Ok(o)
        })?;

        Ok(pkey)
    }

    fn to_pem_pkcs1(&self) -> Result<String, KeyError> {
        let contents = write_pkcs1(self);

        let p = pem::Pem {
            tag: "RSA PUBLIC KEY".to_owned(),
            contents,
        };

        Ok(pem::encode(&p))
    }

    fn from_pem_pkcs1(pem: &str) -> Result<Self, KeyError> {
        let data = pem::parse(pem).map_err(|_| KeyError::Pem)?;

        if data.tag != "RSA PUBLIC KEY" {
            return Err(KeyError::Kind.into());
        }

        let pkey = parse_pkcs1(&data.contents)?;

        Ok(pkey)
    }
}

fn write_pkcs1(rsa: &RSAPublicKey) -> Vec<u8> {
    yasna::construct_der(|writer| {
        writer.write_sequence(|writer| {
            writer.next().write_biguint(&from_dig(rsa.n()));
            writer.next().write_biguint(&from_dig(rsa.e()));
        })
    })
}

fn parse_pkcs1(bytes: &[u8]) -> Result<RSAPublicKey, yasna::ASN1Error> {
    let (n, e) = yasna::parse_der(bytes, |reader| {
        reader.read_sequence(|reader| {
            let n = reader.next().read_biguint()?;
            let e = reader.next().read_biguint()?;

            Ok((n, e))
        })
    })?;

    RSAPublicKey::new(to_dig(&n), to_dig(&e))
        .map_err(|_| yasna::ASN1Error::new(yasna::ASN1ErrorKind::Invalid))
}