trazaeo 0.5.2

use crate::envelope::{verify_attestation, PublishEnvelope};
use crate::error::TrazaeoResult;
use crate::trust::{TrustPolicy, TrustStatus};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum VerificationMode {
    Sampled,
    Full,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum VerificationState {
    Passed,
    Failed,
    NotEvaluated,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum VerifiabilityTier {
    Core,
    CoreWithStorageBinding,
    CoreWithProofLogCommitment,
    CoreWithStorageBindingAndProofLogCommitment,
    Failed,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct VerificationReport {
    pub mode: VerificationMode,
    pub signature_state: VerificationState,
    pub trust_state: VerificationState,
    pub binding_state: VerificationState,
    pub reproducibility_state: VerificationState,
    pub lineage_state: VerificationState,
    pub storage_state: VerificationState,
    pub proof_log_state: VerificationState,
    pub reasons: Vec<String>,
}

/// Verifies publish envelope.
pub fn verify_publish_envelope(
    envelope: &PublishEnvelope,
    mode: VerificationMode,
    trust_policy: &TrustPolicy,
) -> VerificationReport {
    let mut reasons = Vec::new();

    let payload = envelope.canonical_attestation_payload_bytes();
    let signature_state = if envelope.attestations.is_empty() {
        reasons.push("missing attestations".to_string());
        VerificationState::Failed
    } else if envelope
        .attestations
        .iter()
        .all(|att| verify_attestation(att, &payload))
    {
        VerificationState::Passed
    } else {
        reasons.push("invalid attestation signature".to_string());
        VerificationState::Failed
    };

    let trust_state = match trust_policy.trust_status(&envelope.key_id) {
        TrustStatus::Trusted => VerificationState::Passed,
        TrustStatus::Revoked => {
            reasons.push("key is revoked".to_string());
            VerificationState::Failed
        }
        TrustStatus::Untrusted => {
            reasons.push("key is untrusted".to_string());
            VerificationState::Failed
        }
    };

    let binding_state = if envelope.input_refs.is_empty() || envelope.output_refs.is_empty() {
        reasons.push("missing input/output refs".to_string());
        VerificationState::Failed
    } else {
        VerificationState::Passed
    };

    let reproducibility_state = if mode == VerificationMode::Full {
        if envelope.verification_policy_id.trim().is_empty() {
            reasons.push("missing verification policy id".to_string());
            VerificationState::Failed
        } else {
            VerificationState::Passed
        }
    } else {
        VerificationState::NotEvaluated
    };

    let lineage_state = if envelope.lineage_refs.is_empty() {
        reasons.push("missing lineage refs".to_string());
        VerificationState::Failed
    } else {
        VerificationState::Passed
    };

    VerificationReport {
        mode,
        signature_state,
        trust_state,
        binding_state,
        reproducibility_state,
        lineage_state,
        storage_state: VerificationState::NotEvaluated,
        proof_log_state: VerificationState::NotEvaluated,
        reasons,
    }
}

/// Returns a report updated with storage binding verification.
pub fn verify_storage_binding(
    report: VerificationReport,
    envelope: &PublishEnvelope,
    stored_uri: &str,
) -> VerificationReport {
    let mut updated = report;
    if envelope.output_refs.iter().any(|value| value == stored_uri) {
        updated.storage_state = VerificationState::Passed;
    } else {
        updated.storage_state = VerificationState::Failed;
        updated
            .reasons
            .push("stored object uri not present in output refs".to_string());
    }
    updated
}

/// Returns a report updated with proof-log verification status.
pub fn apply_proof_log_result(
    report: VerificationReport,
    result: TrazaeoResult<()>,
) -> VerificationReport {
    let mut updated = report;
    match result {
        Ok(()) => updated.proof_log_state = VerificationState::Passed,
        Err(err) => {
            updated.proof_log_state = VerificationState::Failed;
            updated.reasons.push(err.to_string());
        }
    }
    updated
}

/// Handles verifiability tier.
pub fn verifiability_tier(report: &VerificationReport) -> VerifiabilityTier {
    if !is_report_success(report) {
        return VerifiabilityTier::Failed;
    }
    match (report.storage_state, report.proof_log_state) {
        (VerificationState::Passed, VerificationState::Passed) => {
            VerifiabilityTier::CoreWithStorageBindingAndProofLogCommitment
        }
        (VerificationState::Passed, _) => VerifiabilityTier::CoreWithStorageBinding,
        (_, VerificationState::Passed) => VerifiabilityTier::CoreWithProofLogCommitment,
        _ => VerifiabilityTier::Core,
    }
}

/// Checks whether report success.
pub fn is_report_success(report: &VerificationReport) -> bool {
    let signature_ok = report.signature_state == VerificationState::Passed;
    let trust_ok = report.trust_state == VerificationState::Passed;
    let binding_ok = report.binding_state == VerificationState::Passed;
    let repro_ok = report.reproducibility_state == VerificationState::Passed
        || report.reproducibility_state == VerificationState::NotEvaluated;
    let lineage_ok = report.lineage_state == VerificationState::Passed;
    let storage_ok = report.storage_state == VerificationState::Passed
        || report.storage_state == VerificationState::NotEvaluated;
    let proof_log_ok = report.proof_log_state == VerificationState::Passed
        || report.proof_log_state == VerificationState::NotEvaluated;
    signature_ok && trust_ok && binding_ok && repro_ok && lineage_ok && storage_ok && proof_log_ok
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::envelope::{make_attestation, Attestation};
    use crate::error::TrazaeoError;
    use crate::trust::TrustPolicy;

    /// Handles publish.
    fn publish() -> PublishEnvelope {
        let mut p = PublishEnvelope {
            schema_version: "1.0.0".to_string(),
            envelope_type: "publish".to_string(),
            issued_at: "2026-01-01T00:00:00Z".to_string(),
            subject_id: "publish-1".to_string(),
            dataset_id: "sst".to_string(),
            dataset_version: "v1".to_string(),
            input_refs: vec!["obj://in".to_string()],
            output_refs: vec!["obj://out".to_string()],
            published_artifacts: vec![crate::checkpoint::CheckpointArtifact {
                artifact_id: "artifact-1".to_string(),
                content_root_hash: "hash".to_string(),
                content_descriptor_ref: None,
                content_descriptor_hash: None,
                media_type: "application/octet-stream".to_string(),
            }],
            primary_artifact_id: "artifact-1".to_string(),
            checkpoint_manifest_ref: "checkpoint://1".to_string(),
            checkpoint_manifest_hash: "checkpoint-hash".to_string(),
            checkpoint_id: "checkpoint-1".to_string(),
            checkpoint_log_root_hash: "checkpoint-log-root".to_string(),
            lineage_refs: vec!["capture://1".to_string()],
            verification_policy_id: "verify-default".to_string(),
            attestations: vec![],
            key_id: "key-1".to_string(),
            stac_refs: vec![],
            reward_context_ref: None,
            reward_context_hash: None,
            provenance_start_mode: "transport_capture".to_string(),
            bootstrap_origin_label: None,
            reward_eligible: false,
        };
        const TEST_SIGNING_KEY_HEX: &str =
            "4f3edf983ac636a65a842ce7c78d9aa706d3b113bce036f9a4f5762b76f70f18";
        let key_seed = make_attestation("s", TEST_SIGNING_KEY_HEX, "2026-01-01T00:00:00Z", b"")
            .expect("derive key id");
        p.key_id = key_seed.key_id.clone();
        p.attestations = vec![Attestation {
            signer_id: "s".to_string(),
            key_id: key_seed.key_id.clone(),
            signature: String::new(),
            signed_at: "2026-01-01T00:00:00Z".to_string(),
        }];
        let payload = p.canonical_attestation_payload_bytes();
        let attestation =
            make_attestation("s", TEST_SIGNING_KEY_HEX, "2026-01-01T00:00:00Z", &payload)
                .expect("build attestation");
        p.attestations = vec![attestation];
        p
    }

    /// Returns trust information for policy for.
    fn trust_policy_for(key_id: &str) -> TrustPolicy {
        let mut p = TrustPolicy::new();
        p.allow_key(key_id);
        p
    }

    /// Tests that sampled verification success.
    #[test]
    fn sampled_verification_success() {
        let report = verify_publish_envelope(
            &publish(),
            VerificationMode::Sampled,
            &trust_policy_for(&publish().key_id),
        );
        assert!(is_report_success(&report));
        assert_eq!(
            report.reproducibility_state,
            VerificationState::NotEvaluated
        );
        assert_eq!(report.lineage_state, VerificationState::Passed);
    }

    /// Tests that full verification success.
    #[test]
    fn full_verification_success() {
        let report = verify_publish_envelope(
            &publish(),
            VerificationMode::Full,
            &trust_policy_for(&publish().key_id),
        );
        assert!(is_report_success(&report));
        assert_eq!(report.reproducibility_state, VerificationState::Passed);
        assert_eq!(report.lineage_state, VerificationState::Passed);
        assert_eq!(report.storage_state, VerificationState::NotEvaluated);
        assert_eq!(report.proof_log_state, VerificationState::NotEvaluated);
    }

    /// Tests that verification fails without attestation.
    #[test]
    fn verification_fails_without_attestation() {
        let mut p = publish();
        p.attestations.clear();
        let report =
            verify_publish_envelope(&p, VerificationMode::Sampled, &trust_policy_for(&p.key_id));
        assert!(!is_report_success(&report));
        assert_eq!(report.signature_state, VerificationState::Failed);
    }

    /// Tests that verification fails when key untrusted.
    #[test]
    fn verification_fails_when_key_untrusted() {
        let report =
            verify_publish_envelope(&publish(), VerificationMode::Sampled, &TrustPolicy::new());
        assert_eq!(report.trust_state, VerificationState::Failed);
    }

    /// Tests that verification fails when signature tampered.
    #[test]
    fn verification_fails_when_signature_tampered() {
        let mut p = publish();
        p.dataset_version = "v2".to_string();
        let report =
            verify_publish_envelope(&p, VerificationMode::Sampled, &trust_policy_for(&p.key_id));
        assert_eq!(report.signature_state, VerificationState::Failed);
    }

    /// Tests that verification fails without lineage refs.
    #[test]
    fn verification_fails_without_lineage_refs() {
        let mut p = publish();
        p.lineage_refs.clear();
        let report =
            verify_publish_envelope(&p, VerificationMode::Sampled, &trust_policy_for(&p.key_id));
        assert_eq!(report.lineage_state, VerificationState::Failed);
        assert!(!is_report_success(&report));
    }

    /// Tests that storage binding can raise verifiability.
    #[test]
    fn storage_binding_can_raise_verifiability() {
        let report = verify_publish_envelope(
            &publish(),
            VerificationMode::Sampled,
            &trust_policy_for(&publish().key_id),
        );
        let mut bound_publish = publish();
        bound_publish.output_refs = vec!["s3://bucket/output".to_string()];
        let report = verify_storage_binding(report, &bound_publish, "s3://bucket/output");
        assert_eq!(report.storage_state, VerificationState::Passed);
        assert_eq!(
            verifiability_tier(&report),
            VerifiabilityTier::CoreWithStorageBinding
        );
    }

    /// Tests that proof log failure is recorded.
    #[test]
    fn proof_log_failure_is_recorded() {
        let report = verify_publish_envelope(
            &publish(),
            VerificationMode::Sampled,
            &trust_policy_for(&publish().key_id),
        );
        let report = apply_proof_log_result(
            report,
            Err(TrazaeoError::external(
                "apply proof log result",
                "missing proof log receipt",
            )),
        );
        assert_eq!(report.proof_log_state, VerificationState::Failed);
        assert!(!is_report_success(&report));
        assert_eq!(verifiability_tier(&report), VerifiabilityTier::Failed);
    }
}