Encoding and transport

use std::collections::HashMap;
use std::process::ExitCode;

use statesync_gkr::compiler::{
    AssetId, LayerStrategy, LeafPayload, LeafState, MerklePath, PublicInputs, SmtOpKind,
    SmtOperation, SmtParams, SmtWitness,
};
use statesync_gkr::primitives::field::{BaseField, PrimeCharacteristicRing};
use statesync_gkr::primitives::hash::{Digest, HashGadget, Poseidon2Gadget};
use statesync_gkr::wrap::encoding::{DecodeError, decode_inner_proof};
use statesync_gkr::{PreparedSync, StateSyncGkrConfig, StateSyncProver, SyncRequest};

/// One message on the wire: the requester's identifier and the encoded proof.
#[derive(Clone, Debug)]
struct ProofMessage {
    request_id: u64,
    proof: Vec<u8>,
}

/// The verifying side: its own prover instance and preparation, and the
/// original requests it is waiting for, keyed by request identifier.
struct Verifier {
    engine: StateSyncProver,
    prepared: PreparedSync,
    pending: HashMap<u64, SyncRequest>,
}

impl Verifier {
    fn new(config: StateSyncGkrConfig) -> Result<Self, String> {
        let engine = StateSyncProver::new(config);
        let prepared = engine
            .prepare(SmtOpKind::Membership)
            .map_err(|error| format!("verifier preparation failed: {error:?}"))?;
        Ok(Self {
            engine,
            prepared,
            pending: HashMap::new(),
        })
    }

    /// Record a request before anyone is asked to prove it.
    fn expect(&mut self, request_id: u64, request: SyncRequest) {
        self.pending.insert(request_id, request);
    }

    /// Check one received message against the original request it names.
    fn receive(&mut self, message: &ProofMessage) -> String {
        let Some(request) = self.pending.get(&message.request_id) else {
            return format!("request {}: unknown or already settled", message.request_id);
        };
        // The original request selects the preparation; the bytes do not.
        if request.operation.kind() != self.prepared.kind() {
            return format!("request {}: no preparation for its kind", message.request_id);
        }
        if self
            .engine
            .verify_encoded_sync_op(&self.prepared, request, &message.proof)
        {
            self.pending.remove(&message.request_id);
            return format!("request {}: accepted", message.request_id);
        }
        format!(
            "request {}: rejected ({})",
            message.request_id,
            self.diagnose(request, &message.proof)
        )
    }

    /// Explain a rejection. Diagnostics only: the verdict is the `false` above.
    fn diagnose(&self, request: &SyncRequest, bytes: &[u8]) -> String {
        let envelope = match decode_inner_proof(bytes) {
            Ok(envelope) => envelope,
            Err(DecodeError::DeclaredLengthMismatch { .. }) => {
                return "received length differs from the declared length".to_owned();
            }
            Err(error) => return format!("malformed encoding: {error:?}"),
        };
        match self.engine.circuit_identity(&self.prepared) {
            Ok(expected) if expected == envelope.identity => {}
            Ok(_) => return "proof was made for another circuit or configuration".to_owned(),
            Err(error) => return format!("no local circuit identity: {error:?}"),
        }
        if envelope.public_inputs != request.public_inputs {
            return "proof is for different public inputs than the original request".to_owned();
        }
        "identity and public inputs match; the request or the proof failed verification".to_owned()
    }
}

/// The proving side: prove one request and encode the result for transport.
fn prove_message(
    prover: &StateSyncProver,
    prepared: &PreparedSync,
    request_id: u64,
    request: &SyncRequest,
) -> Result<ProofMessage, String> {
    let result = prover
        .prove_sync_op_prepared(prepared, request)
        .map_err(|error| format!("request {request_id}: proving failed: {error:?}"))?;
    let proof = prover
        .encode_sync_result(prepared, &result)
        .map_err(|error| format!("request {request_id}: encoding failed: {error:?}"))?;
    Ok(ProofMessage { request_id, proof })
}

fn field(value: u32) -> BaseField {
    BaseField::from_u32(value)
}

/// A synthetic membership request. A real service reads the root, the leaf
/// and the path from its own state store.
fn membership_request(params: &SmtParams, key: u64) -> Result<SyncRequest, String> {
    let hasher = Poseidon2Gadget::new(params.leaf_max_fields as usize);
    let key = AssetId(key);
    let payload = LeafPayload {
        sync_state: vec![field(100), field(7)],
        identity_digest: [9_u8; 32],
    };
    let leaf = LeafState::Occupied(payload.clone());
    let path = MerklePath {
        siblings: (0..params.depth)
            .map(|level| Digest([field(level * 13 + 1); 8]))
            .collect(),
    };
    let root = path
        .compute_root(&hasher, params, key, &leaf)
        .map_err(|error| format!("root construction failed: {error:?}"))?;
    let value_digest = hasher
        .hash_leaf(&leaf.encode())
        .map_err(|error| format!("leaf hashing failed: {error:?}"))?;
    let operation = SmtOperation::Membership { key, payload };
    let op_kind_tag = PublicInputs::kind_tag(operation.kind());
    Ok(SyncRequest {
        operation,
        witness: SmtWitness { leaf, path },
        public_inputs: PublicInputs {
            old_root: root,
            new_root: root,
            op_kind_tag,
            asset_id: key,
            value_digest,
        },
    })
}

fn run() -> Result<(), String> {
    let config = StateSyncGkrConfig {
        smt: SmtParams {
            depth: 4,
            ..Default::default()
        },
        layer_strategy: LayerStrategy::A,
        ..Default::default()
    };
    let params = config.smt;

    // Each side prepares from the same configuration, as two processes would.
    let prover = StateSyncProver::new(config.clone());
    let prover_prepared = prover
        .prepare(SmtOpKind::Membership)
        .map_err(|error| format!("prover preparation failed: {error:?}"))?;
    let mut verifier = Verifier::new(config)?;

    // The verifying side records each request under its identifier first.
    let request_1 = membership_request(&params, 5)?;
    let request_2 = membership_request(&params, 6)?;
    let request_3 = membership_request(&params, 7)?;
    let request_4 = membership_request(&params, 8)?;
    verifier.expect(1, request_1.clone());
    verifier.expect(2, request_2.clone());
    verifier.expect(3, request_3.clone());
    verifier.expect(4, request_4);

    // The proving side answers requests 1 to 3.
    let reply_1 = prove_message(&prover, &prover_prepared, 1, &request_1)?;
    let reply_2 = prove_message(&prover, &prover_prepared, 2, &request_2)?;
    let reply_3 = prove_message(&prover, &prover_prepared, 3, &request_3)?;

    // Replies can arrive in any order.
    println!("{}", verifier.receive(&reply_2));
    println!("{}", verifier.receive(&reply_1));
    // A settled identifier is not accepted twice.
    println!("{}", verifier.receive(&reply_1));
    // Damaged in transit: the last byte is missing.
    let mut damaged = reply_3.clone();
    damaged.proof.pop();
    println!("{}", verifier.receive(&damaged));
    // A valid proof delivered under another request's identifier.
    let misrouted = ProofMessage {
        request_id: 4,
        proof: reply_3.proof.clone(),
    };
    println!("{}", verifier.receive(&misrouted));
    // The intact reply still verifies.
    println!("{}", verifier.receive(&reply_3));
    Ok(())
}

fn main() -> ExitCode {
    match run() {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("transport=FAIL: {error}");
            ExitCode::FAILURE
        }
    }
}

request 2: accepted
request 1: accepted
request 1: unknown or already settled
request 3: rejected (received length differs from the declared length)
request 4: rejected (proof is for different public inputs than the original request)
request 3: accepted