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

#[derive(Debug, Serialize, Deserialize)]
pub struct DeviceAttestationInputs {
    pub id: String,
    pub key: String,
    pub attestation_sig: 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, 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-----"#,
];

pub fn verify_device_attestation(inputs: &DeviceAttestationInputs) -> 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_string = format!("{}:{}", inputs.id, inputs.key);
    let payload = payload_string.as_bytes();

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

    if let Some(chain_pem) = &inputs.certificate_chain {
        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 VerificationResult {
                ok: false,
                reason: Some("Invalid certificate chain".to_string()),
            };
        }

        leaf_pubkey_bytes = Some(certs[0].public_key().subject_public_key.data.to_vec());

        let required_ou = match inputs.trust_profile.as_str() {
            "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 VerificationResult {
                ok: false,
                reason: Some(format!(
                    "Certificate chain does not match the requested trust profile: {}",
                    inputs.trust_profile
                )),
            };
        }

        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 VerificationResult {
                    ok: false,
                    reason: Some(format!("Signature verification failed at chain level {}", i)),
                };
            }
        }

        if let Some(created_ts) = inputs.created {
            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 VerificationResult {
                        ok: false,
                        reason: Some(format!(
                            "Temporal birth check failed: created ({}) is outside cert window",
                            created_ts
                        )),
                    };
                }
            }
        }
    }

    if let Some(pubkey) = leaf_pubkey_bytes {
        if verify_signature_robust(&signature_bytes, payload, &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() {
                    if verify_signature_robust(
                        &signature_bytes,
                        payload,
                        &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_external_identity(inputs: &ExternalIdentityInputs) -> 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);
    }

    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);
            return vk.verify(msg, &s).is_ok();
        }
    }

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

    false
}
