use std::{
    env,
    str::FromStr,
    time::{Duration, SystemTime},
};

use ed25519_dalek::SigningKey;
use keyfork_derive_util::{
    request::{DerivationAlgorithm, DerivationRequest},
    DerivationIndex, DerivationPath, ExtendedPrivateKey, PrivateKey,
};
use keyforkd_client::Client;
use sequoia_openpgp::{
    packet::{
        key::{Key4, PrimaryRole, SubordinateRole},
        signature::SignatureBuilder,
        Key, UserID,
    },
    types::{KeyFlags, SignatureType},
    Cert, Packet,
};

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("Bad character: {0}")]
    BadChar(char),
}

#[derive(Debug)]
struct KeyType {
    _inner: KeyFlags,
}

impl Default for KeyType {
    fn default() -> Self {
        Self {
            _inner: KeyFlags::empty(),
        }
    }
}

impl KeyType {
    fn certify(mut self) -> Self {
        self._inner = self._inner.set_certification();
        self
    }

    fn sign(mut self) -> Self {
        self._inner = self._inner.set_signing();
        self
    }

    fn encrypt(mut self) -> Self {
        self._inner = self._inner.set_transport_encryption();
        self._inner = self._inner.set_storage_encryption();
        self
    }

    fn authenticate(mut self) -> Self {
        self._inner = self._inner.set_authentication();
        self
    }

    fn inner(&self) -> &KeyFlags {
        &self._inner
    }
}

impl FromStr for KeyType {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        s.chars().try_fold(Self::default(), |s, ch| match ch {
            'C' | 'c' => Ok(s.certify()),
            'S' | 's' => Ok(s.sign()),
            'E' | 'e' => Ok(s.encrypt()),
            'A' | 'a' => Ok(s.authenticate()),
            ch => Err(Error::BadChar(ch)),
        })
    }
}

fn validate(
    path: &str,
    subkey_format: &str,
    default_userid: &str,
) -> Result<(DerivationPath, Vec<KeyType>, UserID), Box<dyn std::error::Error>> {
    let mut pgp_u32 = [0u8; 4];
    pgp_u32[1..].copy_from_slice(&"pgp".bytes().collect::<Vec<u8>>());
    let index = DerivationIndex::new(u32::from_be_bytes(pgp_u32), true)?;

    let path = DerivationPath::from_str(path)?;
    assert_eq!(2, path.len(), "Expected path of m/{index}/account_id'");

    let given_index = path.iter().next().expect("checked .len() above");
    assert_eq!(
        &index, given_index,
        "Expected derivation path starting with m/{index}, got: {given_index}",
    );

    let subkey_format = subkey_format
        .split(',')
        .map(KeyType::from_str)
        .collect::<Result<Vec<_>, Error>>()?;
    assert!(
        subkey_format[0].inner().for_certification(),
        "First key must be able to certify"
    );

    Ok((path, subkey_format, UserID::from(default_userid)))
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let mut args = env::args();
    let program_name = args.next().expect("program name");
    let args = args.collect::<Vec<_>>();
    let (path, subkey_format, default_userid) = match args.as_slice() {
        [path, subkey_format, default_userid] => validate(path, subkey_format, default_userid)?,
        _ => panic!("Usage: {program_name} path subkey_format default_userid"),
    };

    let request = DerivationRequest::new(DerivationAlgorithm::Ed25519, &path);
    let derived_data = Client::discover_socket()?.request(&request)?;
    let xprv = ExtendedPrivateKey::<SigningKey>::try_from(&derived_data)?;

    let derived_key = xprv.derive_child(&DerivationIndex::new(0, true)?)?;
    let epoch = SystemTime::UNIX_EPOCH + Duration::from_secs(1);
    let one_day = SystemTime::now() + Duration::from_secs(60 * 60 * 24);
    let primary_key = Key::from(
        Key4::<_, PrimaryRole>::import_secret_ed25519(
            &PrivateKey::to_bytes(derived_key.private_key()),
            epoch,
        )
        .unwrap(),
    );

    let cert =
        Cert::from_packets(vec![Packet::SecretKey(primary_key.clone())].into_iter()).unwrap();
    let builder = SignatureBuilder::new(SignatureType::PositiveCertification)
        .set_signature_creation_time(epoch)?
        .set_key_flags(subkey_format[0].inner().clone())?;
    let binding = default_userid.bind(&mut primary_key.clone().into_keypair()?, &cert, builder)?;
    let cert = cert.insert_packets(vec![Packet::from(default_userid), binding.into()])?;

    let policy = sequoia_openpgp::policy::StandardPolicy::new();
    let mut keypair = cert
        .primary_key()
        .key()
        .clone()
        .parts_into_secret()?
        .into_keypair()?;
    let signatures = cert.set_expiration_time(&policy, None, &mut keypair, Some(one_day))?;
    let mut cert = cert.insert_packets(signatures)?;

    for (index, subkey_flags) in subkey_format.iter().enumerate().skip(1) {
        let index = u32::try_from(index)?;
        let derived_key = xprv.derive_child(&DerivationIndex::new(index, true)?)?;
        let subkey = Key::from(
            Key4::<_, SubordinateRole>::import_secret_ed25519(
                &PrivateKey::to_bytes(derived_key.private_key()),
                epoch,
            )
            .unwrap(),
        );
        let builder = if subkey_flags.inner().for_signing() {
            SignatureBuilder::new(SignatureType::SubkeyBinding)
                .set_key_flags(subkey_flags.inner().clone())?
                .set_signature_creation_time(epoch)?
                .set_key_validity_period(one_day.duration_since(epoch)?)?
                .set_embedded_signature(
                    SignatureBuilder::new(SignatureType::PrimaryKeyBinding)
                        .set_signature_creation_time(epoch)?
                        .sign_primary_key_binding(
                            &mut subkey.clone().into_keypair()?,
                            &primary_key,
                            &subkey,
                        )?,
                )?
        } else {
            SignatureBuilder::new(SignatureType::SubkeyBinding)
                .set_key_flags(subkey_flags.inner().clone())?
                .set_signature_creation_time(epoch)?
                .set_key_validity_period(one_day.duration_since(epoch)?)?
        };
        let binding =
            builder.sign_subkey_binding(&mut primary_key.clone().into_keypair()?, None, &subkey)?;
        cert = cert.insert_packets(vec![
            Packet::SecretSubkey(subkey.clone()),
            Packet::Signature(binding),
        ])?;
    }

    use sequoia_openpgp::{
        armor::{Kind, Writer},
        serialize::Marshal,
    };

    let mut w = Writer::new(std::io::stdout(), Kind::Message)?;

    for packet in cert.into_packets() {
        packet.serialize(&mut w)?;
    }

    w.finalize()?;

    Ok(())
}