trazaeo 0.5.7

use crate::error::{TrazaeoError, TrazaeoResult};
use crate::utils::Hash;
use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
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,
    pub attestor_signature: String,
}

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

#[derive(Debug, Clone)]
pub struct LocalAttestorSigner {
    signing_key: SigningKey,
}

impl LocalAttestorSigner {
    pub fn from_seed(seed: [u8; 32]) -> Self {
        Self {
            signing_key: SigningKey::from_bytes(&seed),
        }
    }

    pub fn attestor_pubkey(&self) -> [u8; 32] {
        self.signing_key.verifying_key().to_bytes()
    }
}

pub fn attestor_key_ref(attestor_pubkey: &[u8; 32]) -> String {
    hex::encode(attestor_pubkey)
}

fn commit_anchor_v1_signing_bytes(
    cluster: &str,
    program_id: &str,
    payload: &CommitAnchorV1Payload,
) -> Vec<u8> {
    let mut bytes = Vec::with_capacity(179);
    bytes.extend_from_slice(b"trazaeo:commit_anchor_v1");
    bytes.extend_from_slice(&(cluster.len() as u32).to_le_bytes());
    bytes.extend_from_slice(cluster.as_bytes());
    bytes.extend_from_slice(&(program_id.len() as u32).to_le_bytes());
    bytes.extend_from_slice(program_id.as_bytes());
    bytes.push(payload.schema_version);
    bytes.extend_from_slice(&payload.anchored_envelope_hash);
    bytes.extend_from_slice(&payload.anchored_checkpoint_hash);
    bytes.extend_from_slice(&payload.anchored_root_hash);
    bytes.extend_from_slice(&payload.anchored_unix_seconds.to_le_bytes());
    bytes.extend_from_slice(&payload.prev_anchor_hash);
    bytes
}

/// 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: &LocalAttestorSigner,
    payload: CommitAnchorV1Payload,
) -> TrazaeoResult<TxResult> {
    if payload.schema_version != 1 {
        return Err(TrazaeoError::invalid_input(
            "commit anchor v1",
            "unsupported schema_version",
        ));
    }

    let attestor_pubkey = attestor.attestor_pubkey();
    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 signing_bytes =
        commit_anchor_v1_signing_bytes(&client.config.cluster, &client.config.program_id, &payload);
    let attestor_signature = attestor.signing_key.sign(&signing_bytes);

    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(&attestor_pubkey);
    tx_hasher.update(&attestor_signature.to_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,
        attestor_signature: hex::encode(attestor_signature.to_bytes()),
    };

    *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()))
}

/// Verifies that an anchor transaction was authorized by the account attestor.
pub fn verify_anchor_transaction_signature(tx: &TxResult, account: &AnchorAccountV1) -> bool {
    let payload = CommitAnchorV1Payload {
        schema_version: account.schema_version,
        anchored_envelope_hash: account.anchored_envelope_hash,
        anchored_checkpoint_hash: account.anchored_checkpoint_hash,
        anchored_root_hash: account.anchored_root_hash,
        anchored_unix_seconds: account.anchored_unix_seconds,
        prev_anchor_hash: account.prev_anchor_hash,
    };
    let signing_bytes = commit_anchor_v1_signing_bytes(&tx.cluster, &tx.program_id, &payload);
    let Ok(signature_bytes) = hex::decode(&tx.attestor_signature) else {
        return false;
    };
    let Ok(signature_array) = <[u8; 64]>::try_from(signature_bytes.as_slice()) else {
        return false;
    };
    let Ok(verifying_key) = VerifyingKey::from_bytes(&account.attestor_pubkey) else {
        return false;
    };
    let signature = Signature::from_bytes(&signature_array);
    verifying_key.verify(&signing_bytes, &signature).is_ok()
}

/// 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],
    };
    let signer = LocalAttestorSigner::from_seed(attestor_pubkey);
    commit_anchor_v1(client, &signer, 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 = LocalAttestorSigner::from_seed([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 = LocalAttestorSigner::from_seed([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]);
    }
}