host-chain-core 0.3.3

//! Generic chain state machine — shared between WASM and native.
//!
//! Tracks chain connection status, processes JSON-RPC responses, and persists
//! chain databases. The state machine is generic over [`ChainStore`] so the
//! same logic works with localStorage (WASM) and the filesystem (native).

use std::collections::HashMap;

use crate::chain::{
    parse_block_number, ChainExtra, ChainId, ChainState, ChainStatus, REQ_ID_PARA_DB_SAVE,
    REQ_ID_RELAY_DB_SAVE,
};
use crate::store::ChainStore;

/// Per-chain state tracked by the state machine.
struct ChainEntry {
    status: ChainStatus,
    para_db_key: String,
    health_id: u64,
}

/// Pure state machine for chain connections — no networking, no async.
///
/// Processes JSON-RPC responses from smoldot (or smoldot-js on WASM),
/// tracks chain status, and persists chain databases via a [`ChainStore`].
pub struct ChainStateMachine<S: ChainStore> {
    chains: HashMap<ChainId, ChainEntry>,
    store: S,
}

impl<S: ChainStore> ChainStateMachine<S> {
    pub fn new(store: S) -> Self {
        Self {
            chains: HashMap::new(),
            store,
        }
    }

    /// Register a chain for tracking. Initial state is `Connecting`.
    pub fn register_chain(&mut self, chain: ChainId) {
        let para_db_key = chain.para_db_key();
        self.chains.insert(
            chain,
            ChainEntry {
                status: ChainStatus {
                    id: chain,
                    name: chain.display_name(),
                    state: ChainState::Connecting,
                    extra: ChainExtra::None,
                },
                para_db_key,
                health_id: 1000,
            },
        );
    }

    /// Unregister a chain (sets state to Disconnected).
    pub fn unregister_chain(&mut self, chain: ChainId) {
        if let Some(entry) = self.chains.get_mut(&chain) {
            entry.status.state = ChainState::Disconnected;
        }
    }

    /// Set the chain state, preserving existing `extra`.
    pub fn set_state(&mut self, chain: ChainId, state: ChainState) {
        if let Some(entry) = self.chains.get_mut(&chain) {
            entry.status.state = state;
        }
    }

    /// Set both state and extra data.
    pub fn set_state_with_extra(&mut self, chain: ChainId, state: ChainState, extra: ChainExtra) {
        if let Some(entry) = self.chains.get_mut(&chain) {
            entry.status.state = state;
            entry.status.extra = extra;
        }
    }

    /// Access the underlying store.
    pub fn store(&self) -> &S {
        &self.store
    }

    /// Process a JSON-RPC response from smoldot for a given chain.
    pub fn process_response(&mut self, chain: ChainId, text: &str) {
        let v: serde_json::Value = match serde_json::from_str(text) {
            Ok(v) => v,
            Err(_) => return,
        };

        let entry = match self.chains.get_mut(&chain) {
            Some(e) => e,
            None => return,
        };

        // Check for para DB save response.
        if let Some(id) = v.get("id").and_then(|i| i.as_u64()) {
            if id == REQ_ID_PARA_DB_SAVE {
                if let Some(db) = v.get("result").and_then(|r| r.as_str()) {
                    self.store.save(&entry.para_db_key, db);
                    log::info!("{chain:?}: saved para DB ({} bytes)", db.len());
                } else if let Some(err) = v.get("error") {
                    log::warn!("{chain:?}: para DB save returned error: {err}");
                }
                return;
            }
        }

        // Check for system_health response.
        if let Some(result) = v.get("result") {
            if let (Some(peers), Some(is_syncing)) = (
                result.get("peers").and_then(|p| p.as_u64()),
                result.get("isSyncing").and_then(|s| s.as_bool()),
            ) {
                let current_block = match &entry.status.state {
                    ChainState::Live { best_block, .. }
                    | ChainState::Syncing { best_block, .. } => *best_block,
                    _ => 0,
                };

                entry.status.state = if is_syncing {
                    ChainState::Syncing {
                        best_block: current_block,
                        peers: peers as u32,
                    }
                } else {
                    ChainState::Live {
                        best_block: current_block,
                        peers: peers as u32,
                    }
                };
                return;
            }
        }

        // Check for chain_newHead subscription notification.
        if let Some(block) = parse_block_number(text) {
            let (current_peers, current_syncing) = match &entry.status.state {
                ChainState::Live { peers, .. } => (*peers, false),
                ChainState::Syncing { peers, .. } => (*peers, true),
                ChainState::Connecting => (0, true),
                _ => (0, false),
            };

            entry.status.state = if current_syncing && current_peers > 0 {
                ChainState::Syncing {
                    best_block: block,
                    peers: current_peers,
                }
            } else {
                ChainState::Live {
                    best_block: block,
                    peers: current_peers,
                }
            };
        }
    }

    /// Process a JSON-RPC response for a relay chain DB save.
    pub fn process_relay_response(&mut self, chain: ChainId, text: &str) {
        if let Ok(v) = serde_json::from_str::<serde_json::Value>(text) {
            if v.get("id").and_then(|i| i.as_u64()) == Some(REQ_ID_RELAY_DB_SAVE) {
                if let Some(db) = v.get("result").and_then(|r| r.as_str()) {
                    self.store.save(chain.relay_db_key(), db);
                    log::info!("{chain:?}: saved relay DB ({} bytes)", db.len());
                } else if let Some(err) = v.get("error") {
                    log::warn!("{chain:?}: relay DB save returned error: {err}");
                }
            }
        }
    }

    /// Set the chain state to Error.
    pub fn set_error(&mut self, chain: ChainId, msg: String) {
        if let Some(entry) = self.chains.get_mut(&chain) {
            entry.status.state = ChainState::Error(msg);
        }
    }

    pub fn status(&self, chain: ChainId) -> ChainStatus {
        self.chains
            .get(&chain)
            .map(|e| e.status.clone())
            .unwrap_or_else(|| ChainStatus::disconnected(chain))
    }

    pub fn all_statuses(&self) -> Vec<ChainStatus> {
        ChainId::all().iter().map(|&id| self.status(id)).collect()
    }

    /// Generate the chain_subscribeNewHeads JSON-RPC request.
    pub fn subscribe_new_heads_request() -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "chain_subscribeNewHeads",
            "params": []
        })
        .to_string()
    }

    /// Generate a system_health JSON-RPC request with an incrementing ID.
    /// Returns `None` if the chain is not registered.
    pub fn health_check_request(&mut self, chain: ChainId) -> Option<String> {
        let entry = self.chains.get_mut(&chain)?;
        entry.health_id += 1;
        Some(
            serde_json::json!({
                "jsonrpc": "2.0",
                "id": entry.health_id,
                "method": "system_health",
                "params": []
            })
            .to_string(),
        )
    }

    /// Generate a parachain DB save request.
    pub fn para_db_save_request() -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": REQ_ID_PARA_DB_SAVE,
            "method": "chainHead_unstable_finalizedDatabase",
            "params": []
        })
        .to_string()
    }

    /// Generate a relay chain DB save request.
    pub fn relay_db_save_request() -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": REQ_ID_RELAY_DB_SAVE,
            "method": "chainHead_unstable_finalizedDatabase",
            "params": []
        })
        .to_string()
    }

    /// Load persisted relay chain DB.
    pub fn load_relay_db(&self, chain: ChainId) -> String {
        self.store.load(chain.relay_db_key())
    }

    /// Load persisted parachain DB.
    pub fn load_para_db(&self, chain: ChainId) -> String {
        self.store.load(&chain.para_db_key())
    }

    /// Get chain specs for a smoldot chain. Returns (relay_spec, para_spec).
    pub fn chain_specs(chain: ChainId) -> Option<(&'static str, &'static str)> {
        chain.chain_specs()
    }

    // -- Statement store RPC request generators --

    /// Generate a statement_submit JSON-RPC request.
    pub fn statement_submit_request(encoded_hex: &str, request_id: u64) -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "method": "statement_submit",
            "params": [encoded_hex]
        })
        .to_string()
    }

    /// Generate a statement_subscribeStatement JSON-RPC request.
    pub fn statement_subscribe_request(request_id: u64) -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "method": "statement_subscribeStatement",
            "params": ["any"]
        })
        .to_string()
    }

    /// Generate a statement_unsubscribeStatement JSON-RPC request.
    pub fn statement_unsubscribe_request(sub_id: &str, request_id: u64) -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "method": "statement_unsubscribeStatement",
            "params": [sub_id]
        })
        .to_string()
    }

    /// Generate a statement_broadcastsStatement JSON-RPC request.
    pub fn statement_broadcasts_request(topic_hexes: &[String], request_id: u64) -> String {
        serde_json::json!({
            "jsonrpc": "2.0",
            "id": request_id,
            "method": "statement_broadcastsStatement",
            "params": [topic_hexes]
        })
        .to_string()
    }

    /// Parse a statement_subscribeStatement notification.
    /// Returns hex-encoded statement strings found in the notification,
    /// or an empty vec if this is not a statement notification.
    pub fn parse_statement_notification(text: &str) -> Vec<String> {
        let v: serde_json::Value = match serde_json::from_str(text) {
            Ok(v) => v,
            Err(_) => return Vec::new(),
        };

        if v.get("method").and_then(|m| m.as_str()) != Some("statement_subscribeStatement") {
            return Vec::new();
        }

        let result = match v.pointer("/params/result") {
            Some(r) => r,
            None => return Vec::new(),
        };

        let stmts = result
            .pointer("/data/statements")
            .or_else(|| result.pointer("/newStatements/statements"))
            .or_else(|| result.get("statements"));

        match stmts.and_then(|s| s.as_array()) {
            Some(arr) => arr
                .iter()
                .filter_map(|v| v.as_str().map(String::from))
                .collect(),
            None => Vec::new(),
        }
    }
}

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

    struct InMemoryStore {
        data: RefCell<HashMap<String, String>>,
    }

    impl InMemoryStore {
        fn new() -> Self {
            Self {
                data: RefCell::new(HashMap::new()),
            }
        }
    }

    impl ChainStore for InMemoryStore {
        fn load(&self, key: &str) -> String {
            self.data.borrow().get(key).cloned().unwrap_or_default()
        }

        fn save(&self, key: &str, data: &str) {
            self.data
                .borrow_mut()
                .insert(key.to_string(), data.to_string());
        }
    }

    fn make_sm() -> ChainStateMachine<InMemoryStore> {
        let store = InMemoryStore::new();
        let mut sm = ChainStateMachine::new(store);
        sm.register_chain(ChainId::PaseoAssetHub);
        sm
    }

    #[test]
    fn register_sets_connecting() {
        let sm = make_sm();
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(status.state, ChainState::Connecting));
    }

    #[test]
    fn unregister_sets_disconnected() {
        let mut sm = make_sm();
        sm.unregister_chain(ChainId::PaseoAssetHub);
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(status.state, ChainState::Disconnected));
    }

    #[test]
    fn unregistered_chain_returns_disconnected() {
        let sm = make_sm();
        let status = sm.status(ChainId::Ethereum);
        assert!(matches!(status.state, ChainState::Disconnected));
    }

    #[test]
    fn health_response_sets_live() {
        let mut sm = make_sm();
        let resp = r#"{"jsonrpc":"2.0","id":1001,"result":{"peers":5,"isSyncing":false}}"#;
        sm.process_response(ChainId::PaseoAssetHub, resp);
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(status.state, ChainState::Live { peers: 5, .. }));
    }

    #[test]
    fn health_response_sets_syncing() {
        let mut sm = make_sm();
        let resp = r#"{"jsonrpc":"2.0","id":1001,"result":{"peers":3,"isSyncing":true}}"#;
        sm.process_response(ChainId::PaseoAssetHub, resp);
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(status.state, ChainState::Syncing { peers: 3, .. }));
    }

    #[test]
    fn new_head_updates_block_number() {
        let mut sm = make_sm();
        // First set to Live via health
        let health = r#"{"jsonrpc":"2.0","id":1001,"result":{"peers":5,"isSyncing":false}}"#;
        sm.process_response(ChainId::PaseoAssetHub, health);
        // Then new head notification
        let head =
            r#"{"jsonrpc":"2.0","method":"chain_newHead","params":{"result":{"number":"0x1a4"}}}"#;
        sm.process_response(ChainId::PaseoAssetHub, head);
        let status = sm.status(ChainId::PaseoAssetHub);
        match status.state {
            ChainState::Live {
                best_block, peers, ..
            } => {
                assert_eq!(best_block, 0x1a4);
                assert_eq!(peers, 5);
            }
            other => panic!("expected Live, got {other:?}"),
        }
    }

    #[test]
    fn para_db_save_stores_to_store() {
        let mut sm = make_sm();
        let resp = format!(
            r#"{{"jsonrpc":"2.0","id":{},"result":"saved-db-content"}}"#,
            REQ_ID_PARA_DB_SAVE,
        );
        sm.process_response(ChainId::PaseoAssetHub, &resp);
        assert_eq!(sm.store().load("PaseoAssetHub"), "saved-db-content");
    }

    #[test]
    fn relay_db_save_stores_to_store() {
        let mut sm = make_sm();
        let resp = format!(
            r#"{{"jsonrpc":"2.0","id":{},"result":"relay-db-content"}}"#,
            REQ_ID_RELAY_DB_SAVE,
        );
        sm.process_relay_response(ChainId::PaseoAssetHub, &resp);
        assert_eq!(
            sm.store().load(ChainId::PaseoAssetHub.relay_db_key()),
            "relay-db-content"
        );
    }

    #[test]
    fn set_state_preserves_extra() {
        let mut sm = make_sm();
        sm.set_state_with_extra(
            ChainId::PaseoAssetHub,
            ChainState::Live {
                best_block: 100,
                peers: 5,
            },
            ChainExtra::Eth {
                finalized_block: 50,
                gas_price_gwei: 20,
            },
        );
        // set_state (without extra) should preserve the Eth extra
        sm.set_state(
            ChainId::PaseoAssetHub,
            ChainState::Live {
                best_block: 200,
                peers: 3,
            },
        );
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(
            status.extra,
            ChainExtra::Eth {
                finalized_block: 50,
                gas_price_gwei: 20
            }
        ));
        assert!(matches!(
            status.state,
            ChainState::Live {
                best_block: 200,
                peers: 3
            }
        ));
    }

    #[test]
    fn set_state_with_extra_updates_both() {
        let mut sm = make_sm();
        sm.set_state_with_extra(
            ChainId::PaseoAssetHub,
            ChainState::Live {
                best_block: 100,
                peers: 5,
            },
            ChainExtra::Btc {
                tip_height: 800000,
                fee_rate_sat_vb: 10,
            },
        );
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(
            status.state,
            ChainState::Live {
                best_block: 100,
                peers: 5
            }
        ));
        assert!(matches!(
            status.extra,
            ChainExtra::Btc {
                tip_height: 800000,
                fee_rate_sat_vb: 10
            }
        ));
    }

    #[test]
    fn process_response_unregistered_chain_is_noop() {
        let mut sm = make_sm();
        let resp = r#"{"jsonrpc":"2.0","id":1001,"result":{"peers":5,"isSyncing":false}}"#;
        // Ethereum is not registered — should not panic or change anything
        sm.process_response(ChainId::Ethereum, resp);
        let status = sm.status(ChainId::Ethereum);
        assert!(matches!(status.state, ChainState::Disconnected));
    }

    #[test]
    fn health_check_request_unregistered_returns_none() {
        let mut sm = make_sm();
        assert!(sm.health_check_request(ChainId::Ethereum).is_none());
    }

    #[test]
    fn health_check_request_id_starts_above_1000_and_increments() {
        let mut sm = make_sm();
        let req1 = sm.health_check_request(ChainId::PaseoAssetHub).unwrap();
        let req2 = sm.health_check_request(ChainId::PaseoAssetHub).unwrap();

        let v1: serde_json::Value = serde_json::from_str(&req1).unwrap();
        let v2: serde_json::Value = serde_json::from_str(&req2).unwrap();

        let id1 = v1["id"].as_u64().unwrap();
        let id2 = v2["id"].as_u64().unwrap();

        assert!(id1 > 1000);
        assert_eq!(id2, id1 + 1);
    }

    #[test]
    fn para_db_save_error_does_not_change_state() {
        let mut sm = make_sm();
        // Set to Live first
        let health = r#"{"jsonrpc":"2.0","id":1001,"result":{"peers":5,"isSyncing":false}}"#;
        sm.process_response(ChainId::PaseoAssetHub, health);

        // DB save error response — should not trigger health or new-head branches
        let error_resp = format!(
            r#"{{"jsonrpc":"2.0","id":{},"error":{{"code":-32000,"message":"db error"}}}}"#,
            REQ_ID_PARA_DB_SAVE,
        );
        sm.process_response(ChainId::PaseoAssetHub, &error_resp);

        // State should still be Live with peers=5
        let status = sm.status(ChainId::PaseoAssetHub);
        assert!(matches!(status.state, ChainState::Live { peers: 5, .. }));
        // And nothing should have been saved to the store
        assert_eq!(sm.store().load("PaseoAssetHub"), "");
    }
}