vex-chora 1.7.0

use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use vex_core::segment::AuthorityData;

/// Response from the CHORA Authority.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChoraResponse {
    pub authority: AuthorityData,
    pub signature: String,
    pub intent_hash: Option<String>,
    pub authority_class: Option<String>,
}

/// Trait for external authority clients.
/// This ensures VEX remains neutral and can support multiple witness providers.
#[async_trait]
pub trait AuthorityClient: Send + Sync + std::fmt::Debug {
    async fn request_attestation(
        &self,
        payload: &[u8],
        nonce: &str,
    ) -> Result<ChoraResponse, String>;
    async fn verify_witness_signature(
        &self,
        payload: &[u8],
        signature: &[u8],
    ) -> Result<bool, String>;
    async fn verify_continuation_token(
        &self,
        token: &vex_core::ContinuationToken,
        expected_aid: Option<&str>,
        expected_intent_hash: Option<&str>,
        expected_circuit_id: Option<&str>,
        expected_nonce: Option<&str>,
        expected_source_capsule_root: Option<&str>,
    ) -> Result<bool, String>;

    /// Request a resolution for an escalated action (The EMS Path).
    /// This is used when a gate returns ESCALATE, or when a manual override is needed.
    async fn request_escalation(
        &self,
        escalation_id: &str,
        _context: &serde_json::Value,
    ) -> Result<ChoraResponse, String>;
}

/// A Mock Authority Client for test/dev environments.
/// Generates deterministic signatures based on a test key.
#[derive(Debug)]
pub struct MockChoraClient;

#[async_trait]
impl AuthorityClient for MockChoraClient {
    async fn request_attestation(
        &self,
        payload: &[u8],
        _nonce: &str,
    ) -> Result<ChoraResponse, String> {
        use ed25519_dalek::{Signer, SigningKey};
        use sha2::{Digest, Sha256};

        // SHA-256 for witness_receipt
        let mut hasher = Sha256::new();
        hasher.update(payload);
        let hash = hasher.finalize();
        let _witness_receipt = hex::encode(hash);

        let authority = AuthorityData {
            capsule_id: "chora-mock-id".into(),
            outcome: "ALLOW".into(),
            reason_code: "OK".into(),
            nonce: "42".into(),
            trace_root: "00".repeat(32), // Mocked trace root
            escalation_id: None,
            binding_status: Some("SHADOW".to_string()),
            continuation_token: Some(vex_core::ContinuationToken {
                payload: vex_core::ContinuationPayload {
                    schema: "chora.continuation.token.v3.1".to_string(),
                    ledger_event_id: "mock-ledger-id".to_string(),
                    source_capsule_root: "mock-root".to_string(),
                    resolution_event_id: Some("mock-resolve-id".to_string()),
                    capabilities: vec![],
                    nonce: "mock-nonce".to_string(),
                    execution_target: vex_core::segment::ExecutionTarget {
                        aid: "mock-aid-01".to_string(),
                        circuit_id: "mock-circuit-id".to_string(),
                        intent_hash: hex::encode(Sha256::digest(payload)),
                    },
                    iat: 1711630000,
                    exp: 1711977599, // Way in the future
                    issuer: "chora-gate-mock".to_string(),
                },
                signature: "mock-sig".to_string(),
            }),
            authority_class: None,
            gate_sensors: vex_core::segment::SchemaValue(serde_json::Value::Null),
            metadata: vex_core::segment::SchemaValue(serde_json::Value::Null),
        };

        // Generate mock signature
        let signing_key = SigningKey::from_bytes(&[0u8; 32]);
        let sig = signing_key.sign(payload);
        let signature = hex::encode(sig.to_bytes());

        Ok(ChoraResponse {
            authority,
            signature,
            intent_hash: Some(hex::encode(Sha256::digest(payload))),
            authority_class: Some("ALLOW_PATH".to_string()),
        })
    }

    async fn verify_witness_signature(
        &self,
        payload: &[u8],
        signature: &[u8],
    ) -> Result<bool, String> {
        use ed25519_dalek::{Signature, Verifier, VerifyingKey};
        let verifying_key = VerifyingKey::from_bytes(&[0u8; 32]).map_err(|e| e.to_string())?;
        let sig = Signature::from_bytes(signature.try_into().map_err(|_| "Invalid sig length")?);
        Ok(verifying_key.verify(payload, &sig).is_ok())
    }

    async fn verify_continuation_token(
        &self,
        _token: &vex_core::ContinuationToken,
        _expected_aid: Option<&str>,
        _expected_intent_hash: Option<&str>,
        expected_circuit_id: Option<&str>,
        _expected_nonce: Option<&str>,
        _expected_source_capsule_root: Option<&str>,
    ) -> Result<bool, String> {
        // Mock always returns true for test/dev, but logs circuit_id binding
        if let Some(cid) = expected_circuit_id {
            tracing::debug!("Mock: Verifying token with circuit_id binding: {}", cid);
        }
        Ok(true)
    }

    async fn request_escalation(
        &self,
        escalation_id: &str,
        _context: &serde_json::Value,
    ) -> Result<ChoraResponse, String> {
        tracing::info!(
            "Mock: Escalation {} resolved via EMS manual path.",
            escalation_id
        );
        // Return a mock ALLOW with a fresh token
        self.request_attestation(b"escalated-payload", "mock-ems-nonce")
            .await
    }
}

/// A real HTTP Authority Client for production environments.
/// Connects to any CHORA-compatible gate node via HTTP.
/// Configure with CHORA_GATE_URL and CHORA_API_KEY environment variables.
#[derive(Debug)]
pub struct HttpChoraClient {
    client: reqwest::Client,
    base_url: String,
    api_key: String,
}

impl HttpChoraClient {
    pub fn new(base_url: String, api_key: String) -> Self {
        Self {
            client: reqwest::Client::builder()
                .timeout(std::time::Duration::from_secs(30))
                .build()
                .unwrap_or_default(),
            base_url,
            api_key,
        }
    }

    fn gate_url(&self) -> String {
        let base = self.base_url.trim_end_matches('/');
        format!("{}/gate", base)
    }

    fn public_key_url(&self) -> String {
        let base = self.base_url.trim_end_matches('/');
        format!("{}/public_key", base)
    }
}

#[derive(Debug, Deserialize)]
struct ChoraApiAuthority {
    capsule_id: String,
    outcome: String,
    reason_code: String,
    #[serde(default)]
    nonce: Option<String>,
    #[serde(default)]
    trace_root: Option<String>,
    /// New Phase 2 fields
    #[serde(default)]
    escalation_id: Option<String>,
    #[serde(default)]
    binding_status: Option<String>,
    #[serde(default, alias = "signed_token")]
    pub continuation_token: Option<vex_core::ContinuationToken>,
    #[serde(default)]
    pub authority_class: Option<String>,
    #[serde(default)]
    gate_sensors: vex_core::segment::SchemaValue,
    #[serde(default)]
    metadata: vex_core::segment::SchemaValue,
}

#[derive(Debug, Deserialize)]
struct ChoraApiResponse {
    #[serde(alias = "signed_payload")]
    authority: Option<ChoraApiAuthority>,
    #[serde(default)]
    capsule_id: Option<String>,
    #[serde(default)]
    outcome: Option<String>,
    #[serde(default)]
    reason_code: Option<String>,
    #[serde(default)]
    signature: Option<String>,
    #[serde(default, alias = "receipt_hash")]
    witness_receipt: Option<String>,
    #[serde(default, rename = "capsule_root")]
    _capsule_root: Option<String>,
    #[serde(default)]
    intent_hash: Option<String>,
    /// New Phase 2 fields (flat)
    #[serde(default)]
    escalation_id: Option<String>,
    #[serde(default)]
    binding_status: Option<String>,
    #[serde(default, alias = "signed_token")]
    pub continuation_token: Option<vex_core::ContinuationToken>,
    #[serde(default)]
    pub authority_class: Option<String>,
}

#[async_trait]
impl AuthorityClient for HttpChoraClient {
    async fn request_attestation(
        &self,
        payload: &[u8],
        _nonce: &str,
    ) -> Result<ChoraResponse, String> {
        use sha2::{Digest, Sha256};

        // Derive a SHA-256 based capsule_id from the payload
        let mut hasher = Sha256::new();
        hasher.update(payload);
        let hash = hasher.finalize();
        let payload_hash = hex::encode(hash);

        let intent = serde_json::json!({
            "schema": "vex.intent.v1",
            "data": if payload.starts_with(b"{") {
                serde_json::from_slice::<serde_json::Value>(payload).unwrap_or_else(|_| serde_json::json!({ "raw_payload": payload_hash }))
            } else {
                serde_json::json!({
                    "type": "Transparent",
                    "payload_hash": payload_hash,
                    "payload_b64": base64::Engine::encode(&base64::engine::general_purpose::STANDARD, payload),
                })
            }
        });

        let body = serde_json::json!({
            "confidence": 0.95,
            "intent": intent,
            "bridge_identity": "vex-titan-01"
        });

        let resp = self
            .client
            .post(self.gate_url())
            .header("X-API-Key", &self.api_key)
            .header("Content-Type", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|e| format!("CHORA HTTP request failed: {}", e))?;

        let status = resp.status();
        let text = resp
            .text()
            .await
            .unwrap_or_else(|_| "empty response".to_string());

        if !status.is_success() {
            return Err(format!("CHORA gate returned status {}: {}", status, text));
        }

        let api_resp: ChoraApiResponse = serde_json::from_str(&text)
            .map_err(|e| format!("CHORA response parse failed: {} (raw: {})", e, text))?;

        // Support both response shapes: nested signed_payload (authority) or flat fields
        let capsule_id = api_resp
            .authority
            .as_ref()
            .map(|a| a.capsule_id.clone())
            .or_else(|| api_resp.capsule_id.clone())
            .unwrap_or_else(|| payload_hash.clone());

        let outcome = api_resp
            .authority
            .as_ref()
            .map(|a| a.outcome.clone())
            .or_else(|| api_resp.outcome.clone())
            .unwrap_or_else(|| "ALLOW".to_string());

        let reason_code = api_resp
            .authority
            .as_ref()
            .map(|a| a.reason_code.clone())
            .or_else(|| api_resp.reason_code.clone())
            .unwrap_or_else(|| "OK".to_string());

        let nonce = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.nonce.clone())
            .unwrap_or_else(|| "0".to_string());

        let trace_root = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.trace_root.clone())
            .or_else(|| api_resp.witness_receipt.clone())
            .unwrap_or_else(|| payload_hash.clone());

        let escalation_id = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.escalation_id.clone())
            .or_else(|| api_resp.escalation_id.clone());

        let binding_status = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.binding_status.clone())
            .or_else(|| api_resp.binding_status.clone());

        let continuation_token = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.continuation_token.clone())
            .or_else(|| api_resp.continuation_token.clone());

        let authority_class = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.authority_class.clone())
            .or_else(|| api_resp.authority_class.clone());

        let authority = AuthorityData {
            capsule_id,
            outcome,
            reason_code,
            nonce,
            trace_root,
            escalation_id,
            binding_status,
            continuation_token,
            authority_class: authority_class.clone(),
            gate_sensors: api_resp
                .authority
                .as_ref()
                .map(|a| a.gate_sensors.clone())
                .unwrap_or_default(),
            metadata: api_resp
                .authority
                .as_ref()
                .map(|a| a.metadata.clone())
                .unwrap_or_default(),
        };

        let signature = api_resp
            .signature
            .or(api_resp.witness_receipt)
            .unwrap_or(payload_hash);

        Ok(ChoraResponse {
            authority,
            signature,
            intent_hash: api_resp.intent_hash,
            authority_class,
        })
    }

    async fn verify_witness_signature(
        &self,
        payload: &[u8],
        signature: &[u8],
    ) -> Result<bool, String> {
        use ed25519_dalek::{Signature, Verifier, VerifyingKey};

        // Fetch the CHORA node's Ed25519 public key via GET /public_key
        let resp = self
            .client
            .get(self.public_key_url())
            .header("x-api-key", &self.api_key)
            .send()
            .await
            .map_err(|e| format!("CHORA public_key fetch failed: {}", e))?;

        // Live endpoint can return Hex, Quoted Hex, or PEM
        let raw_text = resp.text().await.map_err(|e| e.to_string())?;
        let key_data = if raw_text.contains("-----BEGIN PUBLIC KEY-----") {
            // Simple PEM extraction
            let base64_part = raw_text
                .lines()
                .filter(|line| !line.starts_with("---"))
                .collect::<String>()
                .replace(" ", "")
                .replace("\n", "")
                .replace("\r", "");
            base64::Engine::decode(&base64::engine::general_purpose::STANDARD, base64_part)
                .map_err(|e| format!("PEM Base64 decode failed: {}", e))?
        } else {
            let hex_key: String = serde_json::from_str(&raw_text)
                .unwrap_or_else(|_| raw_text.trim_matches('"').to_string());
            hex::decode(&hex_key).map_err(|e| format!("Public key hex decode failed: {}", e))?
        };

        // If PEM/PKCS8, Ed25519 raw key is the last 32 bytes
        let raw_key: [u8; 32] = if key_data.len() > 32 {
            key_data[key_data.len() - 32..]
                .try_into()
                .map_err(|_| "Invalid key slice".to_string())?
        } else {
            key_data
                .try_into()
                .map_err(|_| "Invalid Ed25519 public key length".to_string())?
        };

        let verifying_key = VerifyingKey::from_bytes(&raw_key).map_err(|e| e.to_string())?;

        let sig_bytes: [u8; 64] = signature
            .try_into()
            .map_err(|_| "Signature must be 64 bytes".to_string())?;
        let sig = Signature::from_bytes(&sig_bytes);

        Ok(verifying_key.verify(payload, &sig).is_ok())
    }

    async fn verify_continuation_token(
        &self,
        token: &vex_core::ContinuationToken,
        expected_aid: Option<&str>,
        expected_intent_hash: Option<&str>,
        expected_circuit_id: Option<&str>,
        expected_nonce: Option<&str>,
        expected_source_capsule_root: Option<&str>,
    ) -> Result<bool, String> {
        // 1. Context Binding (Silicon Seal)
        if let Some(target) = expected_aid {
            if token.payload.execution_target.aid != target {
                return Err(format!(
                    "Context Mismatch: token.execution_target.aid ({}) != current target ({})",
                    token.payload.execution_target.aid, target
                ));
            }
        }

        if let Some(ih) = expected_intent_hash {
            if token.payload.execution_target.intent_hash != ih {
                return Err(format!(
                    "Context Mismatch: token.execution_target.intent_hash ({}) != current intent_hash ({})",
                    token.payload.execution_target.intent_hash, ih
                ));
            }
        }

        if let Some(cid) = expected_circuit_id {
            if token.payload.execution_target.circuit_id != cid {
                return Err(format!(
                    "Circuit Mismatch: token.execution_target.circuit_id ({}) != expected ({})",
                    token.payload.execution_target.circuit_id, cid
                ));
            }
        }

        if let Some(enc) = expected_nonce {
            if token.payload.nonce != enc {
                return Err(format!(
                    "Nonce Mismatch: token.nonce ({}) != expected ({})",
                    token.payload.nonce, enc
                ));
            }
        }

        if let Some(escr) = expected_source_capsule_root {
            if token.payload.source_capsule_root != escr {
                return Err(format!(
                    "Custody Mismatch: token.source_capsule_root ({}) != expected ({})",
                    token.payload.source_capsule_root, escr
                ));
            }
        }

        // 2. Lifecycle Validation
        token.payload.validate_lifecycle(chrono::Utc::now())?;

        // 3. Cryptographic Verification (Silicon Signature)
        let resp = self
            .client
            .get(self.public_key_url())
            .header("x-api-key", &self.api_key)
            .send()
            .await
            .map_err(|e| format!("CHORA public_key fetch failed: {}", e))?;

        let hex_key: String = resp.json().await.unwrap_or_else(|_| "00".repeat(32)); // Default to fail

        token.verify_v3(&hex_key)
    }

    async fn request_escalation(
        &self,
        escalation_id: &str,
        context: &serde_json::Value,
    ) -> Result<ChoraResponse, String> {
        let ems_url = format!(
            "{}/ems/{}",
            self.base_url.trim_end_matches('/'),
            escalation_id
        );

        let resp = self
            .client
            .post(&ems_url)
            .header("x-api-key", &self.api_key)
            .json(context)
            .send()
            .await
            .map_err(|e| format!("CHORA EMS request failed: {}", e))?;

        if !resp.status().is_success() {
            return Err(format!("CHORA EMS returned error: {}", resp.status()));
        }

        let api_resp: ChoraApiResponse = resp
            .json()
            .await
            .map_err(|e| format!("CHORA EMS response parse failed: {}", e))?;

        let authority_class = api_resp
            .authority
            .as_ref()
            .and_then(|a| a.authority_class.clone())
            .or_else(|| api_resp.authority_class.clone());

        // Re-use logic from request_attestation (Simplified here)
        let authority = AuthorityData {
            capsule_id: api_resp.capsule_id.unwrap_or_default(),
            outcome: api_resp.outcome.unwrap_or_else(|| "ALLOW".to_string()),
            reason_code: api_resp.reason_code.unwrap_or_else(|| "OK".to_string()),
            nonce: "0".to_string(),
            trace_root: api_resp.witness_receipt.clone().unwrap_or_default(),
            escalation_id: api_resp.escalation_id,
            binding_status: api_resp.binding_status,
            continuation_token: api_resp.continuation_token,
            authority_class: authority_class.clone(),
            gate_sensors: vex_core::segment::SchemaValue(serde_json::Value::Null),
            metadata: vex_core::segment::SchemaValue(serde_json::Value::Null),
        };

        Ok(ChoraResponse {
            authority,
            signature: api_resp.signature.unwrap_or_default(),
            intent_hash: api_resp.intent_hash,
            authority_class,
        })
    }
}

/// Factory: creates a real HttpChoraClient.
/// Used by vex-server to avoid importing the concrete type directly.
pub fn make_authority_client(url: String, api_key: String) -> std::sync::Arc<dyn AuthorityClient> {
    std::sync::Arc::new(HttpChoraClient::new(url, api_key))
}

/// Factory: creates a MockChoraClient for local dev / CI.
pub fn make_mock_client() -> std::sync::Arc<dyn AuthorityClient> {
    std::sync::Arc::new(MockChoraClient)
}

#[cfg(test)]
mod tests {
    use ed25519_dalek::{Signer, Verifier};
    #[tokio::test]
    async fn test_george_token_verification() {
        let token_json = r#"{
            "payload": {
                "schema": "chora.continuation.token.v3.1",
                "issuer": "chora-gate-v3.1",
                "iat": 1710867447,
                "exp": 1710868047,
                "ledger_event_id": "763331c3-3b82-485d-b449-0f6f033a5203",
                "resolution_event_id": "ems-resolve-763331c3-3b82-485d-b449-0f6f033a5203",
                "source_capsule_root": "ef7e9de0b541489e249ce4f7c6f49c078d5537be512592c52215e7441222037d",
                "capabilities": [],
                "nonce": "nonce-750001",
                "execution_target": {
                    "aid": "0x1234567890abcdef",
                    "circuit_id": "attest-rs.audit.v2",
                    "intent_hash": "669234f875c1ddc4629d6241ac8157b7cb180b55fd0f08dd05dac732811737b3"
                }
            },
            "signature": "6e6e140025ce60e471a903d787db85c65b5474d743c6e6cda2901b66e63b52e047a4a781815d6df6e9714f8c383a79d8d08661aba774f949f0c242a5884dd201"
        }"#;

        let token: vex_core::ContinuationToken = serde_json::from_str(token_json).unwrap();

        let mut csprng = rand::thread_rng();
        let signing_key = ed25519_dalek::SigningKey::generate(&mut csprng);
        let verifying_key = signing_key.verifying_key();

        // Sign the updated payload
        let jcs_bytes = serde_jcs::to_vec(&token.payload).unwrap();
        let signature = signing_key.sign(&jcs_bytes);

        // Verification logic check
        assert!(
            verifying_key.verify(&jcs_bytes, &signature).is_ok(),
            "Signature verification failed for Canonical JCS encoded Token V3"
        );
    }
}