trazaeo 0.5.4

use crate::content::CONTENT_COMMITMENT_PROFILE_BLAKE3;
use crate::envelope::{Attestation, CaptureEnvelope};
use crate::hashing::{compute_rolling_hash, compute_root_from_hashes};
use crate::utils::Chunk;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CapturedSegment {
    pub segment_id: String,
    pub payload: Vec<u8>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CaptureSessionInput {
    pub schema_version: String,
    pub issued_at: String,
    pub subject_id: String,
    pub capture_role: String,
    pub capture_actor_id: String,
    pub capture_system_id: String,
    pub capture_window: String,
    pub input_refs: Vec<String>,
    pub output_refs: Vec<String>,
    pub content_descriptor_ref: Option<String>,
    pub content_descriptor_hash: Option<String>,
    pub key_id: String,
    pub policy_profile_id: Option<String>,
}

/// Builds capture envelope.
pub fn build_capture_envelope(
    input: &CaptureSessionInput,
    segments: &[CapturedSegment],
    attestation: Attestation,
) -> CaptureEnvelope {
    let segment_ids: Vec<String> = segments.iter().map(|m| m.segment_id.clone()).collect();
    let chunks: Vec<Chunk> = segments
        .iter()
        .map(|m| Chunk {
            data: m.payload.clone(),
        })
        .collect();

    let segment_hashes: Vec<String> = segments
        .iter()
        .map(|m| hex::encode(blake3::hash(&m.payload).as_bytes()))
        .collect();
    let segment_hash_values = segments
        .iter()
        .map(|m| crate::utils::Hash(*blake3::hash(&m.payload).as_bytes()))
        .collect::<Vec<_>>();

    let rolling = compute_rolling_hash(&chunks);
    let rolling_hash_state = hex::encode(rolling.0);
    let content_root_hash = compute_root_from_hashes(&segment_hash_values);
    let chunk_size = segments.iter().map(|m| m.payload.len()).max().unwrap_or(1);

    CaptureEnvelope {
        schema_version: input.schema_version.clone(),
        envelope_type: "capture".to_string(),
        issued_at: input.issued_at.clone(),
        subject_id: input.subject_id.clone(),
        capture_role: input.capture_role.clone(),
        capture_actor_id: input.capture_actor_id.clone(),
        capture_system_id: input.capture_system_id.clone(),
        capture_window: input.capture_window.clone(),
        segment_ids,
        input_refs: input.input_refs.clone(),
        output_refs: input.output_refs.clone(),
        segment_hashes,
        rolling_hash_state: Some(rolling_hash_state),
        content_root_hash: hex::encode(content_root_hash.0),
        content_commitment_profile: CONTENT_COMMITMENT_PROFILE_BLAKE3.to_string(),
        chunk_size,
        leaf_count: segments.len(),
        content_descriptor_ref: input.content_descriptor_ref.clone(),
        content_descriptor_hash: input.content_descriptor_hash.clone(),
        attestations: vec![attestation],
        key_id: input.key_id.clone(),
        policy_profile_id: input.policy_profile_id.clone(),
    }
}

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

    /// Handles attestation.
    fn attestation() -> Attestation {
        Attestation {
            signer_id: "station-signer".to_string(),
            key_id: "key-1".to_string(),
            signature: "sig".to_string(),
            signed_at: "2026-01-01T00:00:00Z".to_string(),
        }
    }

    /// Handles input.
    fn input() -> CaptureSessionInput {
        CaptureSessionInput {
            schema_version: "1.0.0".to_string(),
            issued_at: "2026-01-01T00:00:00Z".to_string(),
            subject_id: "capture-session-1".to_string(),
            capture_role: "transport".to_string(),
            capture_actor_id: "station-a".to_string(),
            capture_system_id: "rx-1".to_string(),
            capture_window: "2026-01-01T00:00:00Z/2026-01-01T00:05:00Z".to_string(),
            input_refs: vec![],
            output_refs: vec!["obj://raw/1".to_string()],
            content_descriptor_ref: None,
            content_descriptor_hash: None,
            key_id: "key-1".to_string(),
            policy_profile_id: Some("policy-1".to_string()),
        }
    }

    /// Tests that build capture envelope emits expected fields.
    #[test]
    fn build_capture_envelope_emits_expected_fields() {
        let messages = vec![
            CapturedSegment {
                segment_id: "seg-1".to_string(),
                payload: b"abc".to_vec(),
            },
            CapturedSegment {
                segment_id: "seg-2".to_string(),
                payload: b"def".to_vec(),
            },
        ];

        let envelope = build_capture_envelope(&input(), &messages, attestation());
        assert_eq!(envelope.segment_ids, vec!["seg-1", "seg-2"]);
        assert_eq!(envelope.segment_hashes.len(), 2);
        assert!(envelope.validate().is_ok());
    }

    /// Tests that build capture envelope is deterministic.
    #[test]
    fn build_capture_envelope_is_deterministic() {
        let messages = vec![CapturedSegment {
            segment_id: "seg-1".to_string(),
            payload: b"abc".to_vec(),
        }];

        let a = build_capture_envelope(&input(), &messages, attestation());
        let b = build_capture_envelope(&input(), &messages, attestation());
        assert_eq!(a.content_root_hash, b.content_root_hash);
        assert_eq!(a.segment_hashes, b.segment_hashes);
    }

    /// Tests that source capture mode is supported for device-side signing.
    #[test]
    fn build_capture_envelope_supports_source_capture() {
        let mut input = input();
        input.capture_role = "source".to_string();
        input.capture_actor_id = "sensor-1".to_string();
        input.capture_system_id = "sensor-pipeline-1".to_string();

        let segments = vec![CapturedSegment {
            segment_id: "frame-1".to_string(),
            payload: b"telemetry".to_vec(),
        }];

        let envelope = build_capture_envelope(&input, &segments, attestation());
        assert_eq!(envelope.capture_role, "source");
        assert_eq!(envelope.capture_actor_id, "sensor-1");
        assert!(envelope.validate().is_ok());
    }
}