//! `HandoffEnvelope`: signed state-transfer payload between agent runs.
//!
//! Spec § 2.6: the signing payload is canonical JSON over the fields
//! `from_run`, `merkle_root`, `state_hash`, `source_identity`, `tenant`,
//! `issued_at`, and `expires_at`. Receivers MUST call `verify()` before
//! adopting any envelope — the type system enforces this via `HandoffProof`.

use crate::types::{AgentId, TenantId};
use ed25519_dalek::{Signature, Verifier, VerifyingKey};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

/// Merkle-prefix proof anchoring a handoff envelope to a specific point in
/// the source agent's audit chain.
///
/// Phase B (v0.3) scope: carries the hex-encoded Merkle root; verifiers
/// recompute against the local provenance store where reconstruction is
/// available. Multi-deployment proof transfer is out of v0.3 — the receiver
/// must trust that the source run's chain can be inspected when needed.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub struct MerklePrefixProof {
    /// Hex-encoded Merkle root at the point of handoff.
    pub root: String,
    /// Optional sequence-number anchor.
    pub at_seq: Option<u64>,
}

/// Source-agent Ed25519 signature over the canonical-JSON payload.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub struct EnvelopeSignature {
    /// Hex-encoded 64-byte Ed25519 signature.
    pub signature_hex: String,
    /// Hex-encoded 32-byte verifying key. Receivers MUST resolve this against
    /// a trusted source-identity registry before relying on the signature.
    pub verifying_key_hex: String,
}

/// Signed state-transfer envelope. Receivers must call `Handoff::verify`
/// before adopting; the returned `HandoffProof` is the only path to the
/// inner state.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub struct HandoffEnvelope {
    /// Source `RunId`.
    pub from_run: String,
    /// Merkle prefix proof anchoring the envelope to the source run.
    pub merkle_prefix_proof: MerklePrefixProof,
    /// CBOR-serialised, redacted state snapshot.
    pub redacted_state: Vec<u8>,
    /// Source-agent signature over the canonical payload.
    pub source_signature: EnvelopeSignature,
    /// Source-agent identity.
    pub source_identity: AgentId,
    /// Tenant tag at packaging time.
    pub tenant: Option<TenantId>,
    /// Issued-at timestamp (RFC 3339).
    pub issued_at: String,
    /// Expiry timestamp (RFC 3339).
    pub expires_at: String,
}

/// Proof returned by `Handoff::verify` once an envelope is accepted.
///
/// `#[non_exhaustive]` prevents external crates from constructing this via
/// struct literal, enforcing that the only path to a `HandoffProof` is
/// through `Handoff::verify`. Internal construction (within `klieo-ops`)
/// uses the standard struct literal.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct HandoffProof {
    /// Source run id.
    pub from_run: String,
    /// Source agent identity.
    pub source_identity: AgentId,
    /// Tenant tag.
    pub tenant: Option<TenantId>,
    /// Decoded redacted state (raw bytes); receiver decodes CBOR.
    pub redacted_state: Vec<u8>,
}

/// Canonical JSON payload used for signing.
///
/// Shape is pinned to spec § 2.6 — adding fields is a SemVer-major break
/// for the handoff envelope wire format.
///
/// Note on canonicalisation: `klieo-provenance` does not currently expose
/// a standalone `canonical_json` function (it uses an internal custom
/// serialiser in `EvidenceBundle::canonical_payload`). We therefore
/// canonicalise here by serialising a `BTreeMap` whose keys are sorted
/// lexicographically — `serde_json` preserves insertion order for maps
/// built from iterators, but `BTreeMap` guarantees sorted iteration.
/// This matches the spec contract for this field set.  A follow-up task
/// should expose `klieo_provenance::canonical_json` publicly (SemVer-minor
/// addition) and wire it here.
#[derive(Debug)]
pub(crate) struct CanonicalPayload {
    pub from_run: String,
    pub merkle_root: String,
    /// Hex of blake3(redacted_state).
    pub state_hash: String,
    pub source_identity: String,
    pub tenant: Option<String>,
    pub issued_at: String,
    pub expires_at: String,
}

impl CanonicalPayload {
    pub(crate) fn from_envelope(env: &HandoffEnvelope) -> Self {
        let state_hash = blake3::hash(&env.redacted_state).to_hex().to_string();
        Self {
            from_run: env.from_run.clone(),
            merkle_root: env.merkle_prefix_proof.root.clone(),
            state_hash,
            source_identity: env.source_identity.0.clone(),
            tenant: env.tenant.as_ref().map(|t| t.0.clone()),
            issued_at: env.issued_at.clone(),
            expires_at: env.expires_at.clone(),
        }
    }

    /// Serialise to canonical JSON bytes via a sorted `BTreeMap`.
    pub(crate) fn to_bytes(&self) -> Vec<u8> {
        let mut map: BTreeMap<&str, serde_json::Value> = BTreeMap::new();
        map.insert("expires_at", self.expires_at.clone().into());
        map.insert("from_run", self.from_run.clone().into());
        map.insert("issued_at", self.issued_at.clone().into());
        map.insert("merkle_root", self.merkle_root.clone().into());
        map.insert("source_identity", self.source_identity.clone().into());
        map.insert("state_hash", self.state_hash.clone().into());
        map.insert(
            "tenant",
            match &self.tenant {
                Some(t) => t.clone().into(),
                None => serde_json::Value::Null,
            },
        );
        serde_json::to_vec(&map).expect("BTreeMap<&str, Value> is always JSON-serialisable")
    }
}

/// Verify an envelope's signature against the embedded verifying key.
///
/// Receivers MUST additionally confirm that the verifying key matches a
/// trusted source-agent identity via their own registry. This function
/// verifies only the cryptographic math.
pub(crate) fn verify_signature(env: &HandoffEnvelope) -> Result<(), String> {
    let sig_bytes = hex::decode(&env.source_signature.signature_hex)
        .map_err(|e| format!("signature_hex decode: {e}"))?;
    let sig = Signature::from_slice(&sig_bytes).map_err(|e| format!("signature parse: {e}"))?;
    let vk_bytes = hex::decode(&env.source_signature.verifying_key_hex)
        .map_err(|e| format!("verifying_key_hex decode: {e}"))?;
    let vk_array: [u8; 32] = vk_bytes
        .try_into()
        .map_err(|_| "verifying key must be exactly 32 bytes".to_string())?;
    let vk =
        VerifyingKey::from_bytes(&vk_array).map_err(|e| format!("verifying key parse: {e}"))?;
    let canonical_bytes = CanonicalPayload::from_envelope(env).to_bytes();
    vk.verify(&canonical_bytes, &sig)
        .map_err(|e| format!("signature math invalid: {e}"))?;
    Ok(())
}
