// Copyright (C) 2026 Doan Luu Duc Tai <doan@luuduct.ai>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#![doc = "Experimental Sequoia post-quantum browser compatibility probe."]

use anyhow::{Context, Result, ensure};
use pgp::Profile;
use pgp::cert::prelude::*;
use pgp::crypto::{Password, SessionKey};
use pgp::parse::{
    Parse,
    stream::{
        DecryptionHelper, DecryptorBuilder, MessageLayer, MessageStructure, VerificationHelper,
    },
};
use pgp::policy::StandardPolicy;
use pgp::serialize::{
    Serialize,
    stream::{Encryptor, LiteralWriter, Message, Signer},
};
use pgp::types::{KeyFlags, SymmetricAlgorithm};
use sequoia_openpgp as pgp;
use std::io::Write;

const PASSWORD: &str = "compatibility-test-only-not-a-real-secret";
const TEXT: &[u8] = "stack-chat/v1: UTF-8 message — doan ↔ visitor".as_bytes();

#[doc = "Exercises v6 PQ generation, protected-key import, signing and two-recipient encryption."]
#[doc = "# Errors"]
#[doc = "Returns a diagnostic if a required operation or integrity check fails."]
pub fn compatibility_probe() -> std::result::Result<String, String> {
    probe().map_err(|error| format!("{error:#}"))
}

fn identity() -> Result<Cert> {
    let (cert, _revocation) = CertBuilder::new()
        .set_profile(Profile::RFC9580)?
        .set_cipher_suite(CipherSuite::MLDSA87_Ed448)
        .set_primary_key_flags(KeyFlags::empty().set_certification().set_signing())
        .add_transport_encryption_subkey()
        .set_password(Some(Password::from(PASSWORD)))
        .generate()?;
    let mut encoded = Vec::new();
    cert.as_tsk().serialize(&mut encoded)?;
    Cert::from_bytes(&encoded)
}

fn encrypt(sender: &Cert, other: &Cert) -> Result<Vec<u8>> {
    let policy = StandardPolicy::new();
    let recipients = [sender, other].into_iter().flat_map(|cert| {
        cert.keys()
            .with_policy(&policy, None)
            .supported()
            .alive()
            .revoked(false)
            .for_transport_encryption()
    });
    let signer = sender
        .primary_key()
        .key()
        .parts_as_secret()?
        .clone()
        .decrypt_secret(&Password::from(PASSWORD))?
        .into_keypair()?;
    let mut ciphertext = Vec::new();
    let message = Message::new(&mut ciphertext);
    let encrypted = Encryptor::for_recipients(message, recipients)
        .symmetric_algo(SymmetricAlgorithm::AES256)
        .build()?;
    let authenticated_stream = Signer::new(encrypted, signer)?.build()?;
    let mut literal = LiteralWriter::new(authenticated_stream).build()?;
    literal.write_all(TEXT)?;
    literal.finalize()?;
    Ok(ciphertext)
}

struct Helper<'a> {
    recipient: &'a Cert,
    sender: &'a Cert,
    password: &'a str,
}

impl VerificationHelper for Helper<'_> {
    fn get_certs(&mut self, _ids: &[pgp::KeyHandle]) -> Result<Vec<Cert>> {
        Ok(vec![self.sender.clone().strip_secret_key_material()])
    }
    fn check(&mut self, structure: MessageStructure<'_>) -> Result<()> {
        let mut signed = false;
        for layer in structure {
            if let MessageLayer::SignatureGroup { results } = layer {
                ensure!(!results.is_empty(), "Missing signature");
                for result in results {
                    result.map_err(|error| {
                        anyhow::anyhow!("Signature verification failed: {error}")
                    })?;
                }
                signed = true;
            }
        }
        ensure!(signed, "Unsigned message rejected");
        Ok(())
    }
}

impl DecryptionHelper for Helper<'_> {
    fn decrypt(
        &mut self,
        pkesks: &[pgp::packet::PKESK],
        _skesks: &[pgp::packet::SKESK],
        algorithm: Option<SymmetricAlgorithm>,
        decrypt: &mut dyn FnMut(Option<SymmetricAlgorithm>, &SessionKey) -> bool,
    ) -> Result<Option<Cert>> {
        let policy = StandardPolicy::new();
        let key = self
            .recipient
            .keys()
            .secret()
            .with_policy(&policy, None)
            .for_transport_encryption()
            .next()
            .context("Missing encryption subkey")?
            .key()
            .clone();
        let mut pair = key
            .decrypt_secret(&Password::from(self.password))?
            .into_keypair()?;
        for packet in pkesks {
            if packet
                .decrypt(&mut pair, algorithm)
                .is_some_and(|(cipher, key)| decrypt(cipher, &key))
            {
                return Ok(Some(self.recipient.clone().strip_secret_key_material()));
            }
        }
        anyhow::bail!("No decryptable recipient packet")
    }
}

fn decrypt(ciphertext: &[u8], recipient: &Cert, sender: &Cert, password: &str) -> Result<Vec<u8>> {
    let policy = StandardPolicy::new();
    let helper = Helper {
        recipient,
        sender,
        password,
    };
    let mut reader =
        DecryptorBuilder::from_bytes(ciphertext)?.with_policy(&policy, None, helper)?;
    let mut plaintext = Vec::new();
    std::io::copy(&mut reader, &mut plaintext)?;
    Ok(plaintext)
}

fn probe() -> Result<String> {
    let backend = pgp::crypto::backend();
    ensure!(
        backend.starts_with("OpenSSL "),
        "Expected OpenSSL backend, got {backend}"
    );
    let sender = identity()?;
    let visitor = identity()?;
    ensure!(
        sender.primary_key().key().pk_algo() == pgp::types::PublicKeyAlgorithm::MLDSA87_Ed448,
        "Unexpected signing algorithm"
    );
    ensure!(
        sender.primary_key().key().version() == 6,
        "Expected v6 identity"
    );
    ensure!(
        sender
            .keys()
            .any(|key| key.key().pk_algo() == pgp::types::PublicKeyAlgorithm::MLKEM1024_X448),
        "Missing ML-KEM-1024+X448"
    );
    let ciphertext = encrypt(&sender, &visitor)?;
    ensure!(
        decrypt(&ciphertext, &sender, &sender, PASSWORD)? == TEXT,
        "Sender round-trip mismatch"
    );
    ensure!(
        decrypt(&ciphertext, &visitor, &sender, PASSWORD)? == TEXT,
        "Visitor round-trip mismatch"
    );
    ensure!(
        decrypt(&ciphertext, &visitor, &sender, "wrong-password").is_err(),
        "Wrong password accepted"
    );
    ensure!(
        decrypt(&ciphertext, &visitor, &visitor, PASSWORD).is_err(),
        "Wrong signer accepted"
    );
    let mut altered = ciphertext;
    let byte = altered.last_mut().context("Empty ciphertext")?;
    *byte ^= 1;
    ensure!(
        decrypt(&altered, &visitor, &sender, PASSWORD).is_err(),
        "Tampered ciphertext accepted"
    );
    Ok(format!(
        "PASS: backend={backend}; RFC 9980 ML-KEM-1024+X448, ML-DSA-87+Ed448, v6 keys, protected key import, both recipients, wrong password, wrong signer, tampering"
    ))
}

#[cfg(test)]
mod tests {
    #[test]
    fn required_crypto_round_trip() -> anyhow::Result<()> {
        super::probe()?;
        Ok(())
    }
}