// SPDX-License-Identifier: Apache-2.0
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use ed25519_dalek::{self as ed25519, Verifier};
use ml_dsa::MlDsa65;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use x509_parser::prelude::*;

const DAC_VALIDITY_ANCHOR_GRACE_SECONDS: i64 = 3600;

enum CertificateValidityAnchor {
    Created,
    DacNotBeforePlusGrace,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct DeviceAttestationInputs {
    pub id: String,
    pub key: String,
    pub attestation_sig: String,
    pub ctr: Option<u64>,
    pub vendor: Option<String>,
    pub record_id: Option<String>,
    pub certificate_chain: Option<String>,
    pub created: Option<i64>,
    pub attestation_alg: Option<String>,
    pub attestation_payload_version: Option<u32>,
    pub trust_profile: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct VerificationResult {
    pub ok: bool,
    pub reason: Option<String>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct CertificateChainValidationResult {
    pub ok: bool,
    pub reason: Option<String>,
    pub cert_public_keys: Option<Vec<Vec<u8>>>,
    pub leaf_vendor: Option<String>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ExternalIdentityInputs {
    pub endorser_id: String,
    pub root_fingerprint: String,
    pub cert_chain_der: Vec<String>,
    pub signature: String,
    pub expected_payload: String,
    pub trusted_fingerprints: Vec<String>,
}

const TRUSTED_ROOT_CERTS_PEM: &[&str] = &[r#"-----BEGIN CERTIFICATE-----

DxQWGRk=
-----END CERTIFICATE-----"#];

use crate::revocation::RevocationManager;

#[derive(Debug, Serialize, Deserialize)]
pub struct HeartbeatAttestationInputs {
    pub id: String,
    pub heartbeat_sig: String,
    pub last_sync_utc: i64,
    pub ctr: Option<u64>,
    pub record_id: Option<String>,
    pub certificate_chain: Option<String>,
    pub trust_profile: String,
}

pub fn build_record_attestation_payload(id: &str, key: &str, ctr: Option<u64>, vendor: Option<&str>, record_id: Option<&str>) -> String {
    match (ctr, vendor, record_id) {
        (Some(ctr), Some(vendor), Some(record_id)) if !record_id.is_empty() => {
            format!("attestation:{}:{}:{}:{}:{}", id, key, ctr, vendor, record_id)
        }
        _ => format!("{}:{}", id, key),
    }
}

pub fn build_record_heartbeat_payload(
    id: &str,
    last_sync_utc: i64,
    ctr: Option<u64>,
    record_id: Option<&str>,
) -> String {
    match (ctr, record_id) {
        (Some(ctr), Some(record_id)) if !record_id.is_empty() => {
            format!("heartbeat:{}:{}:{}:{}", id, last_sync_utc, ctr, record_id)
        }
        _ => format!("heartbeat:{}:{}", id, last_sync_utc),
    }
}

pub fn verify_device_attestation(
    inputs: &DeviceAttestationInputs,
    revocation_manager: Option<&RevocationManager>,
) -> VerificationResult {
    let signature_bytes = match BASE64.decode(inputs.attestation_sig.trim()) {
        Ok(bytes) => bytes,
        Err(_) => {
            return VerificationResult {
                ok: false,
                reason: Some("attestationSig is not valid base64".to_string()),
            }
        }
    };

    let payload_candidates = [
        build_record_attestation_payload(
            &inputs.id,
            &inputs.key,
            inputs.ctr,
            inputs.vendor.as_deref(),
            inputs.record_id.as_deref(),
        ),
        format!("{}:{}", inputs.id, inputs.key),
    ];

    let mut leaf_pubkey_bytes: Option<Vec<u8>> = None;

    if let Some(chain_pem) = &inputs.certificate_chain {
        let chain_result = validate_certificate_chain_with_anchor(
            chain_pem,
            None,
            &inputs.trust_profile,
            revocation_manager,
            CertificateValidityAnchor::DacNotBeforePlusGrace,
        );
        if !chain_result.ok {
            return VerificationResult {
                ok: false,
                reason: chain_result.reason,
            };
        }
        if let Some(cert_public_keys) = chain_result.cert_public_keys {
            leaf_pubkey_bytes = cert_public_keys.first().cloned();
        }

        if let Some(expected_vendor) = &inputs.vendor {
            if let Some(cert_vendor) = chain_result.leaf_vendor {
                if cert_vendor != *expected_vendor {
                    return VerificationResult {
                        ok: false,
                        reason: Some(format!("Vendor mismatch: expected {}, got {}", expected_vendor, cert_vendor)),
                    };
                }
            } else {
                return VerificationResult {
                    ok: false,
                    reason: Some("Certificate does not contain a vendor organization".to_string()),
                };
            }
        }
    }

    if let Some(pubkey) = leaf_pubkey_bytes {
        for payload in payload_candidates.iter() {
            if verify_signature_robust(&signature_bytes, payload.as_bytes(), &pubkey) {
                return VerificationResult {
                    ok: true,
                    reason: None,
                };
            }
        }
    } else {
        for root_pem in TRUSTED_ROOT_CERTS_PEM {
            if let Some(p) = x509_parser::pem::Pem::iter_from_buffer(root_pem.as_bytes())
                .filter_map(|r| r.ok())
                .next()
            {
                if let Ok(root_cert) = p.parse_x509() {
                    for payload in payload_candidates.iter() {
                        if verify_signature_robust(
                            &signature_bytes,
                            payload.as_bytes(),
                            &root_cert.public_key().subject_public_key.data,
                        ) {
                            return VerificationResult {
                                ok: true,
                                reason: None,
                            };
                        }
                    }
                }
            }
        }
    }

    VerificationResult {
        ok: false,
        reason: Some("Attestation verification failed".to_string()),
    }
}

pub fn verify_heartbeat_attestation(
    inputs: &HeartbeatAttestationInputs,
    revocation_manager: Option<&RevocationManager>,
) -> VerificationResult {
    let signature_bytes = match BASE64.decode(inputs.heartbeat_sig.trim()) {
        Ok(bytes) => bytes,
        Err(_) => {
            return VerificationResult {
                ok: false,
                reason: Some("heartbeat_sig is not valid base64".to_string()),
            }
        }
    };

    let payload_candidates = [
        build_record_heartbeat_payload(
            &inputs.id,
            inputs.last_sync_utc,
            inputs.ctr,
            inputs.record_id.as_deref(),
        ),
        format!("heartbeat:{}:{}", inputs.id, inputs.last_sync_utc),
    ];

    let mut leaf_pubkey_bytes: Option<Vec<u8>> = None;
    let mut authority_pubkey_bytes: Option<Vec<u8>> = None;

    if let Some(chain_pem) = &inputs.certificate_chain {
        // We use a looser temporal check for heartbeats as they are short-lived
        let chain_result = validate_certificate_chain(
            chain_pem,
            None,
            &inputs.trust_profile,
            revocation_manager,
        );
        if !chain_result.ok {
            return VerificationResult {
                ok: false,
                reason: chain_result.reason,
            };
        }
        if let Some(cert_public_keys) = chain_result.cert_public_keys {
            leaf_pubkey_bytes = cert_public_keys.first().cloned();
            authority_pubkey_bytes = cert_public_keys.get(1).cloned();
        }
    }

    if let Some(pubkey) = leaf_pubkey_bytes {
        for payload in payload_candidates.iter() {
            if verify_signature_robust(&signature_bytes, payload.as_bytes(), &pubkey) {
                return VerificationResult {
                    ok: true,
                    reason: None,
                };
            }
        }
    }

    if let Some(pubkey) = authority_pubkey_bytes {
        let legacy_payload = format!("heartbeat:{}:{}", inputs.id, inputs.last_sync_utc);
        if verify_signature_robust(&signature_bytes, legacy_payload.as_bytes(), &pubkey) {
            return VerificationResult {
                ok: true,
                reason: None,
            };
        }
    }

    VerificationResult {
        ok: false,
        reason: Some("Heartbeat verification failed".to_string()),
    }
}

pub fn validate_certificate_chain(
    chain_pem: &str,
    created: Option<i64>,
    trust_profile: &str,
    revocation_manager: Option<&RevocationManager>,
) -> CertificateChainValidationResult {
    validate_certificate_chain_with_anchor(
        chain_pem,
        created,
        trust_profile,
        revocation_manager,
        CertificateValidityAnchor::Created,
    )
}

fn validate_certificate_chain_with_anchor(
    chain_pem: &str,
    created: Option<i64>,
    trust_profile: &str,
    revocation_manager: Option<&RevocationManager>,
    validity_anchor: CertificateValidityAnchor,
) -> CertificateChainValidationResult {
    let pems: Vec<_> = x509_parser::pem::Pem::iter_from_buffer(chain_pem.as_bytes())
        .filter_map(|r| r.ok())
        .collect();

    let certs: Vec<_> = pems.iter().filter_map(|p| p.parse_x509().ok()).collect();

    if certs.is_empty() {
        return CertificateChainValidationResult {
            ok: false,
            reason: Some("Invalid certificate chain".to_string()),
            cert_public_keys: None,
            leaf_vendor: None,
        };
    }

    let required_ou = match trust_profile {
        "dev" => "lukuid-dev",
        "test" => "lukuid-test",
        _ => "lukuid-production",
    };

    let has_valid_profile = certs
        .iter()
        .skip(1)
        .any(|cert| cert.subject().to_string().contains(required_ou));

    if !has_valid_profile {
        return CertificateChainValidationResult {
            ok: false,
            reason: Some(format!(
                "Certificate chain does not match the requested trust profile: {}",
                trust_profile
            )),
            cert_public_keys: None,
            leaf_vendor: None,
        };
    }

    for i in 0..certs.len() {
        let current = &certs[i];
        let mut verified = false;

        if i + 1 < certs.len() {
            let next = &certs[i + 1];
            verified = verify_signature_robust(
                current.signature_value.as_ref(),
                current.tbs_certificate.as_ref(),
                &next.public_key().subject_public_key.data,
            );
        } else {
            for root_pem in TRUSTED_ROOT_CERTS_PEM {
                if let Some(p) = x509_parser::pem::Pem::iter_from_buffer(root_pem.as_bytes())
                    .filter_map(|r| r.ok())
                    .next()
                {
                    if let Ok(root_cert) = p.parse_x509() {
                        if verify_signature_robust(
                            current.signature_value.as_ref(),
                            current.tbs_certificate.as_ref(),
                            &root_cert.public_key().subject_public_key.data,
                        ) {
                            verified = true;
                            break;
                        }
                    }
                }
            }
        }

        if !verified {
                return CertificateChainValidationResult {
                ok: false,
                reason: Some(format!("Signature verification failed at chain level {}", i)),
                cert_public_keys: None,
                leaf_vendor: None,
            };
        }
    }

    if let Some(rm) = revocation_manager {
        for cert in &certs {
            if rm.is_revoked(cert) {
                return CertificateChainValidationResult {
                    ok: false,
                    reason: Some(format!("Certificate is revoked: {}", rm.get_fingerprint(cert))),
                    cert_public_keys: None,
                    leaf_vendor: None,
                };
            }
        }
    }

    let temporal_reference = match validity_anchor {
        CertificateValidityAnchor::Created => created,
        CertificateValidityAnchor::DacNotBeforePlusGrace => certs
            .first()
            .map(|cert| cert.validity().not_before.timestamp() + DAC_VALIDITY_ANCHOR_GRACE_SECONDS),
    };

    if let Some(created_ts) = temporal_reference {
        for cert in &certs {
            let valid_from = cert.validity().not_before.timestamp();
            let valid_to = cert.validity().not_after.timestamp();
            if created_ts < valid_from || created_ts > valid_to {
                return CertificateChainValidationResult {
                    ok: false,
                    reason: Some(format!(
                        "Temporal birth check failed: created ({}) is outside cert window",
                        created_ts
                    )),
                    cert_public_keys: None,
                    leaf_vendor: None,
                };
            }
        }
    }

    let cert_public_keys = certs
        .iter()
        .map(|cert| cert.public_key().subject_public_key.data.to_vec())
        .collect();

    let mut leaf_vendor = None;
    if let Some(leaf) = certs.first() {
        for rdn in leaf.subject().iter() {
            for attr in rdn.iter() {
                if attr.attr_type().to_string() == "2.5.4.10" { // OID for organizationName
                    if let Ok(value) = attr.attr_value().as_str() {
                        leaf_vendor = Some(value.to_string());
                    } else if let Ok(value) = std::str::from_utf8(attr.attr_value().data) {
                        leaf_vendor = Some(value.to_string());
                    }
                }
            }
        }
    }

    CertificateChainValidationResult {
        ok: true,
        reason: None,
        cert_public_keys: Some(cert_public_keys),
        leaf_vendor,
    }
}

pub fn public_key_matches_ed25519_certificate_leaf(
    certificate_public_key: &[u8],
    public_key_base64: &str,
) -> bool {
    match BASE64.decode(public_key_base64.trim()) {
        Ok(decoded) => decoded.len() == 32 && decoded == certificate_public_key,
        Err(_) => false,
    }
}

pub fn verify_payload_signature_with_public_key(
    certificate_public_key: &[u8],
    signature_base64: &str,
    payload: &[u8],
) -> VerificationResult {
    let signature_bytes = match BASE64.decode(signature_base64.trim()) {
        Ok(bytes) => bytes,
        Err(_) => {
            return VerificationResult {
                ok: false,
                reason: Some("signature is not valid base64".to_string()),
            }
        }
    };

    if verify_signature_robust(&signature_bytes, payload, certificate_public_key) {
        VerificationResult {
            ok: true,
            reason: None,
        }
    } else {
        VerificationResult {
            ok: false,
            reason: Some("Signature verification failed against certificate public key".to_string()),
        }
    }
}

pub fn verify_external_identity(
    inputs: &ExternalIdentityInputs,
    revocation_manager: Option<&RevocationManager>,
) -> VerificationResult {
    let signature_bytes = match BASE64.decode(inputs.signature.trim()) {
        Ok(bytes) => bytes,
        Err(_) => {
            return VerificationResult {
                ok: false,
                reason: Some("signature is not valid base64".to_string()),
            }
        }
    };

    if inputs.cert_chain_der.is_empty() {
        return VerificationResult {
            ok: false,
            reason: Some("Certificate chain is empty".to_string()),
        };
    }

    let mut cert_ders = Vec::new();
    for der_b64 in &inputs.cert_chain_der {
        let der = match BASE64.decode(der_b64.trim()) {
            Ok(d) => d,
            Err(_) => {
                return VerificationResult {
                    ok: false,
                    reason: Some("Invalid base64 in cert_chain_der".to_string()),
                }
            }
        };
        cert_ders.push(der);
    }

    // Revocation Check: Check every certificate in the chain
    if let Some(rm) = revocation_manager {
        for der in &cert_ders {
            if let Ok((_, cert)) = X509Certificate::from_der(der) {
                if rm.is_revoked(&cert) {
                    return VerificationResult {
                        ok: false,
                        reason: Some(format!(
                            "External identity certificate is revoked: {}",
                            rm.get_fingerprint(&cert)
                        )),
                    };
                }
            }
        }
    }

    let actual_root_fingerprint = {
        let mut hasher = Sha256::new();
        hasher.update(&cert_ders[cert_ders.len() - 1]);
        hasher
            .finalize()
            .iter()
            .map(|b| format!("{:02x}", b))
            .collect::<String>()
    };

    if actual_root_fingerprint != inputs.root_fingerprint.to_lowercase() {
        return VerificationResult {
            ok: false,
            reason: Some(format!(
                "Root fingerprint mismatch: expected {}, got {}",
                inputs.root_fingerprint, actual_root_fingerprint
            )),
        };
    }

    let mut is_trusted = inputs
        .trusted_fingerprints
        .iter()
        .any(|f| f.to_lowercase() == actual_root_fingerprint);

    if !is_trusted {
        for root_pem in TRUSTED_ROOT_CERTS_PEM {
            let pems: Vec<_> = x509_parser::pem::Pem::iter_from_buffer(root_pem.as_bytes())
                .filter_map(|r| r.ok())
                .collect();
            if let Some(p) = pems.first() {
                let mut hasher = Sha256::new();
                hasher.update(&p.contents);
                let root_f = hasher
                    .finalize()
                    .iter()
                    .map(|b| format!("{:02x}", b))
                    .collect::<String>();
                if root_f == actual_root_fingerprint {
                    if let Ok(leaf) = X509Certificate::from_der(&cert_ders[0]) {
                        if leaf
                            .1
                            .extensions()
                            .iter()
                            .any(|ext| ext.oid.to_string() == "1.3.6.1.4.1.65432.1.4")
                        {
                            is_trusted = true;
                            break;
                        }
                    }
                }
            }
        }
    }

    if !is_trusted {
        return VerificationResult {
            ok: false,
            reason: Some(format!(
                "Root fingerprint {} is not in the trusted list",
                actual_root_fingerprint
            )),
        };
    }

    let leaf_spki = match X509Certificate::from_der(&cert_ders[0]) {
        Ok(c) => c.1.public_key().subject_public_key.data.to_vec(),
        Err(_) => {
            return VerificationResult {
                ok: false,
                reason: Some("Failed to parse leaf certificate".to_string()),
            }
        }
    };

    if verify_signature_robust(
        &signature_bytes,
        inputs.expected_payload.as_bytes(),
        &leaf_spki,
    ) {
        VerificationResult {
            ok: true,
            reason: None,
        }
    } else {
        VerificationResult {
            ok: false,
            reason: Some("External identity signature verification failed".to_string()),
        }
    }
}

fn verify_signature_robust(sig: &[u8], msg: &[u8], pubkey_spki: &[u8]) -> bool {
    // 1. Try ML-DSA-65
    if sig.len() == 3309 {
        let pub_key_bytes = if pubkey_spki.len() > 1952 {
            &pubkey_spki[pubkey_spki.len() - 1952..]
        } else {
            pubkey_spki
        };
        if pub_key_bytes.len() == 1952 {
            // Using ml_dsa crate
            let vk = ml_dsa::VerifyingKey::<MlDsa65>::decode(pub_key_bytes.try_into().unwrap());
            if let Some(s) = ml_dsa::Signature::<MlDsa65>::decode(sig.try_into().unwrap()) {
                // ml-dsa crate 0.1.0-rc.8 implements signature::Verifier
                use ml_dsa::signature::Verifier;
                return vk.verify(msg, &s).is_ok();
            }
        }
    }

    // 2. Try Ed25519
    let ed_pub_bytes = if pubkey_spki.len() == 33 && pubkey_spki[0] == 0 {
        &pubkey_spki[1..33]
    } else if pubkey_spki.len() > 32 {
        &pubkey_spki[pubkey_spki.len() - 32..]
    } else {
        pubkey_spki
    };
    if sig.len() == 64 && ed_pub_bytes.len() == 32 {
        if let Ok(vk) = ed25519::VerifyingKey::from_bytes(ed_pub_bytes.try_into().unwrap()) {
            let mut sig_arr = [0u8; 64];
            sig_arr.copy_from_slice(sig);
            let s = ed25519::Signature::from_bytes(&sig_arr);
            if vk.verify(msg, &s).is_ok() {
                return true;
            }
        }
    }

    // 3. Try ECDSA (P-256) via ring
    if pubkey_spki.len() >= 64 {
        use ring::signature;
        if sig.len() == 64 {
            let peer_public_key =
                signature::UnparsedPublicKey::new(&signature::ECDSA_P256_SHA256_FIXED, pubkey_spki);
            if peer_public_key.verify(msg, sig).is_ok() {
                return true;
            }
        }
        let peer_public_key =
            signature::UnparsedPublicKey::new(&signature::ECDSA_P256_SHA256_ASN1, pubkey_spki);
        if peer_public_key.verify(msg, sig).is_ok() {
            return true;
        }
    }

    false
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_verify_signature_robust_p256_fixed() {
        use ring::signature::{self, EcdsaKeyPair, KeyPair};
        let rng = ring::rand::SystemRandom::new();
        let pkcs8_bytes = signature::EcdsaKeyPair::generate_pkcs8(
            &signature::ECDSA_P256_SHA256_FIXED_SIGNING,
            &rng,
        ).unwrap();

        let key_pair = EcdsaKeyPair::from_pkcs8(
            &signature::ECDSA_P256_SHA256_FIXED_SIGNING,
            pkcs8_bytes.as_ref(),
            &rng,
        ).unwrap();

        let msg = b"hello world";
        let sig = key_pair.sign(&rng, msg).unwrap();
        let pubkey = key_pair.public_key().as_ref();

        assert_eq!(sig.as_ref().len(), 64);
        assert!(pubkey.len() >= 64);

        assert!(verify_signature_robust(sig.as_ref(), msg, pubkey));
    }
}
