Storage Traits

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

impl<T> Backend for T
where
    T: SignalStore + AppSyncStore + ProtocolStore + MsgSecretStore + DeviceStore + Send + Sync
{}

/// Store an identity key for a remote address
async fn put_identity(&self, address: &str, key: [u8; 32]) -> Result<()>;

/// Load an identity key for a remote address
async fn load_identity(&self, address: &str) -> Result<Option<Vec<u8>>>;

/// Delete an identity key
async fn delete_identity(&self, address: &str) -> Result<()>;

/// Delete several identity keys in one backend operation (default loops over delete_identity)
async fn delete_identities_batch(&self, addresses: &[Arc<str>]) -> Result<()>;

/// Get an encrypted session for an address
async fn get_session(&self, address: &str) -> Result<Option<Vec<u8>>>;

/// Store an encrypted session
async fn put_session(&self, address: &str, session: &[u8]) -> Result<()>;

/// Delete a session
async fn delete_session(&self, address: &str) -> Result<()>;

/// Delete several sessions in one backend operation (default loops over delete_session)
async fn delete_sessions_batch(&self, addresses: &[Arc<str>]) -> Result<()>;

/// Check if a session exists (default implementation uses get_session)
async fn has_session(&self, address: &str) -> Result<bool>;

/// Store a pre-key
async fn store_prekey(&self, id: u32, record: &[u8], uploaded: bool) -> Result<()>;

/// Store multiple pre-keys in a single batch operation (default loops over store_prekey)
async fn store_prekeys_batch(&self, keys: &[(u32, Vec<u8>)], uploaded: bool) -> Result<()>;

/// Load a pre-key by ID
async fn load_prekey(&self, id: u32) -> Result<Option<Vec<u8>>>;

/// Load multiple pre-keys by ID in a single batch operation (default loops over load_prekey)
async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>>;

/// Remove a pre-key
async fn remove_prekey(&self, id: u32) -> Result<()>;

/// Remove several pre-keys in one backend operation (default loops over remove_prekey)
async fn remove_prekeys_batch(&self, ids: &[u32]) -> Result<()>;

/// Get the highest pre-key ID currently stored
async fn get_max_prekey_id(&self) -> Result<u32>;

/// Store a signed pre-key
async fn store_signed_prekey(&self, id: u32, record: &[u8]) -> Result<()>;

/// Load a signed pre-key by ID
async fn load_signed_prekey(&self, id: u32) -> Result<Option<Vec<u8>>>;

/// Load all signed pre-keys (returns id, record pairs)
async fn load_all_signed_prekeys(&self) -> Result<Vec<(u32, Vec<u8>)>>;

/// Remove a signed pre-key
async fn remove_signed_prekey(&self, id: u32) -> Result<()>;

/// Store a sender key for group messaging
async fn put_sender_key(&self, address: &str, record: &[u8]) -> Result<()>;

/// Get a sender key
async fn get_sender_key(&self, address: &str) -> Result<Option<Vec<u8>>>;

/// Delete a sender key
async fn delete_sender_key(&self, address: &str) -> Result<()>;

/// Delete several sender keys in one backend operation (default loops over delete_sender_key)
async fn delete_sender_keys_batch(&self, addresses: &[Arc<str>]) -> Result<()>;

/// Get an app state sync key by ID
async fn get_sync_key(&self, key_id: &[u8]) -> Result<Option<AppStateSyncKey>>;

/// Set an app state sync key
async fn set_sync_key(&self, key_id: &[u8], key: AppStateSyncKey) -> Result<()>;

/// Get the latest sync key ID
async fn get_latest_sync_key_id(&self) -> Result<Option<Vec<u8>>>;

/// Get the app state version for a collection, or `None` if it has never synced.
async fn get_version(&self, name: &str) -> Result<Option<HashState>>;

/// Forget a collection's version, returning it to the never-synced state.
async fn delete_version(&self, name: &str) -> Result<()>;

/// Set the app state version for a collection
async fn set_version(&self, name: &str, state: HashState) -> Result<()>;

/// Store mutation MACs for a version
async fn put_mutation_macs(
    &self,
    name: &str,
    version: u64,
    mutations: &[AppStateMutationMAC],
) -> Result<()>;

/// Get a mutation MAC by index
async fn get_mutation_mac(&self, name: &str, index_mac: &[u8]) -> Result<Option<Vec<u8>>>;

/// Batch variant of get_mutation_mac: fetch many previous-MAC values in one
/// call, returning index_mac -> value_mac. The default loops over
/// get_mutation_mac; backends with set-membership queries (SQL `IN (...)`)
/// should override to avoid an N+1 (one round-trip per mutation) during sync.
async fn get_mutation_macs(
    &self,
    name: &str,
    index_macs: &[[u8; 32]],
) -> Result<HashMap<[u8; 32], Vec<u8>>>;

/// Delete mutation MACs by their index MACs
async fn delete_mutation_macs(&self, name: &str, index_macs: &[Vec<u8>]) -> Result<()>;

/// Persist one applied patch — its new version, the index MACs it removed,
/// and the MACs it added — as a unit. Use the default if your backend has no
/// transactions: it runs `set_version`, then `delete_mutation_macs`, then
/// `put_mutation_macs`, so your current behavior stays unchanged. Override it
/// with a single transaction otherwise.
async fn commit_patch(
    &self,
    name: &str,
    state: HashState,
    removed_index_macs: &[Vec<u8>],
    added: &[AppStateMutationMAC],
) -> Result<()>;

/// Get sender key distribution status for all known devices in a group.
/// Returns (device_jid_string, has_key) pairs.
async fn get_sender_key_devices(&self, group_jid: &str) -> Result<Vec<(String, bool)>>;

/// Set sender key status for devices. Use has_key=true after successful
/// SKDM distribution (WA Web: markHasSenderKey), or has_key=false to mark
/// devices as needing fresh SKDM (WA Web: markForgetSenderKey).
async fn set_sender_key_status(&self, group_jid: &str, entries: &[(&str, bool)]) -> Result<()>;

/// Clear all sender key device tracking for a group (on sender key rotation).
async fn clear_sender_key_devices(&self, group_jid: &str) -> Result<()>;

/// Get a mapping by LID
async fn get_lid_mapping(&self, lid: &str) -> Result<Option<LidPnMappingEntry>>;

/// Get a mapping by phone number (returns the most recent LID)
async fn get_pn_mapping(&self, phone: &str) -> Result<Option<LidPnMappingEntry>>;

/// Store or update a LID-PN mapping
async fn put_lid_mapping(&self, entry: &LidPnMappingEntry) -> Result<()>;

/// Get all LID-PN mappings (for cache warm-up)
async fn get_all_lid_mappings(&self) -> Result<Vec<LidPnMappingEntry>>;

/// Save the base key for a session address during retry collision detection
async fn save_base_key(&self, address: &str, message_id: &str, base_key: &[u8]) -> Result<()>;

/// Check if the current session has the same base key as the saved one
async fn has_same_base_key(
    &self,
    address: &str,
    message_id: &str,
    current_base_key: &[u8],
) -> Result<bool>;

/// Delete a base key entry
async fn delete_base_key(&self, address: &str, message_id: &str) -> Result<()>;

/// Delete base keys recorded before `cutoff_timestamp` (unix seconds).
/// Returns the count deleted. Defaulted to `Ok(0)` — the keepalive sweep
/// calls it unconditionally for every backend, so a backend that never
/// implements it just never prunes.
async fn delete_expired_base_keys(&self, cutoff_timestamp: i64) -> Result<u32> {
    // default: Ok(0)
}

/// Update the device list for a user (called after usync responses)
async fn update_device_list(&self, record: DeviceListRecord) -> Result<()>;

/// Batched variant — update the device list for many users in one call,
/// used by the parallelized group encrypt fan-out so the device-registry
/// write doesn't serialize per-recipient.
async fn update_device_lists(&self, records: Vec<DeviceListRecord>) -> Result<()>;

/// Get all known devices for a user
async fn get_devices(&self, user: &str) -> Result<Option<DeviceListRecord>>;

/// Batched variant of `get_devices`: the records stored under any of
/// `users`, in no particular order, with absent users left out. Backends
/// should override with one query — the default loops over `get_devices`
/// for correctness.
async fn get_devices_batch(&self, users: &[&str]) -> Result<Vec<DeviceListRecord>>;

/// Get a trusted contact token for a JID (stored under LID)
async fn get_tc_token(&self, jid: &str) -> Result<Option<TcTokenEntry>>;

/// Get the trusted contact tokens for several JIDs at once, in the order
/// asked, `None` where the backend holds no row. The default loops over
/// `get_tc_token`; override with a single `WHERE jid IN (...)` query.
async fn get_tc_tokens(&self, jids: &[String]) -> Result<Vec<Option<TcTokenEntry>>> {
    // default: one get_tc_token call per JID
}

/// Store or update a trusted contact token for a JID
async fn put_tc_token(&self, jid: &str, entry: &TcTokenEntry) -> Result<()>;

/// Delete a trusted contact token for a JID
async fn delete_tc_token(&self, jid: &str) -> Result<()>;

/// Get all JIDs that have stored tc tokens
async fn get_all_tc_token_jids(&self) -> Result<Vec<String>>;

/// Delete tc tokens that have no live state left. A row is removed only when
/// its received token is expired-or-absent (token_timestamp < token_cutoff, or
/// empty) AND its sender bucket is expired-or-absent (sender_timestamp <
/// sender_cutoff, or null) — so recent sender-side rate-limit state is never
/// dropped just because the received token expired. Returns count deleted.
async fn delete_expired_tc_tokens(&self, token_cutoff: i64, sender_cutoff: i64) -> Result<u32>;

/// Advance sender_timestamp for a contact toward `sender_timestamp`, inserting
/// a byte-less placeholder when no entry exists and preserving any existing
/// token bytes. The stored value only ever moves forward (max), so concurrent
/// writers (post-send issuance, history sync) converge regardless of ordering.
/// Must be atomic w.r.t. `put_tc_token`/`store_received_tc_token` — the default
/// is a read-modify-write; override with a single atomic upsert (`SqliteStore`
/// and `InMemoryBackend` both do).
async fn touch_tc_token_sender_timestamp(&self, jid: &str, sender_timestamp: i64) -> Result<()> {
    // default: read-modify-write via get_tc_token + put_tc_token
}

/// Store a token received from a contact, preserving any existing
/// `sender_timestamp`. The symmetric counterpart of
/// `touch_tc_token_sender_timestamp` — each writer owns its own field, so the
/// notification path never drops a sender bucket the issuance path wrote
/// concurrently.
///
/// **Newer-wins**: the stored `(token, token_timestamp)` pair is overwritten
/// only when the existing token is a byte-less placeholder, or the incoming
/// `token_timestamp` is at least as new as the stored one — a stale write must
/// never clobber a fresher real token. `SqliteStore` enforces this atomically
/// inside an `IMMEDIATE` transaction; `InMemoryBackend` under its state lock.
/// This is what lets concurrent history-sync chunks and the
/// privacy-notification path converge on the same row without an external
/// lock. The default read-modify-write below is a best-effort for
/// third-party backends — same atomicity caveat as
/// `touch_tc_token_sender_timestamp`.
async fn store_received_tc_token(&self, jid: &str, token: &[u8], token_timestamp: i64) -> Result<()> {
    // default: read-modify-write via get_tc_token + put_tc_token, newer-wins
}

/// Store a sent message's serialized payload for retry handling.
/// Called after each send_message(); the payload is the protobuf-encoded Message.
async fn store_sent_message(
    &self,
    chat_jid: &str,
    message_id: &str,
    payload: &[u8],
) -> Result<()>;

/// Retrieve and delete a sent message (atomic take). Returns serialized payload.
/// Called when a retry receipt arrives; consuming prevents double-retry.
async fn take_sent_message(
    &self,
    chat_jid: &str,
    message_id: &str,
) -> Result<Option<Vec<u8>>>;

/// Delete sent messages older than cutoff (unix timestamp seconds).
/// Returns count deleted.
async fn delete_expired_sent_messages(
    &self,
    cutoff_timestamp: i64,
) -> Result<u32>;

/// Store a single secret with no expiry (expires_at = 0, kept forever).
async fn put_msg_secret(
    &self,
    chat: &str,
    sender: &str,
    msg_id: &str,
    secret: &[u8; 32], // MessageSecret, the protocol-fixed message-secret size
) -> Result<()>;

/// Batch upsert. Each MsgSecretEntry carries its own absolute expires_at
/// deadline and parent message_ts. On conflict the later deadline wins
/// (0 = never) and the later non-zero message_ts wins. Returns count stored.
async fn put_msg_secrets(&self, entries: Vec<MsgSecretEntry>) -> Result<usize>;

/// Look up a secret for decryption.
async fn get_msg_secret(
    &self,
    chat: &str,
    sender: &str,
    msg_id: &str,
) -> Result<Option<Vec<u8>>>;

/// Look up a secret together with the parent message's event time, used to
/// enforce the per-add-on edit window.
async fn get_msg_secret_with_ts(
    &self,
    chat: &str,
    sender: &str,
    msg_id: &str,
) -> Result<Option<(Vec<u8>, i64)>>;

/// Delete secrets whose expires_at <= cutoff. Rows with expires_at = 0 are
/// kept forever. Returns count removed.
async fn delete_expired_msg_secrets(&self, cutoff_timestamp: i64) -> Result<u32>;

pub type MessageSecret = [u8; 32]; // wacore::reporting_token::MESSAGE_SECRET_SIZE

pub struct MsgSecretEntry {
    pub chat: Arc<str>,     // shared per conversation instead of one String per row
    pub sender: Arc<str>,
    pub msg_id: Arc<str>,
    pub secret: MessageSecret,
    pub expires_at: i64,
    pub message_ts: i64,
}

CREATE TABLE msg_secrets (
    chat TEXT NOT NULL,
    sender TEXT NOT NULL,
    msg_id TEXT NOT NULL,
    secret BLOB NOT NULL,
    device_id INTEGER NOT NULL DEFAULT 1,
    created_at INTEGER NOT NULL DEFAULT (strftime('%s', 'now')),
    expires_at INTEGER NOT NULL DEFAULT 0,  -- absolute unix-seconds deadline (0 = never)
    message_ts INTEGER NOT NULL DEFAULT 0,  -- parent message event time (edit-window check)
    PRIMARY KEY (chat, sender, msg_id, device_id)
);
CREATE INDEX idx_msg_secrets_expires ON msg_secrets (device_id, expires_at);

/// Save device data
async fn save(&self, device: &Device) -> Result<()>;

/// Load device data
async fn load(&self) -> Result<Option<Device>>;

/// Check if a device exists
async fn exists(&self) -> Result<bool>;

/// Create a new device row and return its generated device_id
async fn create(&self) -> Result<i32>;

/// Create a snapshot of the database state
/// Optional: label with name, save extra_content (e.g. failing message)
async fn snapshot_db(&self, name: &str, extra_content: Option<&[u8]>) -> Result<()>;

async fn resource_report(&self) -> wacore::stats::StorageResourceReport {
    wacore::stats::StorageResourceReport::default()
}

async fn maintenance(&self) -> Result<()> {
    Ok(())
}

use whatsapp_rust::store::SqliteStore;

// Basic usage - creates/opens database at path
let store = SqliteStore::new("whatsapp.db").await?;

// With device_id for multi-device support
let store = SqliteStore::new_for_device("whatsapp.db", 1).await?;

// Using sqlite:// URL format
let store = SqliteStore::new("sqlite://path/to/db.sqlite").await?;

pub fn shared(&self) -> SharedSqlite;

pub struct SharedSqlite { /* ... */ }

impl SharedSqlite {
    /// Acquires a semaphore permit and runs `f` on a pooled connection via `spawn_blocking`.
    pub async fn run<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut SqliteConnection) -> Result<T> + Send + 'static,
        T: Send + 'static;

    /// Runs `f` as a read-only, consistent snapshot — a diesel deferred
    /// transaction that never asks for the write lock.
    pub async fn read<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&mut SqliteConnection) -> Result<T> + Send + 'static,
        T: Send + 'static;
}

pub fn share_for_device(&self, device_id: i32) -> Self;

use whatsapp_rust_sqlite_storage::SqliteStore;

let device_1 = SqliteStore::new_for_device("whatsapp.db", 1).await?;
// One pool, two connections (write + read at the defaults), two sessions.
let device_2 = device_1.share_for_device(2);

PRAGMA journal_mode = WAL;
PRAGMA busy_timeout = 30000;
PRAGMA synchronous = NORMAL;
PRAGMA cache_size = 512;
PRAGMA temp_store = memory;
PRAGMA foreign_keys = ON;
PRAGMA journal_size_limit = 33554432; -- 32 MiB, as of whatsapp-rust#1411

use whatsapp_rust_sqlite_storage::{SqliteStore, SqliteStoreConfig};

let config = SqliteStoreConfig::default().with_mmap_size(64 * 1024 * 1024); // 64 MiB
let store = SqliteStore::with_config("whatsapp.db", config).await?;

pub type ConnectionInitHook = Arc<
    dyn Fn(&mut SqliteConnection) -> std::result::Result<(), Box<dyn std::error::Error + Send + Sync>>
        + Send
        + Sync,
>;

use whatsapp_rust_sqlite_storage::{SqliteStore, SqliteStoreConfig};
use diesel::prelude::*;

let config = SqliteStoreConfig::default().with_connection_init(move |conn| {
    diesel::sql_query("PRAGMA key = 'my-passphrase';").execute(conn)?;
    // Verify the key: this fails on a wrongly-keyed database.
    diesel::sql_query("SELECT count(*) FROM sqlite_master;").execute(conn)?;
    Ok(())
});
let store = SqliteStore::with_config("whatsapp.db", config).await?;

use whatsapp_rust_sqlite_storage::SqliteStore;
use whatsapp_rust::Bot;
use std::sync::Arc;

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Create store
    let backend = Arc::new(SqliteStore::new("whatsapp.db").await?);
    
    // Use with Bot builder
    let mut bot = Bot::builder()
        .with_backend(backend.clone())
        .with_transport_factory(transport_factory)
        .with_http_client(http_client)
        .with_runtime(TokioRuntime)
        .on_event(|event, client| async move { /* handle events */ })
        .build()
        .await?;
    
    // Store implements all traits
    // SignalStore
    backend.put_identity("address@s.whatsapp.net", [0u8; 32]).await?;
    let identity = backend.load_identity("address@s.whatsapp.net").await?;
    
    // AppSyncStore
    let version = backend.get_version("regular").await?;
    
    // ProtocolStore
    let devices = backend.get_sender_key_devices("group@g.us").await?;
    
    // DeviceStore
    if backend.exists().await? {
        let device = backend.load().await?;
    }
    
    Ok(())
}

#[async_trait]
pub trait CacheStore: Send + Sync + 'static {
    /// Retrieve a cached value by namespace and key.
    async fn get(&self, namespace: &str, key: &str) -> anyhow::Result<Option<Vec<u8>>>;

    /// Store a value with an optional TTL.
    /// When `ttl` is `None`, the entry persists until explicitly deleted.
    async fn set(
        &self,
        namespace: &str,
        key: &str,
        value: &[u8],
        ttl: Option<Duration>,
    ) -> anyhow::Result<()>;

    /// Delete a single key from the given namespace.
    async fn delete(&self, namespace: &str, key: &str) -> anyhow::Result<()>;

    /// Delete all keys in a namespace.
    async fn clear(&self, namespace: &str) -> anyhow::Result<()>;

    /// Approximate entry count (diagnostics only). Default returns 0.
    async fn entry_count(&self, _namespace: &str) -> anyhow::Result<u64> {
        Ok(0)
    }
}

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};

let redis = Arc::new(MyRedisCacheStore::new("redis://localhost:6379"));
let config = CacheConfig {
    cache_stores: CacheStores::all(redis),
    ..Default::default()
};

use async_trait::async_trait;
use redis::aio::ConnectionManager;
use wacore::store::traits::*;
use wacore::store::error::Result;

pub struct RedisStore {
    client: ConnectionManager,
    device_id: i32,
}

impl RedisStore {
    pub async fn new(redis_url: &str) -> Result<Self> {
        let client = redis::Client::open(redis_url)
            .map_err(|e| StoreError::Connection(Box::new(e)))?;
        let conn = client.get_connection_manager().await
            .map_err(|e| StoreError::Connection(Box::new(e)))?;
        
        Ok(Self {
            client: conn,
            device_id: 1,
        })
    }
}

#[async_trait]
impl SignalStore for RedisStore {
    async fn put_identity(&self, address: &str, key: [u8; 32]) -> Result<()> {
        let mut conn = self.client.clone();
        let key_name = format!("identity:{}:{}", self.device_id, address);
        redis::cmd("SET")
            .arg(key_name)
            .arg(&key[..])
            .query_async(&mut conn)
            .await
            .map_err(|e| StoreError::Database(Box::new(e)))?;
        Ok(())
    }
    
    async fn load_identity(&self, address: &str) -> Result<Option<Vec<u8>>> {
        let mut conn = self.client.clone();
        let key_name = format!("identity:{}:{}", self.device_id, address);
        let result: Option<Vec<u8>> = redis::cmd("GET")
            .arg(key_name)
            .query_async(&mut conn)
            .await
            .map_err(|e| StoreError::Database(Box::new(e)))?;
        Ok(result)
    }
    
    // Implement remaining SignalStore methods...
}

#[async_trait]
impl AppSyncStore for RedisStore {
    // Implement all AppSyncStore methods...
}

#[async_trait]
impl ProtocolStore for RedisStore {
    // Implement all ProtocolStore methods...
}

#[async_trait]
impl DeviceStore for RedisStore {
    // Implement all DeviceStore methods...
}

// Backend is automatically implemented!

pub struct AppStateSyncKey {
    pub key_data: Vec<u8>,
    pub fingerprint: Vec<u8>,
    pub timestamp: i64,
}

pub struct LidPnMappingEntry {
    pub lid: String,              // LID user part
    pub phone_number: String,     // Phone number user part
    pub created_at: i64,          // Unix timestamp
    pub updated_at: i64,          // Unix timestamp
    pub learning_source: String,  // e.g. "usync", "peer_pn_message"
}

pub struct TcTokenEntry {
    pub token: Vec<u8>,                    // Raw token bytes
    pub token_timestamp: i64,              // When token was received
    pub sender_timestamp: Option<i64>,     // When we sent our token
}

pub struct DeviceListRecord {
    pub user: String,               // User part of JID
    pub devices: Vec<DeviceInfo>,   // Known devices
    pub timestamp: i64,             // Last update timestamp
    pub phash: Option<String>,      // Participant hash from usync
    pub raw_id: Option<u32>,        // ADV raw_id for identity change detection
}

pub struct DeviceInfo {
    pub device_id: u32,           // 0 = primary, 1+ = companions
    pub key_index: Option<u32>,   // Key index if known
    pub is_hosted: bool,          // Hosted PN/LID address space (see below)
}
impl DeviceInfo {
    pub const fn new(device_id: u32, key_index: Option<u32>) -> Self
    pub const fn with_hosting(mut self, is_hosted: bool) -> Self
}

pub enum StoreError {
    Io(#[from] std::io::Error),
    Serialization(#[source] Box<dyn std::error::Error + Send + Sync>),
    Validation(String),
    Connection(#[source] Box<dyn std::error::Error + Send + Sync>),
    Database(#[source] Box<dyn std::error::Error + Send + Sync>),
    RetriesExhausted { op: String },
    Migration(#[source] Box<dyn std::error::Error + Send + Sync>),
    InvalidConfig(String),
    DeviceNotFound(i32),
}

pub type Result<T> = std::result::Result<T, StoreError>;