import type { PrekeyStore } from './store.js';

interface IdentityRecord {
  identitySigningKey: Uint8Array;
  identityDHKey: Uint8Array;
}

interface SignedPreKeyRecord {
  keyId: number;
  publicKey: Uint8Array;
  signature: Uint8Array;
}

interface OneTimePreKeyRecord {
  keyId: number;
  publicKey: Uint8Array;
}

/**
 * In-memory PrekeyStore for testing and embedded use.
 */
export class MemoryPrekeyStore implements PrekeyStore {
  private identities = new Map<string, IdentityRecord>();
  private signedPreKeys = new Map<string, SignedPreKeyRecord>();
  private oneTimePreKeys = new Map<string, OneTimePreKeyRecord[]>();
  private lastActivity = new Map<string, number>();

  async saveIdentity(address: string, identitySigningKey: Uint8Array, identityDHKey: Uint8Array): Promise<void> {
    this.identities.set(address, { identitySigningKey, identityDHKey });
  }

  async getIdentity(address: string): Promise<IdentityRecord | null> {
    return this.identities.get(address) ?? null;
  }

  async saveSignedPreKey(address: string, keyId: number, publicKey: Uint8Array, signature: Uint8Array): Promise<void> {
    this.signedPreKeys.set(address, { keyId, publicKey, signature });
  }

  async getSignedPreKey(address: string): Promise<SignedPreKeyRecord | null> {
    return this.signedPreKeys.get(address) ?? null;
  }

  async saveOneTimePreKeys(address: string, keys: Array<{ keyId: number; publicKey: Uint8Array }>): Promise<void> {
    const existing = this.oneTimePreKeys.get(address) ?? [];
    existing.push(...keys);
    this.oneTimePreKeys.set(address, existing);
  }

  async consumeOneTimePreKey(address: string): Promise<OneTimePreKeyRecord | null> {
    const keys = this.oneTimePreKeys.get(address);
    if (!keys || keys.length === 0) return null;
    return keys.shift()!;
  }

  async getOneTimePreKeyCount(address: string): Promise<number> {
    return this.oneTimePreKeys.get(address)?.length ?? 0;
  }

  async deleteAll(address: string): Promise<void> {
    this.identities.delete(address);
    this.signedPreKeys.delete(address);
    this.oneTimePreKeys.delete(address);
    this.lastActivity.delete(address);
  }

  // ─── Stale cleanup ────────────────────────────────────

  async touchIdentity(address: string): Promise<void> {
    this.lastActivity.set(address, Date.now());
  }

  async purgeStaleIdentities(olderThanMs: number): Promise<number> {
    const cutoff = Date.now() - olderThanMs;
    const stale: string[] = [];
    for (const [address, ts] of this.lastActivity) {
      if (ts < cutoff) stale.push(address);
    }
    // Also purge identities that have a row but were never touched
    // (last_activity = 0/undefined before first touch)
    for (const address of this.identities.keys()) {
      if (!this.lastActivity.has(address)) stale.push(address);
    }
    for (const address of stale) {
      this.identities.delete(address);
      this.signedPreKeys.delete(address);
      this.oneTimePreKeys.delete(address);
      this.lastActivity.delete(address);
    }
    return stale.length;
  }
}
