Storage

pub struct PersistenceManager {
    device: Arc<tokio::sync::RwLock<Device>>,
    device_snapshot: std::sync::RwLock<Arc<Device>>,
    backend: Arc<dyn Backend>,
    dirty: Arc<AtomicBool>,
    save_notify: Arc<Event>,  // event_listener::Event (runtime-agnostic)
}

impl PersistenceManager {
    pub async fn new(backend: Arc<dyn Backend>) -> Result<Self, StoreError> {
        // Ensure device row exists
        let exists = backend.exists().await?;
        if !exists {
            backend.create().await?;
        }

        // Load existing data or create new
        let device_data = backend.load().await?;
        let device = if let Some(data) = device_data {
            let mut dev = Device::new(backend.clone());
            dev.load_from_serializable(data);
            dev
        } else {
            Device::new(backend.clone())
        };

        let snapshot = Arc::new(device.clone());
        Ok(Self {
            device: Arc::new(tokio::sync::RwLock::new(device)),
            device_snapshot: std::sync::RwLock::new(snapshot),
            backend,
            dirty: Arc::new(AtomicBool::new(false)),
            save_notify: Arc::new(Event::new()),
        })
    }
}

pub fn get_device_snapshot(&self) -> Arc<Device> {
    self.device_snapshot.read().unwrap().clone()
}

let snapshot = client.persistence_manager.get_device_snapshot();
println!("Device ID: {:?}", snapshot.pn);
println!("Push Name: {}", snapshot.push_name);

pub async fn modify_device<F, R>(&self, modifier: F) -> R
where
    F: FnOnce(&mut Device) -> R,
{
    let mut device_guard = self.device.write().await;
    let result = modifier(&mut device_guard);

    // Rebuild the cached snapshot before releasing the write lock
    *self.device_snapshot.write().unwrap() = Arc::new(device_guard.clone());

    self.dirty.store(true, Ordering::Relaxed);
    self.save_notify.notify(1);

    result
}

client.persistence_manager.modify_device(|device| {
    device.push_name = "New Name".to_string();
}).await;

pub async fn process_command(&self, command: DeviceCommand) {
    self.modify_device(|device| {
        apply_command_to_device(device, command);
    }).await;
}

client.persistence_manager
    .process_command(DeviceCommand::SetPushName("New Name".to_string()))
    .await;

pub fn run_background_saver(self: Arc<Self>, runtime: Arc<dyn Runtime>, interval: Duration) {
    runtime.spawn(Box::pin(async move {
        loop {
            let listener = self.save_notify.listen();
            // Race: wake on notification or after interval
            futures_lite::future::or(listener, async {
                // sleep using runtime abstraction
                runtime.sleep(interval).await;
            }).await;

            if let Err(e) = self.save_to_disk().await {
                error!("Error saving device state: {e}");
            }
        }
    }));
}

async fn save_to_disk(&self) -> Result<(), StoreError> {
    if self.dirty.swap(false, Ordering::AcqRel) {
        let device_guard = self.device.read().await;
        let serializable_device = device_guard.to_serializable();
        drop(device_guard);

        self.backend.save(&serializable_device).await?;
    }
    Ok(())
}

let persistence_manager = Arc::new(PersistenceManager::new(backend).await?);
persistence_manager.clone().run_background_saver(runtime.clone(), Duration::from_secs(30));

pub trait Backend: SignalStore + AppSyncStore + ProtocolStore + DeviceStore + Send + Sync {}

use bytes::Bytes; // from the `bytes` crate
use std::sync::Arc;

#[async_trait]
pub trait SignalStore: Send + Sync {
    // Identity Operations
    async fn put_identity(&self, address: &str, key: [u8; 32]) -> Result<()>;
    async fn load_identity(&self, address: &str) -> Result<Option<[u8; 32]>>;
    async fn delete_identity(&self, address: &str) -> Result<()>;
    async fn delete_identities_batch(&self, addresses: &[Arc<str>]) -> Result<()>; // has default impl looping over delete_identity

    // Session Operations
    async fn get_session(&self, address: &str) -> Result<Option<Bytes>>;
    async fn put_session(&self, address: &str, session: &[u8]) -> Result<()>;
    async fn delete_session(&self, address: &str) -> Result<()>;
    async fn delete_sessions_batch(&self, addresses: &[Arc<str>]) -> Result<()>; // has default impl looping over delete_session
    async fn has_session(&self, address: &str) -> Result<bool>; // has default impl using get_session

    // PreKey Operations
    async fn store_prekey(&self, id: u32, record: &[u8], uploaded: bool) -> Result<()>;
    async fn store_prekeys_batch(&self, keys: &[(u32, Bytes)], uploaded: bool) -> Result<()>; // has default impl looping over store_prekey
    async fn load_prekey(&self, id: u32) -> Result<Option<Bytes>>;
    async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Bytes)>>; // has default impl looping over load_prekey

    async fn remove_prekey(&self, id: u32) -> Result<()>;
    async fn remove_prekeys_batch(&self, ids: &[u32]) -> Result<()>; // has default impl looping over remove_prekey
    /// Mark already-stored prekeys as uploaded. UPDATE semantics — a key
    /// consumed (deleted) between the upload snapshot and this call must stay
    /// deleted, never resurrected by an upsert.
    async fn mark_prekeys_uploaded(&self, ids: &[u32]) -> Result<()>;
    async fn get_max_prekey_id(&self) -> Result<u32>;

    // Signed PreKey Operations
    async fn store_signed_prekey(&self, id: u32, record: &[u8]) -> Result<()>;
    async fn load_signed_prekey(&self, id: u32) -> Result<Option<Vec<u8>>>;
    async fn load_all_signed_prekeys(&self) -> Result<Vec<(u32, Vec<u8>)>>;
    async fn remove_signed_prekey(&self, id: u32) -> Result<()>;

    // Sender Key Operations (for groups)
    async fn put_sender_key(&self, address: &str, record: &[u8]) -> Result<()>;
    async fn get_sender_key(&self, address: &str) -> Result<Option<Vec<u8>>>;
    async fn delete_sender_key(&self, address: &str) -> Result<()>;
    async fn delete_sender_keys_batch(&self, addresses: &[Arc<str>]) -> Result<()>; // has default impl looping over delete_sender_key
}

// Store identity key for a contact
backend.put_identity(
    "15551234567@s.whatsapp.net:0",
    identity_key
).await?;

// Load session for decryption
if let Some(session) = backend.get_session("15551234567@s.whatsapp.net:0").await? {
    // Decrypt message using session
}

pub struct SignalStoreCache {
    sessions: Mutex<SessionStoreState>,
    identities: Mutex<ByteStoreState>,
    sender_keys: Mutex<SenderKeyStoreState>,
}

// Session object cache — no per-message serialize/deserialize
struct SessionStoreState {
    cache: HashMap<Arc<str>, Option<SessionRecord>>,  // None = known-absent
    dirty: HashSet<Arc<str>>,                          // Modified keys pending flush
    deleted: HashSet<Arc<str>>,                        // Deleted keys pending flush
}

// Sender key object cache (same pattern as sessions)
struct SenderKeyStoreState {
    cache: HashMap<Arc<str>, Option<SenderKeyRecord>>,  // None = known-absent
    dirty: HashSet<Arc<str>>,                            // Modified keys pending flush
}

// Byte cache for identities
struct ByteStoreState {
    cache: HashMap<Arc<str>, Option<Arc<[u8]>>>,  // None = known-absent
    dirty: HashSet<Arc<str>>,
    deleted: HashSet<Arc<str>>,
}

// Read (loads from backend if not cached, returns SessionRecord object)
let session = cache.get_session(address, &backend).await?;

// Check existence (cache hit is O(1), cold load caches for subsequent get)
let exists = cache.has_session(address, &backend).await?;

// Write (takes ownership, marks as dirty, doesn't hit backend)
cache.put_session(address, record).await;

// Delete (marks for deletion on flush)
cache.delete_session(address).await;

// Persist all dirty state to backend (sessions serialized here)
cache.flush(&backend).await?;

// Clear cache (on disconnect/reconnect, retains allocated capacity)
cache.clear().await;

#[async_trait]
pub trait AppSyncStore: Send + Sync {
    // Sync Keys
    async fn get_sync_key(&self, key_id: &[u8]) -> Result<Option<AppStateSyncKey>>;
    async fn set_sync_key(&self, key_id: &[u8], key: AppStateSyncKey) -> Result<()>;

    // Version Tracking
    // `None` means the collection has never synced — distinct from a collection
    // that synced and is legitimately empty, which sits at version 0 with a record.
    async fn get_version(&self, name: &str) -> Result<Option<HashState>>;
    // Forgets a collection's version, returning it to the never-synced state
    // (used to force a rebuild from a fresh snapshot).
    async fn delete_version(&self, name: &str) -> Result<()>;
    async fn set_version(&self, name: &str, state: HashState) -> Result<()>;

    // Mutation MACs
    async fn put_mutation_macs(
        &self,
        name: &str,
        version: u64,
        mutations: &[AppStateMutationMAC],
    ) -> Result<()>;
    async fn get_mutation_mac(&self, name: &str, index_mac: &[u8]) -> Result<Option<Vec<u8>>>;
    // Batched previous-MAC lookup (v0.6); default loops over get_mutation_mac.
    // index_macs are inline [u8; 32] (an index MAC is always a full
    // HMAC-SHA256 output) — no per-MAC heap allocation on either side.
    async fn get_mutation_macs(
        &self,
        name: &str,
        index_macs: &[[u8; 32]],
    ) -> Result<HashMap<[u8; 32], Vec<u8>>>;
    async fn delete_mutation_macs(&self, name: &str, index_macs: &[Vec<u8>]) -> Result<()>;
    async fn get_latest_sync_key_id(&self) -> Result<Option<Vec<u8>>>;
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AppStateSyncKey {
    pub key_data: Vec<u8>,
    pub fingerprint: Vec<u8>,
    pub timestamp: i64,
}

pub struct HashState {
    pub version: u64,
    pub hash: [u8; 128],
    pub index_value_map: HashMap<String, Vec<u8>>,
    pub mac_mismatch_fatal: bool,
    // Set only when a sync run reaches its terminal page — a bootstrap
    // deferred mid-run leaves this false even though version already moved.
    // `#[serde(default)]`, so rows written before this field existed decode
    // it as `false`. That only forces a re-bootstrap for a legacy row still
    // sitting at version 0 — the ambiguous case this field exists to settle.
    // A legacy row with a real version already has `has_baseline()` true and
    // asks for patches as before.
    pub bootstrapped: bool,
}

#[async_trait]
pub trait ProtocolStore: Send + Sync {
    // Per-Device Sender Key Tracking (matches WA Web's participant.senderKey Map)
    // Updated only AFTER the server ACKs the message stanza,
    // preventing stale entries from network failures mid-send.
    async fn get_sender_key_devices(&self, group_jid: &str) -> Result<Vec<(String, bool)>>;
    async fn set_sender_key_status(&self, group_jid: &str, entries: &[(&str, bool)]) -> Result<()>;
    async fn clear_sender_key_devices(&self, group_jid: &str) -> Result<()>;

    // Clear all sender key device tracking across ALL groups.
    // Provided for completeness; not invoked on identity change — WhatsApp Web's
    // `WAWebUpdateLocalSignalSession` only calls `markForgetSenderKey` per-group/per-device
    // on retry receipts, so a global wipe here would empty the tracker too aggressively.
    async fn clear_all_sender_key_devices(&self) -> Result<()>;

    // LID-PN Mapping (Long-term ID to Phone Number)
    async fn get_lid_mapping(&self, lid: &str) -> Result<Option<LidPnMappingEntry>>;
    async fn get_pn_mapping(&self, phone: &str) -> Result<Option<LidPnMappingEntry>>;
    async fn put_lid_mapping(&self, entry: &LidPnMappingEntry) -> Result<()>;
    async fn get_all_lid_mappings(&self) -> Result<Vec<LidPnMappingEntry>>;

    // Base Key Collision Detection
    async fn save_base_key(&self, address: &str, message_id: &str, base_key: &[u8]) -> Result<()>;
    async fn has_same_base_key(&self, address: &str, message_id: &str, current_base_key: &[u8]) -> Result<bool>;
    async fn delete_base_key(&self, address: &str, message_id: &str) -> Result<()>;

    // Device Registry
    async fn update_device_list(&self, record: DeviceListRecord) -> Result<()>;
    async fn get_devices(&self, user: &str) -> Result<Option<DeviceListRecord>>;

    // Delete a device list record, forcing a network re-fetch on next query.
    async fn delete_devices(&self, user: &str) -> Result<()>;

    // TcToken Storage (Trusted Contact Tokens)
    async fn get_tc_token(&self, jid: &str) -> Result<Option<TcTokenEntry>>;
    async fn put_tc_token(&self, jid: &str, entry: &TcTokenEntry) -> Result<()>;
    async fn delete_tc_token(&self, jid: &str) -> Result<()>;
    async fn get_all_tc_token_jids(&self) -> Result<Vec<String>>;
    // Pruned only when BOTH the received-token cutoff and the sender-bucket
    // cutoff have passed, so recent sender-side state outlives an expired token.
    async fn delete_expired_tc_tokens(&self, token_cutoff: i64, sender_cutoff: i64) -> Result<u32>;
    // Advance-only upsert of the sender-side issuance timestamp; inserts a
    // byte-less placeholder if no entry exists yet. Defaulted (read-modify-write).
    async fn touch_tc_token_sender_timestamp(&self, jid: &str, sender_timestamp: i64) -> Result<()> { /* default */ }
    // Store a token received from a contact without clobbering an existing
    // sender_timestamp. Newer-wins: a stale write (older token_timestamp,
    // unless the existing token is a byte-less placeholder) is a no-op.
    // Defaulted (read-modify-write); SqliteStore/InMemoryBackend enforce it atomically.
    async fn store_received_tc_token(&self, jid: &str, token: &[u8], token_timestamp: i64) -> Result<()> { /* default */ }

    // Sent Message Store (retry support, matches WA Web's getMessageTable)
    async fn store_sent_message(&self, chat_jid: &str, message_id: &str, payload: &[u8]) -> Result<()>;
    async fn take_sent_message(&self, chat_jid: &str, message_id: &str) -> Result<Option<Vec<u8>>>;
    async fn delete_expired_sent_messages(&self, cutoff_timestamp: i64) -> Result<u32>;
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LidPnMappingEntry {
    pub lid: String,
    pub phone_number: String,
    pub created_at: i64,
    pub updated_at: i64,
    pub learning_source: String,
}

pub struct LidPnCache {
    lid_to_entry: TypedCache<Arc<str>, Arc<LidPnEntry>>,
    pn_to_entry: TypedCache<Arc<str>, Arc<LidPnEntry>>,
}

// Get LID for a phone number
let lid = cache.get_current_lid("15551234567").await;

// Get phone number for a LID
let phone = cache.get_phone_number("100000012345678").await;

// Add mapping (only updates PN map if newer timestamp)
cache.add(entry).await;

// Warm up cache from persistent storage on client init
let entries = backend.get_all_lid_mappings().await?;
cache.warm_up(entries).await;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeviceListRecord {
    pub user: String,
    pub devices: Vec<DeviceInfo>,
    pub timestamp: i64,
    pub phash: Option<String>,
    /// ADV raw_id from ADVKeyIndexList — used to detect identity changes.
    /// When this changes, all sessions and sender keys for the user must be cleared.
    pub raw_id: Option<u32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeviceInfo {
    pub device_id: u32,
    pub key_index: Option<u32>,
    /// Whether the device belongs to WhatsApp's hosted PN/LID address space.
    #[serde(default)]
    pub is_hosted: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TcTokenEntry {
    pub token: Vec<u8>,
    pub token_timestamp: i64,
    pub sender_timestamp: Option<i64>,
}

#[async_trait]
pub trait DeviceStore: Send + Sync {
    async fn save(&self, device: &Device) -> Result<()>;
    async fn load(&self) -> Result<Option<Device>>;
    async fn exists(&self) -> Result<bool>;
    async fn create(&self) -> Result<i32>;
    async fn snapshot_db(&self, _name: &str, _extra_content: Option<&[u8]>) -> Result<()>;
}

// wacore/src/store/device.rs
#[derive(Clone, Serialize, Deserialize)]
pub struct Device {
    pub pn: Option<Jid>,                    // Phone number JID
    pub lid: Option<Jid>,                   // Long-term identifier
    pub push_name: String,                  // Display name
    pub registration_id: u32,               // Signal registration ID
    pub adv_secret_key: [u8; 32],          // Advertisement secret
    #[serde(with = "key_pair_serde")]
    pub identity_key: KeyPair,              // Signal identity keypair
    #[serde(with = "key_pair_serde")]
    pub noise_key: KeyPair,                 // Noise protocol keypair
    #[serde(with = "key_pair_serde")]
    pub signed_pre_key: KeyPair,            // Signal signed pre-key
    pub signed_pre_key_id: u32,
    #[serde(with = "BigArray")]
    pub signed_pre_key_signature: [u8; 64],
    #[serde(with = "account_serde", default)]
    pub account: Option<Arc<wa::ADVSignedDeviceIdentity>>,
    pub app_version_primary: u32,
    pub app_version_secondary: u32,
    pub app_version_tertiary: u32,
    pub app_version_last_fetched_ms: i64,
    #[serde(skip)]
    pub device_props: wa::DeviceProps,       // Not persisted
    #[serde(skip)]
    pub client_profile: ClientProfile,       // Noise-handshake identity, not persisted
    #[serde(default)]
    pub edge_routing_info: Option<Vec<u8>>,  // Optimized reconnection
    #[serde(default)]
    pub props_hash: Option<String>,          // A/B experiment tracking
    #[serde(default)]
    pub next_pre_key_id: u32,               // Monotonic counter; advances at generation time (WA Web NEXT_PK_ID)
    #[serde(default)]
    pub first_unupload_pre_key_id: u32,     // First generated-but-unuploaded key (WA Web FIRST_UNUPLOAD_PK_ID); 0 = unset/legacy
    #[serde(default)]
    pub server_has_prekeys: bool,            // Whether server has prekeys uploaded
    #[serde(default)]
    pub nct_salt: Option<Vec<u8>>,           // NCT salt for cstoken computation
    #[serde(skip)]
    pub nct_salt_sync_seen: bool,            // Runtime flag: authoritative sync seen
    #[serde(default)]
    pub server_cert_chain: Option<CachedServerCertChain>, // Cached for Noise IK resume
    #[serde(default)]
    pub login_counter: i32,                  // Per-connect counter sent as ClientPayload.lc
    #[serde(default)]
    pub lid_migrated: bool,                  // Account is 1:1-LID-migrated; gates DM wire addressing
    #[serde(default)]
    pub last_signed_pre_key_rotation_ms: i64, // WA Web RotateKeyJob cadence clock; 0 = pre-feature device, baselined (not rotated) on first connect
    #[serde(default)]
    pub server_client_expiration: Option<ServerClientExpiration>, // Server's <ib><client_expiration> deadline for the running build; None until the server says otherwise
}

/// Cached form of the server's two-cert chain. `leaf.key` is the server
/// static public key consumed by Noise IK; the intermediate is kept solely
/// to mirror WA Web's expiry checks.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedServerCertChain {
    pub intermediate: CachedNoiseCert,
    pub leaf: CachedNoiseCert,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct CachedNoiseCert {
    pub key: [u8; 32],     // 32-byte X25519 public key
    pub not_before: i64,   // Validity window (Unix epoch seconds)
    pub not_after: i64,
}

/// The server's answer to "when does this client build stop being accepted",
/// scoped to the build it was issued against.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ServerClientExpiration {
    pub expires_at: i64,           // Unix seconds; server stops accepting this build after this
    pub version: (u32, u32, u32),  // The (primary, secondary, tertiary) build the deadline applies to
}

-- Device table (named columns, not a serialized blob)
CREATE TABLE device (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    lid TEXT NOT NULL,
    pn TEXT NOT NULL,
    registration_id INTEGER NOT NULL,
    noise_key BLOB NOT NULL,
    identity_key BLOB NOT NULL,
    signed_pre_key BLOB NOT NULL,
    signed_pre_key_id INTEGER NOT NULL,
    signed_pre_key_signature BLOB NOT NULL,
    adv_secret_key BLOB NOT NULL,
    account BLOB,
    push_name TEXT NOT NULL DEFAULT '',
    app_version_primary INTEGER NOT NULL DEFAULT 0,
    app_version_secondary INTEGER NOT NULL DEFAULT 0,
    app_version_tertiary BIGINT NOT NULL DEFAULT 0,
    app_version_last_fetched_ms BIGINT NOT NULL DEFAULT 0,
    edge_routing_info BLOB,
    props_hash TEXT,
    next_pre_key_id INTEGER NOT NULL DEFAULT 0,
    first_unupload_pre_key_id INTEGER NOT NULL DEFAULT 0,  -- WA Web FIRST_UNUPLOAD_PK_ID; 0 = unset/legacy
    server_has_prekeys BOOLEAN NOT NULL DEFAULT 0,
    nct_salt BLOB,
    server_cert_chain BLOB,  -- Cached two-cert chain for Noise IK resume
    login_counter INTEGER NOT NULL DEFAULT 0,  -- Per-connect counter; sent as ClientPayload.lc
    lid_migrated BOOLEAN NOT NULL DEFAULT 0,  -- Account is 1:1-LID-migrated; gates DM wire addressing
    last_signed_pre_key_rotation_ms BIGINT NOT NULL DEFAULT 0,  -- WA Web RotateKeyJob cadence clock; 0 = unset/legacy, baselined on first connect
    server_client_expiration TEXT  -- ServerClientExpiration record as JSON; NULL = no deadline held
);

-- Signal protocol tables
CREATE TABLE identities (
    address TEXT NOT NULL,
    key BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (address, device_id)
);

CREATE TABLE sessions (
    address TEXT NOT NULL,
    record BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (address, device_id)
);

CREATE TABLE prekeys (
    id INTEGER NOT NULL,
    key BLOB NOT NULL,
    uploaded BOOLEAN NOT NULL DEFAULT FALSE,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (id, device_id)
);

CREATE TABLE signed_prekeys (
    id INTEGER NOT NULL,
    record BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (id, device_id)
);

CREATE TABLE sender_keys (
    address TEXT NOT NULL,
    record BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (address, device_id)
);

-- App state sync tables
CREATE TABLE app_state_keys (
    key_id BLOB NOT NULL,
    key_data BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (key_id, device_id)
);

CREATE TABLE app_state_versions (
    name TEXT NOT NULL,
    state_data BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (name, device_id)
);

CREATE TABLE app_state_mutation_macs (
    name TEXT NOT NULL,
    version BIGINT NOT NULL,
    index_mac BLOB NOT NULL,
    value_mac BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (name, index_mac, device_id)
);

-- Protocol tables
-- Unified per-device sender key tracking (matches WA Web's participant.senderKey Map)
CREATE TABLE sender_key_devices (
    group_jid TEXT NOT NULL,
    device_jid TEXT NOT NULL,
    has_key INTEGER NOT NULL DEFAULT 0,
    device_id INTEGER NOT NULL DEFAULT 1,
    updated_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
    PRIMARY KEY (group_jid, device_jid, device_id)
);
CREATE INDEX idx_sender_key_devices_device_jid ON sender_key_devices (device_id, device_jid);

CREATE TABLE lid_pn_mapping (
    lid TEXT NOT NULL,
    phone_number TEXT NOT NULL,
    created_at BIGINT NOT NULL,
    learning_source TEXT NOT NULL,
    updated_at BIGINT NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    PRIMARY KEY (lid, device_id)
);

CREATE TABLE base_keys (
    address TEXT NOT NULL,
    message_id TEXT NOT NULL,
    base_key BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    created_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
    PRIMARY KEY (address, message_id, device_id)
);
CREATE INDEX idx_base_keys_created ON base_keys (device_id, created_at);

CREATE TABLE device_registry (
    user_id TEXT NOT NULL,
    devices_json TEXT NOT NULL,
    timestamp INTEGER NOT NULL,
    phash TEXT,
    device_id INTEGER NOT NULL DEFAULT 1,
    updated_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
    PRIMARY KEY (user_id, device_id)
);

CREATE TABLE tc_tokens (
    jid TEXT NOT NULL,
    token BLOB NOT NULL,
    token_timestamp INTEGER NOT NULL,
    sender_timestamp INTEGER,
    device_id INTEGER NOT NULL DEFAULT 1,
    updated_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
    PRIMARY KEY (jid, device_id)
);

-- Sent message store for retry handling
CREATE TABLE sent_messages (
    chat_jid TEXT NOT NULL,
    message_id TEXT NOT NULL,
    payload BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    created_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
    PRIMARY KEY (chat_jid, message_id, device_id)
);

CREATE INDEX idx_sent_messages_created ON sent_messages (created_at, device_id);

use whatsapp_rust::store::SqliteStore;

// Account 1
let backend1 = Arc::new(SqliteStore::new_for_device("whatsapp.db", 1).await?);
let pm1 = PersistenceManager::new(backend1).await?;

// Account 2
let backend2 = Arc::new(SqliteStore::new_for_device("whatsapp.db", 2).await?);
let pm2 = PersistenceManager::new(backend2).await?;

SELECT * FROM sessions WHERE device_id = 1 AND address = ?;

pub struct SqliteStoreConfig {
    /// Max concurrent operations on the write path (`run`): r2d2 `max_size`
    /// AND the internal semaphore permits, kept in lockstep. Clamped to at
    /// least 1. This bounds `run`, not the store as a whole — connections
    /// from `read_pool_size` sit on top of it and aren't counted here.
    ///
    /// Raising this makes *writes* concurrent, which SQLite does not want: two
    /// deferred transactions that both read and then write deadlock on the
    /// upgrade, and `busy_timeout` cannot break it. Leave it at 1 and reach for
    /// `read_pool_size` instead — that is the knob for concurrency, and it is
    /// safe because WAL readers never contend for the write lock.
    pub pool_size: u32,
    /// Extra connections reserved for read-only work, each free to run while a
    /// write holds the write permit. Defaults to `1`. Without a dedicated
    /// reader, your reads queue behind whatever the write permit is doing —
    /// we measured p50 7.2 ms / p99 22.8 ms during a write-behind flush at
    /// `0`, against 0.17 ms / 4.0 ms with one reader. Set it to `0` to put
    /// every operation back on the single queue, exactly as before this knob
    /// existed. This only helps on a WAL connection; on a non-WAL connection
    /// your reads queue on `pool_size` regardless of this setting.
    ///
    /// Each reader costs one connection's page cache (`cache_size_kib`), so
    /// set this to `0` if you're holding many per-session stores that read
    /// rarely.
    pub read_pool_size: u32,
    /// Per-connection SQLite page cache in KiB.
    pub cache_size_kib: u32,
    /// SQLite busy-timeout.
    pub busy_timeout: Duration,
    /// SQLite PRAGMA synchronous level.
    pub synchronous: Synchronous,
    /// Optional shared r2d2 thread pool. When `None`, the store uses the
    /// process-wide lazy pool. Inject a custom pool (e.g. for tests) via this field.
    pub thread_pool: Option<Arc<scheduled_thread_pool::ScheduledThreadPool>>,
    /// Optional hook run first on every new pooled connection, before the store's own
    /// pragmas, WAL setup, and migrations. See [Connection init hook](/api/store#connection-init-hook).
    pub connection_init: Option<ConnectionInitHook>,
}

use whatsapp_rust::store::{SqliteStore, SqliteStoreConfig};

// Low-memory default — same as SqliteStore::new()
let backend = Arc::new(SqliteStore::with_config("whatsapp.db", SqliteStoreConfig::default()).await?);

// Low-memory default, specific device_id — same as SqliteStore::new_for_device()
let backend = Arc::new(
    SqliteStore::with_config_for_device("whatsapp.db", 1, SqliteStoreConfig::default()).await?,
);

// High-throughput profile for a single busy session: a bigger page cache,
// and read concurrency via read_pool_size rather than raising pool_size
// (which would risk the write-upgrade deadlock described below).
let backend = Arc::new(
    SqliteStore::with_config_for_device(
        "whatsapp.db",
        1,
        SqliteStoreConfig {
            cache_size_kib: 4096,
            read_pool_size: 4,
            ..SqliteStoreConfig::default()
        },
    )
    .await?,
);

use whatsapp_rust_sqlite_storage::{SqliteStore, SqliteStoreConfig};

let config = SqliteStoreConfig {
    // A warm cache for a database far larger than the default assumes.
    cache_size_kib: 16 * 1024,
    // Two readers, so a session lookup never waits out a write-behind flush.
    read_pool_size: 2,
    ..Default::default()
}
// Optional: moves reads onto reclaimable file-backed pages.
.with_mmap_size(256 * 1024 * 1024);

let backend = Arc::new(SqliteStore::with_config("whatsapp.db", config).await?);

pub enum DeviceCommand {
    SetId(Option<Jid>),
    SetLid(Option<Jid>),
    SetPushName(String),
    SetPlatform(String),
    SetAccount(Option<wa::ADVSignedDeviceIdentity>),
    /// Update both prekey watermarks atomically. Both must always move
    /// together — split updates are how the pre-watermark model lost track
    /// of generated keys. Replaces the removed `SetNextPreKeyId` variant.
    SetPreKeyWatermarks {
        next_pre_key_id: u32,
        first_unupload_pre_key_id: u32,
    },
    SetNctSalt(Option<Vec<u8>>),
    SetNctSaltFromHistorySync(Vec<u8>),
    // ... more commands
}

pub fn apply_command_to_device(device: &mut Device, command: DeviceCommand) {
    match command {
        DeviceCommand::SetId(id) => {
            device.pn = id;
        }
        DeviceCommand::SetLid(lid) => {
            device.lid = lid;
        }
        DeviceCommand::SetPushName(name) => {
            device.push_name = name;
        }
        DeviceCommand::SetPlatform(platform) => {
            device.platform = platform;
        }
        DeviceCommand::SetAccount(account) => {
            device.account = account.map(std::sync::Arc::new);
        }
        DeviceCommand::SetPreKeyWatermarks { next_pre_key_id, first_unupload_pre_key_id } => {
            device.next_pre_key_id = next_pre_key_id;
            device.first_unupload_pre_key_id = first_unupload_pre_key_id;
        }
        // ...
    }
}

// ✅ Correct: Use DeviceCommand
client.persistence_manager
    .process_command(DeviceCommand::SetPushName("New Name".to_string()))
    .await;

// ❌ Wrong: Direct modification
let mut device = client.device.write().await;
device.push_name = "New Name".to_string(); // DON'T DO THIS

// Cheap snapshot for read-only access — sync, no .await
let snapshot = client.persistence_manager.get_device_snapshot();
println!("JID: {:?}", snapshot.pn);

// Use process_command for type-safe mutations
client.persistence_manager
    .process_command(DeviceCommand::SetPushName(name))
    .await;

// Use modify_device for multiple changes
client.persistence_manager.modify_device(|device| {
    device.push_name = "New Name".to_string();
    device.platform = "Chrome".to_string();
}).await;

// Create snapshot for debugging crypto failures
client.persistence_manager
    .create_snapshot("decrypt_failure", Some(&failed_message_bytes))
    .await?;

use async_trait::async_trait;
use wacore::store::traits::*;

pub struct PostgresStore {
    pool: sqlx::PgPool,
    device_id: i32,
}

#[async_trait]
impl SignalStore for PostgresStore {
    async fn put_identity(&self, address: &str, key: [u8; 32]) -> Result<()> {
        sqlx::query(
            "INSERT INTO identities (device_id, address, key) 
             VALUES ($1, $2, $3) 
             ON CONFLICT (device_id, address) DO UPDATE SET key = $3"
        )
        .bind(self.device_id)
        .bind(address)
        .bind(&key[..])
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    // ... implement other methods
}

#[async_trait]
impl AppSyncStore for PostgresStore {
    // ... implement methods
}

#[async_trait]
impl ProtocolStore for PostgresStore {
    // ... implement methods
}

#[async_trait]
impl DeviceStore for PostgresStore {
    // ... implement methods
}

// Backend trait automatically implemented

let backend = Arc::new(PostgresStore::new(pool, device_id));
let pm = PersistenceManager::new(backend).await?;

[dependencies]
whatsapp-rust = { version = "0.7", features = ["debug-snapshots"] }

// Trigger snapshot on critical errors
if let Err(e) = decrypt_message(&msg).await {
    client.persistence_manager
        .create_snapshot(
            &format!("decrypt_failure_{}", msg.id),
            Some(&serialized_message)
        )
        .await?;
    return Err(e);
}

snapshots/
├── decrypt_failure_1234567890_whatsapp.db
└── decrypt_failure_1234567890_extra.bin

#[async_trait]
pub trait CacheStore: Send + Sync + 'static {
    async fn get(&self, namespace: &str, key: &str) -> anyhow::Result<Option<Vec<u8>>>;
    async fn set(&self, namespace: &str, key: &str, value: &[u8], ttl: Option<Duration>) -> anyhow::Result<()>;
    async fn delete(&self, namespace: &str, key: &str) -> anyhow::Result<()>;
    async fn clear(&self, namespace: &str) -> anyhow::Result<()>;
    async fn entry_count(&self, namespace: &str) -> anyhow::Result<u64> { Ok(0) }
}

// In-process path (zero overhead)
let cache = TypedCache::from_local(portable_cache);

// Custom store path (serde_json serialization)
let cache = TypedCache::from_store(store, "group", Some(Duration::from_secs(3600)));

let cache: TypedCache<Jid, Arc<GroupInfo>> = config.group_cache.build_typed_ttl(
    config.cache_stores.group_cache.clone(),
    "group",
);

pub struct CacheStores {
    pub group_cache: Option<Arc<dyn CacheStore>>,
    pub device_registry_cache: Option<Arc<dyn CacheStore>>,
    pub lid_pn_cache: Option<Arc<dyn CacheStore>>,
}

use whatsapp_rust::{CacheConfig, CacheStores};
use std::sync::Arc;

let redis = Arc::new(MyRedisCacheStore::new("redis://localhost:6379"));

let bot = Bot::builder()
    .with_backend(backend)
    .with_transport_factory(transport)
    .with_http_client(http_client)
    .with_cache_config(CacheConfig {
        cache_stores: CacheStores {
            group_cache: Some(redis.clone()),
            device_registry_cache: Some(redis.clone()),
            ..Default::default()
        },
        ..Default::default()
    })
    .build()
    .await?;

let config = CacheConfig {
    cache_stores: CacheStores::all(redis.clone()),
    ..Default::default()
};

// Pseudocode for the patching pattern
if let Some(mut cached) = cache.get(&key).await {
    cached.apply_change(notification_data);
    cache.insert(key, cached).await;
}
// If the key isn't cached, the patch is a no-op —
// the next read fetches authoritative state from the server.