package no.zyon.shade.storage

import no.zyon.shade.crypto.CryptoProvider
import no.zyon.shade.types.IdentityKeyPair
import no.zyon.shade.types.OneTimePreKey
import no.zyon.shade.types.SessionState
import no.zyon.shade.types.SignedPreKey

/**
 * In-memory storage for tests and embedded use.
 * Mirrors MemoryStorage in @shade/crypto-web.
 */
class MemoryStorage(private val crypto: CryptoProvider) : StorageProvider {
    private var identity: IdentityKeyPair? = null
    private var registrationId: Int = 0
    private val signedPreKeys = mutableMapOf<Int, SignedPreKey>()
    private val oneTimePreKeys = mutableMapOf<Int, OneTimePreKey>()
    private val sessions = mutableMapOf<String, SessionState>()
    private val trustedIdentities = mutableMapOf<String, ByteArray>()

    override suspend fun getIdentityKeyPair(): IdentityKeyPair? = identity
    override suspend fun saveIdentityKeyPair(keyPair: IdentityKeyPair) { identity = keyPair }
    override suspend fun getLocalRegistrationId(): Int = registrationId
    override suspend fun saveLocalRegistrationId(id: Int) { registrationId = id }

    override suspend fun getSignedPreKey(keyId: Int): SignedPreKey? = signedPreKeys[keyId]
    override suspend fun saveSignedPreKey(key: SignedPreKey) { signedPreKeys[key.keyId] = key }
    override suspend fun removeSignedPreKey(keyId: Int) { signedPreKeys.remove(keyId) }

    override suspend fun getOneTimePreKey(keyId: Int): OneTimePreKey? = oneTimePreKeys[keyId]
    override suspend fun saveOneTimePreKey(key: OneTimePreKey) { oneTimePreKeys[key.keyId] = key }
    override suspend fun removeOneTimePreKey(keyId: Int) { oneTimePreKeys.remove(keyId) }
    override suspend fun getOneTimePreKeyCount(): Int = oneTimePreKeys.size

    override suspend fun getSession(address: String): SessionState? = sessions[address]
    override suspend fun saveSession(address: String, state: SessionState) { sessions[address] = state }
    override suspend fun removeSession(address: String) { sessions.remove(address) }

    override suspend fun isTrustedIdentity(address: String, identityKey: ByteArray): Boolean {
        val stored = trustedIdentities[address] ?: return true // TOFU
        return crypto.constantTimeEqual(stored, identityKey)
    }

    override suspend fun saveTrustedIdentity(address: String, identityKey: ByteArray) {
        trustedIdentities[address] = identityKey
    }
}
