trazaeo 0.5.2

use crate::error::{TrazaeoError, TrazaeoResult};
use crate::utils::Hash;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SolanaConfig {
    pub cluster: String,
    pub program_id: String,
}

#[derive(Debug, Clone)]
pub struct SolanaClient {
    config: SolanaConfig,
    ledger: Arc<Mutex<HashMap<String, AnchorCommitRecord>>>,
    latest_root: Arc<Mutex<Hash>>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CommitAnchorV1Payload {
    pub schema_version: u8,
    pub anchored_envelope_hash: [u8; 32],
    pub anchored_checkpoint_hash: [u8; 32],
    pub anchored_root_hash: [u8; 32],
    pub anchored_unix_seconds: i64,
    pub prev_anchor_hash: [u8; 32],
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnchorAccountV1 {
    pub schema_version: u8,
    pub anchored_envelope_hash: [u8; 32],
    pub anchored_checkpoint_hash: [u8; 32],
    pub anchored_root_hash: [u8; 32],
    pub attestor_pubkey: [u8; 32],
    pub anchored_unix_seconds: i64,
    pub anchored_slot: u64,
    pub prev_anchor_hash: [u8; 32],
    pub bump: u8,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TxResult {
    pub signature: String,
    pub cluster: String,
    pub program_id: String,
    pub slot: u64,
    pub finalized: bool,
    pub anchor_account_pda: String,
}

#[derive(Debug, Clone)]
struct AnchorCommitRecord {
    tx: TxResult,
    account: AnchorAccountV1,
}

/// Handles init solana client.
pub fn init_solana_client(config: &SolanaConfig) -> SolanaClient {
    SolanaClient {
        config: config.clone(),
        ledger: Arc::new(Mutex::new(HashMap::new())),
        latest_root: Arc::new(Mutex::new(Hash([0u8; 32]))),
    }
}

/// Handles cluster name.
pub fn cluster_name(client: &SolanaClient) -> String {
    client.config.cluster.clone()
}

/// Handles program id.
pub fn program_id(client: &SolanaClient) -> String {
    client.config.program_id.clone()
}

/// Handles derive anchor account pda.
pub fn derive_anchor_account_pda(
    program_id: &str,
    attestor_pubkey: &[u8; 32],
    anchored_envelope_hash: &[u8; 32],
) -> (String, u8) {
    let mut hasher = blake3::Hasher::new();
    hasher.update(program_id.as_bytes());
    hasher.update(b"mk1_anchor_v1");
    hasher.update(attestor_pubkey);
    hasher.update(anchored_envelope_hash);
    let digest = hasher.finalize();
    let pda = hex::encode(digest.as_bytes());
    let bump = digest.as_bytes()[0];
    (pda, bump)
}

/// Handles commit anchor v1.
pub fn commit_anchor_v1(
    client: &SolanaClient,
    attestor_pubkey: [u8; 32],
    payload: CommitAnchorV1Payload,
) -> TrazaeoResult<TxResult> {
    if payload.schema_version != 1 {
        return Err(TrazaeoError::invalid_input(
            "commit anchor v1",
            "unsupported schema_version",
        ));
    }

    let (anchor_account_pda, bump) = derive_anchor_account_pda(
        &client.config.program_id,
        &attestor_pubkey,
        &payload.anchored_envelope_hash,
    );

    let mut ledger = client
        .ledger
        .lock()
        .map_err(|_| TrazaeoError::external("commit anchor v1", "failed to acquire ledger lock"))?;

    if ledger
        .values()
        .any(|record| record.tx.anchor_account_pda == anchor_account_pda)
    {
        return Err(TrazaeoError::invalid_input(
            "commit anchor v1",
            "anchor account already initialized",
        ));
    }

    let slot = (ledger.len() as u64) + 1;
    let mut tx_hasher = blake3::Hasher::new();
    tx_hasher.update(client.config.cluster.as_bytes());
    tx_hasher.update(client.config.program_id.as_bytes());
    tx_hasher.update(&slot.to_le_bytes());
    tx_hasher.update(&payload.anchored_envelope_hash);
    tx_hasher.update(&payload.anchored_checkpoint_hash);
    tx_hasher.update(&payload.anchored_root_hash);
    let signature = hex::encode(tx_hasher.finalize().as_bytes());

    let account = AnchorAccountV1 {
        schema_version: payload.schema_version,
        anchored_envelope_hash: payload.anchored_envelope_hash,
        anchored_checkpoint_hash: payload.anchored_checkpoint_hash,
        anchored_root_hash: payload.anchored_root_hash,
        attestor_pubkey,
        anchored_unix_seconds: payload.anchored_unix_seconds,
        anchored_slot: slot,
        prev_anchor_hash: payload.prev_anchor_hash,
        bump,
    };

    let tx = TxResult {
        signature: signature.clone(),
        cluster: client.config.cluster.clone(),
        program_id: client.config.program_id.clone(),
        slot,
        finalized: true,
        anchor_account_pda,
    };

    *client.latest_root.lock().map_err(|_| {
        TrazaeoError::external("commit anchor v1", "failed to acquire latest_root lock")
    })? = Hash(payload.anchored_root_hash);

    ledger.insert(
        signature,
        AnchorCommitRecord {
            tx: tx.clone(),
            account,
        },
    );
    Ok(tx)
}

/// Gets transaction.
pub fn get_transaction(client: &SolanaClient, signature: &str) -> TrazaeoResult<Option<TxResult>> {
    let ledger = client
        .ledger
        .lock()
        .map_err(|_| TrazaeoError::external("get transaction", "failed to acquire ledger lock"))?;
    Ok(ledger.get(signature).map(|record| record.tx.clone()))
}

/// Gets anchor account by pda.
pub fn get_anchor_account_by_pda(
    client: &SolanaClient,
    pda: &str,
) -> TrazaeoResult<Option<AnchorAccountV1>> {
    let ledger = client.ledger.lock().map_err(|_| {
        TrazaeoError::external("get anchor account by pda", "failed to acquire ledger lock")
    })?;
    Ok(ledger
        .values()
        .find(|record| record.tx.anchor_account_pda == pda)
        .map(|record| record.account.clone()))
}

/// Sends root tx.
pub fn send_root_tx(
    client: &SolanaClient,
    root_hash: &Hash,
    meta_ref: &str,
) -> TrazaeoResult<TxResult> {
    let mut attestor_pubkey = [0u8; 32];
    let digest = blake3::hash(meta_ref.as_bytes());
    attestor_pubkey.copy_from_slice(digest.as_bytes());
    let payload = CommitAnchorV1Payload {
        schema_version: 1,
        anchored_envelope_hash: [0u8; 32],
        anchored_checkpoint_hash: [0u8; 32],
        anchored_root_hash: root_hash.0,
        anchored_unix_seconds: 0,
        prev_anchor_hash: [0u8; 32],
    };
    commit_anchor_v1(client, attestor_pubkey, payload)
}

/// Gets chain root.
pub fn get_chain_root(client: &SolanaClient) -> TrazaeoResult<Hash> {
    client
        .latest_root
        .lock()
        .map(|value| value.clone())
        .map_err(|_| TrazaeoError::external("get chain root", "failed to acquire latest_root lock"))
}

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

    /// Handles client.
    fn client() -> SolanaClient {
        init_solana_client(&SolanaConfig {
            cluster: "solana-testnet".to_string(),
            program_id: "program-1".to_string(),
        })
    }

    /// Tests that init client and send tx compile contract.
    #[test]
    fn init_client_and_send_tx_compile_contract() {
        let client = client();
        let root = Hash([1u8; 32]);
        let tx = send_root_tx(&client, &root, "meta://ref").expect("send root tx");
        assert!(!tx.signature.is_empty());
        assert_eq!(get_chain_root(&client).expect("chain root"), root);
    }

    /// Tests that commit anchor v1 rejects duplicate account.
    #[test]
    fn commit_anchor_v1_rejects_duplicate_account() {
        let client = client();
        let payload = CommitAnchorV1Payload {
            schema_version: 1,
            anchored_envelope_hash: [2u8; 32],
            anchored_checkpoint_hash: [5u8; 32],
            anchored_root_hash: [3u8; 32],
            anchored_unix_seconds: 1700000000,
            prev_anchor_hash: [0u8; 32],
        };
        let attestor = [4u8; 32];
        let first = commit_anchor_v1(&client, attestor, payload.clone()).expect("first commit");
        assert!(get_transaction(&client, &first.signature)
            .expect("transaction lookup")
            .is_some());
        assert!(commit_anchor_v1(&client, attestor, payload).is_err());
    }

    /// Tests that get anchor account by pda returns account.
    #[test]
    fn get_anchor_account_by_pda_returns_account() {
        let client = client();
        let attestor = [4u8; 32];
        let payload = CommitAnchorV1Payload {
            schema_version: 1,
            anchored_envelope_hash: [9u8; 32],
            anchored_checkpoint_hash: [7u8; 32],
            anchored_root_hash: [8u8; 32],
            anchored_unix_seconds: 1700000000,
            prev_anchor_hash: [0u8; 32],
        };
        let tx = commit_anchor_v1(&client, attestor, payload).expect("commit");
        let account = get_anchor_account_by_pda(&client, &tx.anchor_account_pda)
            .expect("account lookup")
            .expect("account");
        assert_eq!(account.anchored_checkpoint_hash, [7u8; 32]);
        assert_eq!(account.anchored_root_hash, [8u8; 32]);
    }
}