Shade/packages/shade-core/src/session.ts

import type { CryptoProvider } from './crypto.js';
import type { StorageProvider } from './storage.js';
import type {
  IdentityKeyPair,
  SignedPreKey,
  OneTimePreKey,
  PreKeyBundle,
  PreKeyMessage,
  RatchetMessage,
  ShadeEnvelope,
} from './types.js';
import {
  generateIdentityKeyPair,
  generateSignedPreKey,
  generateOneTimePreKeys,
  createPreKeyBundle,
  processPreKeyBundle,
  processPreKeyMessage,
} from './x3dh.js';
import {
  initSenderSession,
  initReceiverSession,
  ratchetEncrypt,
  ratchetDecrypt,
} from './ratchet.js';
import { NoSessionError } from './errors.js';
import { computeFingerprint, shortFingerprint } from './fingerprint.js';
import { ShadeEventEmitter, shortHash } from './events.js';
import {
  ATTR_ERROR_CODE,
  ATTR_OP,
  ATTR_PEER_HASH,
  ATTR_RESULT,
  NOOP_HOOK,
  peerHash,
  type ObservabilityHook,
  type Span,
} from '@shade/observability';

const enc = new TextEncoder();
const dec = new TextDecoder();

/** Default grace period for retired identities: 7 days */
export const GRACE_PERIOD_MS = 7 * 24 * 60 * 60 * 1000;

/**
 * ShadeSessionManager — the high-level API for using Shade.
 *
 * Wraps X3DH key agreement and Double Ratchet into a simple interface:
 *   - `initialize()` — generate or load identity keys
 *   - `createPreKeyBundle()` — publish to prekey server
 *   - `encrypt(address, plaintext)` — encrypt for a peer
 *   - `decrypt(address, envelope)` — decrypt from a peer
 *
 * Usage:
 * ```ts
 * const manager = new ShadeSessionManager(crypto, storage);
 * await manager.initialize();
 *
 * // To initiate: fetch bundle, then encrypt
 * await manager.initSessionFromBundle('bob', bundle);
 * const envelope = await manager.encrypt('bob', 'Hello!');
 *
 * // To receive: decrypt handles everything
 * const plaintext = await manager.decrypt('alice', envelope);
 * ```
 */
export class ShadeSessionManager {
  private identity: IdentityKeyPair | null = null;
  private registrationId: number = 0;
  private currentSignedPreKeyId: number = 0;
  private readonly events?: ShadeEventEmitter;
  private readonly observability: ObservabilityHook;
  /**
   * Per-address operation chain. Both encrypt and decrypt mutate ratchet
   * state in place (counter, DH key, skipped-keys cache); concurrent
   * operations on the same peer can corrupt the session. We serialize
   * per-peer by chaining promises — operations to different peers stay
   * fully concurrent.
   */
  private readonly peerOpChains = new Map<string, Promise<unknown>>();

  constructor(
    private readonly crypto: CryptoProvider,
    private readonly storage: StorageProvider,
    options: { events?: ShadeEventEmitter; observability?: ObservabilityHook } = {},
  ) {
    if (options.events !== undefined) {
      this.events = options.events;
    }
    this.observability = options.observability ?? NOOP_HOOK;
  }

  /**
   * Wrap a per-peer crypto op in a PII-safe span. The span captures the
   * mutex-acquire latency separately from the inner crypto work so a
   * "ratchet contention" pathology shows up clearly in traces.
   */
  private async withSpan<T>(
    op: 'encrypt' | 'decrypt',
    address: string,
    fn: () => Promise<T>,
  ): Promise<T> {
    const span: Span = this.observability.startSpan(`shade.session.${op}`, {
      [ATTR_OP]: op,
      [ATTR_PEER_HASH]: peerHash(address),
    });
    const lockStart = nowMs();
    try {
      return await this.runUnderPeerLock(address, async () => {
        span.setAttribute('shade.lock.wait_ms', Math.round(nowMs() - lockStart));
        const result = await fn();
        span.setAttribute(ATTR_RESULT, 'ok');
        span.setStatus('ok');
        return result;
      });
    } catch (err) {
      span.setAttribute(ATTR_RESULT, 'error');
      span.setAttribute(ATTR_ERROR_CODE, sessionErrorCodeOf(err));
      span.recordException(err);
      span.setStatus('error');
      throw err;
    } finally {
      span.end();
    }
  }

  /**
   * Run `fn` under the per-address mutex so encrypt/decrypt for the same
   * peer never interleave their session-state mutations.
   */
  private async runUnderPeerLock<T>(address: string, fn: () => Promise<T>): Promise<T> {
    const previous = this.peerOpChains.get(address) ?? Promise.resolve();
    const next = previous.catch(() => undefined).then(fn);
    this.peerOpChains.set(address, next);
    try {
      return await next;
    } finally {
      // Best-effort cleanup so finished chains can be garbage collected
      // when a peer goes idle. If a newer op has chained on, we leave it.
      if (this.peerOpChains.get(address) === next) {
        this.peerOpChains.delete(address);
      }
    }
  }

  /** Get the event emitter (if observability is enabled) */
  getEvents(): ShadeEventEmitter | undefined {
    return this.events;
  }

  // ─── Initialization ────────────────────────────────────────

  /** Initialize: load or generate identity keys and a signed prekey */
  async initialize(): Promise<void> {
    // Load or generate identity
    this.identity = await this.storage.getIdentityKeyPair();
    if (!this.identity) {
      this.identity = await generateIdentityKeyPair(this.crypto);
      await this.storage.saveIdentityKeyPair(this.identity);
    }

    // Load or generate registration ID (cryptographically secure)
    this.registrationId = await this.storage.getLocalRegistrationId();
    if (this.registrationId === 0) {
      // Ensure nonzero (0 is the "unset" sentinel)
      let id = this.crypto.randomUint32();
      if (id === 0) id = 1;
      this.registrationId = id;
      await this.storage.saveLocalRegistrationId(this.registrationId);
    }

    // Generate initial signed prekey if none exists
    const spk = await this.storage.getSignedPreKey(1);
    if (!spk) {
      const signedPreKey = await generateSignedPreKey(this.crypto, this.identity, 1);
      await this.storage.saveSignedPreKey(signedPreKey);
      this.currentSignedPreKeyId = 1;
    } else {
      this.currentSignedPreKeyId = spk.keyId;
    }

    // Emit identity initialization event
    if (this.events) {
      const fingerprint = await this.getIdentityFingerprint();
      this.events.emit('identity.initialized', {
        fingerprint,
        registrationId: this.registrationId,
      });
    }
  }

  /** Get our identity's DH public key (for addressing) */
  getPublicIdentity(): { signingKey: Uint8Array; dhKey: Uint8Array } {
    if (!this.identity) throw new Error('Not initialized');
    return {
      signingKey: this.identity.signingPublicKey,
      dhKey: this.identity.dhPublicKey,
    };
  }

  // ─── Fingerprints (Safety Numbers) ──────────────────────────

  /**
   * Get a human-readable fingerprint of our own identity.
   * Format: 12 groups of 5 decimal digits (60 total).
   * Two parties compare these out-of-band to verify no MITM.
   */
  async getIdentityFingerprint(): Promise<string> {
    if (!this.identity) throw new Error('Not initialized');
    return computeFingerprint(
      this.crypto,
      this.identity.signingPublicKey,
      this.identity.dhPublicKey,
    );
  }

  /** Short 4-group fingerprint for quick comparison */
  async getShortFingerprint(): Promise<string> {
    return shortFingerprint(await this.getIdentityFingerprint());
  }

  /**
   * Get a fingerprint for a remote peer's identity.
   * Throws NoSessionError if we haven't established a session with them.
   */
  async getRemoteFingerprint(address: string): Promise<string> {
    const session = await this.storage.getSession(address);
    if (!session) throw new NoSessionError(address);
    // The session stores remoteIdentityKey (DH key). We need the signing key too,
    // which we don't store per-session. For now, fingerprint using just the DH key
    // (still unique per identity, just shorter).
    // In the future, store remoteIdentitySigningKey alongside.
    return computeFingerprint(
      this.crypto,
      session.remoteIdentityKey,
      session.remoteIdentityKey,
    );
  }

  // ─── Prekey Stock Management ────────────────────────────────

  /**
   * Ensure the one-time prekey stock is above a minimum threshold.
   * If below `min`, generates enough to bring it up to `target`.
   * Returns the number of new keys generated (0 if no action needed).
   */
  async ensurePreKeyStock(min = 5, target = 20): Promise<number> {
    const current = await this.storage.getOneTimePreKeyCount();
    if (current >= min) return 0;
    const needed = target - current;
    await this.generateOneTimePreKeys(needed);
    return needed;
  }

  // ─── Session Reset / Identity Change ────────────────────────

  /**
   * Delete the session for a peer. The next message will trigger a fresh X3DH.
   * Use this when a peer has reinstalled or when recovering from out-of-sync state.
   */
  async resetSession(address: string): Promise<void> {
    await this.storage.removeSession(address);
    this.events?.emit('session.removed', { address });
    // Note: we keep the trusted identity; new session will verify against it.
  }

  /**
   * Accept a changed remote identity. This should only be called after
   * verifying the new identity out-of-band (e.g., comparing fingerprints).
   * After this, any pinned trust for this address is replaced.
   */
  async acceptIdentityChange(address: string, newIdentityKey: Uint8Array): Promise<void> {
    // Capture old hash for the trust.changed event (TOFU semantics make this messy
    // because isTrustedIdentity() compares not retrieves; we just emit the new hash)
    await this.storage.saveTrustedIdentity(address, newIdentityKey);
    await this.storage.removeSession(address);

    if (this.events) {
      const newHash = await shortHash(this.crypto, newIdentityKey);
      this.events.emit('trust.changed', { address, oldKeyHash: '?', newKeyHash: newHash });
      this.events.emit('session.removed', { address });
    }
  }

  /**
   * Check whether a remote identity key matches what we have pinned for an address.
   * Returns true on TOFU (no pinned key yet) or exact match.
   */
  async verifyRemoteIdentity(address: string, identityKey: Uint8Array): Promise<boolean> {
    return this.storage.isTrustedIdentity(address, identityKey);
  }

  // ─── Prekey Management ─────────────────────────────────────

  /** Create a prekey bundle to publish to the prekey server */
  async createPreKeyBundle(): Promise<PreKeyBundle> {
    if (!this.identity) throw new Error('Not initialized');
    const spk = await this.storage.getSignedPreKey(this.currentSignedPreKeyId);
    if (!spk) throw new Error('No signed prekey');

    // Try to include a one-time prekey
    // (In real usage, the prekey server would pick one — here we just check if any exist)
    return createPreKeyBundle(this.registrationId, this.identity, spk);
  }

  /** Generate and store a batch of one-time prekeys */
  async generateOneTimePreKeys(count: number): Promise<OneTimePreKey[]> {
    const existingCount = await this.storage.getOneTimePreKeyCount();
    const startId = existingCount + 1;
    const keys = await generateOneTimePreKeys(this.crypto, startId, count);
    for (const key of keys) {
      await this.storage.saveOneTimePreKey(key);
    }
    this.events?.emit('prekey.generated', {
      count,
      totalAfter: existingCount + count,
    });
    return keys;
  }

  /** Rotate the signed prekey (recommended: every 1-7 days) */
  async rotateSignedPreKey(): Promise<SignedPreKey> {
    if (!this.identity) throw new Error('Not initialized');
    const oldId = this.currentSignedPreKeyId;
    const newId = oldId + 1;
    const spk = await generateSignedPreKey(this.crypto, this.identity, newId);
    await this.storage.saveSignedPreKey(spk);
    // Keep old one for a grace period (sessions may still reference it)
    this.currentSignedPreKeyId = newId;
    this.events?.emit('signed_prekey.rotated', { oldKeyId: oldId, newKeyId: newId });
    return spk;
  }

  // ─── Identity Rotation (with grace period) ─────────────────

  /**
   * Rotate the identity keypair.
   *
   * Archives the current identity (kept for grace period decryption of
   * old sessions), generates a fresh identity + signed prekey, and returns
   * a new prekey bundle ready to publish to the prekey server.
   *
   * Callers should:
   *   1. Call this method
   *   2. Re-publish the new bundle via ShadeFetchTransport.register()
   *   3. Optionally broadcast the rotation to known peers out-of-band
   *
   * The old identity is retained for GRACE_PERIOD_MS so existing sessions
   * continue decrypting. Call `pruneExpiredIdentities()` periodically.
   */
  async rotateIdentity(): Promise<PreKeyBundle> {
    if (!this.identity) throw new Error('Not initialized');

    // Archive current identity
    await this.storage.addRetiredIdentity({
      keyPair: this.identity,
      retiredAt: Date.now(),
    });

    // Generate new identity + save
    this.identity = await generateIdentityKeyPair(this.crypto);
    await this.storage.saveIdentityKeyPair(this.identity);

    // Generate new signed prekey (under the new identity)
    const newSpkId = this.currentSignedPreKeyId + 1;
    const spk = await generateSignedPreKey(this.crypto, this.identity, newSpkId);
    await this.storage.saveSignedPreKey(spk);
    this.currentSignedPreKeyId = newSpkId;

    if (this.events) {
      const newFingerprint = await this.getIdentityFingerprint();
      this.events.emit('identity.rotated', { newFingerprint });
    }

    // Return a fresh bundle for re-publication
    return createPreKeyBundle(this.registrationId, this.identity, spk);
  }

  /**
   * Get all retired identities that are still within the grace period.
   * Used internally for trying previous identities when X3DH fails.
   */
  async getActiveRetiredIdentities(gracePeriodMs = GRACE_PERIOD_MS): Promise<IdentityKeyPair[]> {
    const all = await this.storage.getRetiredIdentities();
    const cutoff = Date.now() - gracePeriodMs;
    return all.filter((r) => r.retiredAt >= cutoff).map((r) => r.keyPair);
  }

  /**
   * Delete retired identities older than the grace period.
   * Call this periodically (e.g., daily cleanup task).
   */
  async pruneExpiredIdentities(gracePeriodMs = GRACE_PERIOD_MS): Promise<void> {
    const cutoff = Date.now() - gracePeriodMs;
    await this.storage.pruneRetiredIdentities(cutoff);
  }

  // ─── Session Establishment ─────────────────────────────────

  /**
   * Initiate a session with a peer by processing their prekey bundle.
   * Call this before the first `encrypt()` to a new peer.
   */
  async initSessionFromBundle(address: string, bundle: PreKeyBundle): Promise<void> {
    const x3dhResult = await processPreKeyBundle(this.crypto, this.storage, bundle);

    const session = await initSenderSession(
      this.crypto,
      x3dhResult.rootKey,
      x3dhResult.remoteIdentityKey,
      x3dhResult.remoteSignedPreKey,
    );

    await this.storage.saveSession(address, session);
    await this.storage.saveTrustedIdentity(address, x3dhResult.remoteIdentityKey);

    // Store X3DH metadata for the first message
    // We stash this on the session object for the first encrypt call
    (session as any).__x3dh = {
      ephemeralPublicKey: x3dhResult.ephemeralPublicKey,
      signedPreKeyId: x3dhResult.signedPreKeyId,
      preKeyId: x3dhResult.preKeyId,
      identityDHKey: this.identity!.dhPublicKey,
      registrationId: this.registrationId,
    };
    await this.storage.saveSession(address, session);

    if (this.events) {
      const remoteHash = await shortHash(this.crypto, x3dhResult.remoteIdentityKey);
      this.events.emit('session.created', { address, remoteIdentityKeyHash: remoteHash });
      this.events.emit('trust.pinned', { address, identityKeyHash: remoteHash });
    }
  }

  // ─── Encrypt / Decrypt ─────────────────────────────────────

  /**
   * Encrypt a message for a peer. Returns a ShadeEnvelope ready to send.
   *
   * The first message to a new peer will be a PreKeyMessage (includes X3DH info).
   * Subsequent messages are standard RatchetMessages.
   */
  async encrypt(address: string, plaintext: string): Promise<ShadeEnvelope> {
    return this.withSpan('encrypt', address, async () => {
      const session = await this.storage.getSession(address);
      if (!session) throw new NoSessionError(address);

      const ratchetMsg = await ratchetEncrypt(this.crypto, session, enc.encode(plaintext));

      this.events?.emit('message.encrypted', {
        address,
        counter: ratchetMsg.counter,
        ciphertextSize: ratchetMsg.ciphertext.length,
      });

      // Check if this is the first message (X3DH metadata attached)
      const x3dh = (session as any).__x3dh;
      if (x3dh) {
        delete (session as any).__x3dh;
        await this.storage.saveSession(address, session);

        const preKeyMsg: PreKeyMessage = {
          registrationId: x3dh.registrationId,
          preKeyId: x3dh.preKeyId,
          signedPreKeyId: x3dh.signedPreKeyId,
          ephemeralKey: x3dh.ephemeralPublicKey,
          identityDHKey: x3dh.identityDHKey,
          message: ratchetMsg,
        };
        return {
          type: 'prekey',
          content: preKeyMsg,
          timestamp: Date.now(),
          senderAddress: address,
        };
      }

      await this.storage.saveSession(address, session);
      return {
        type: 'ratchet',
        content: ratchetMsg,
        timestamp: Date.now(),
        senderAddress: address,
      };
    });
  }

  /**
   * Decrypt a message from a peer. Handles both PreKeyMessage and RatchetMessage.
   */
  async decrypt(address: string, envelope: ShadeEnvelope): Promise<string> {
    return this.withSpan('decrypt', address, async () => {
      if (envelope.type === 'prekey') {
        return this.decryptPreKeyMessage(address, envelope.content as PreKeyMessage);
      }
      return this.decryptRatchetMessage(address, envelope.content as RatchetMessage);
    });
  }

  private async decryptPreKeyMessage(address: string, message: PreKeyMessage): Promise<string> {
    // Process X3DH to establish session
    const x3dhResult = await processPreKeyMessage(this.crypto, this.storage, message);

    // Find the signed prekey that was used
    const spk = await this.storage.getSignedPreKey(message.signedPreKeyId);
    if (!spk) throw new Error(`Signed prekey ${message.signedPreKeyId} not found`);

    const session = initReceiverSession(
      x3dhResult.rootKey,
      x3dhResult.remoteIdentityKey,
      spk.keyPair,
    );

    // Decrypt the embedded first ratchet message
    const plaintext = await ratchetDecrypt(this.crypto, session, x3dhResult.initialMessage);

    await this.storage.saveSession(address, session);
    await this.storage.saveTrustedIdentity(address, x3dhResult.remoteIdentityKey);

    if (this.events) {
      const remoteHash = await shortHash(this.crypto, x3dhResult.remoteIdentityKey);
      this.events.emit('session.created', { address, remoteIdentityKeyHash: remoteHash });
      this.events.emit('trust.pinned', { address, identityKeyHash: remoteHash });
      if (message.preKeyId != null) {
        this.events.emit('prekey.consumed', { keyId: message.preKeyId });
      }
      this.events.emit('message.decrypted', {
        address,
        counter: x3dhResult.initialMessage.counter,
        plaintextSize: plaintext.length,
      });
    }

    return dec.decode(plaintext);
  }

  private async decryptRatchetMessage(address: string, message: RatchetMessage): Promise<string> {
    const session = await this.storage.getSession(address);
    if (!session) throw new NoSessionError(address);

    // Detect DH ratchet step (new remote DH key)
    const willRatchet = !session.dhReceive ||
      !arraysEqual(message.dhPublicKey, session.dhReceive);

    const plaintext = await ratchetDecrypt(this.crypto, session, message);
    await this.storage.saveSession(address, session);

    if (this.events) {
      if (willRatchet) {
        this.events.emit('ratchet.dh_step', { address });
      }
      this.events.emit('message.decrypted', {
        address,
        counter: message.counter,
        plaintextSize: plaintext.length,
      });
    }

    return dec.decode(plaintext);
  }
}

function arraysEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  for (let i = 0; i < a.length; i++) {
    if (a[i] !== b[i]) return false;
  }
  return true;
}

function nowMs(): number {
  return typeof performance !== 'undefined' ? performance.now() : Date.now();
}

function sessionErrorCodeOf(err: unknown): string {
  if (err === null || err === undefined) return 'SHADE_UNKNOWN';
  if (typeof err === 'object') {
    const code = (err as { code?: unknown }).code;
    if (typeof code === 'string' && code.length > 0) return code;
    const name = (err as { name?: unknown }).name;
    if (typeof name === 'string' && name.length > 0) return name;
  }
  return 'SHADE_UNKNOWN';
}