Shade/packages/shade-sdk/src/shade.ts

import type { ShadeEnvelope, StorageProvider } from '@shade/core';
import {
  ShadeSessionManager,
  ShadeEventEmitter,
  NoSessionError,
} from '@shade/core';
import { SubtleCryptoProvider, MemoryStorage } from '@shade/crypto-web';
import { ShadeFetchTransport } from '@shade/transport';
import { BackgroundTasks, type BackgroundHooks } from './background.js';
import { exportBackup, importBackup, backupToString, backupFromString, type BackupBlob } from './backup.js';
import type { ResolvedConfig } from './config.js';

/**
 * The high-level Shade API.
 *
 * Wraps crypto, storage, session management, transport, and optional
 * observer into a single object. Provides magic auto-establish + auto-
 * publish + auto-replenish behavior.
 */
export class Shade {
  private readonly crypto = new SubtleCryptoProvider();
  private readonly events = new ShadeEventEmitter();
  private storage!: StorageProvider;
  private manager!: ShadeSessionManager;
  private transport!: ShadeFetchTransport;
  private background: BackgroundTasks | null = null;
  private address!: string;
  private initialized = false;

  // Per-address mutex to serialize session establishment under concurrent sends
  private establishing = new Map<string, Promise<void>>();
  // Per-address encrypt queue to serialize ratchet mutations
  private encryptChains = new Map<string, Promise<unknown>>();

  // Message handlers
  private messageHandlers: Array<(from: string, plaintext: string) => void> = [];

  constructor(private readonly config: ResolvedConfig) {}

  /**
   * Initialize the SDK:
   *  1. Resolve storage backend
   *  2. Create session manager + generate identity if needed
   *  3. Create transport
   *  4. Generate initial one-time prekeys
   *  5. Register with prekey server
   *  6. Start background tasks
   */
  async initialize(): Promise<void> {
    if (this.initialized) return;

    // Step 1: Storage
    this.storage = await resolveStorage(this.config.storage);

    // Step 2: Session manager with event bus attached
    this.manager = new ShadeSessionManager(this.crypto, this.storage, {
      events: this.events,
    });
    await this.manager.initialize();

    // Step 3: Address (user-provided or persisted UUID)
    this.address = this.config.address ?? (await resolveAddress(this.storage));

    // Step 4: Transport with our signing key
    const identity = await this.storage.getIdentityKeyPair();
    if (!identity) throw new Error('Identity not available after initialize');
    this.transport = new ShadeFetchTransport({
      baseUrl: this.config.prekeyServer,
      crypto: this.crypto,
      signingPrivateKey: identity.signingPrivateKey,
    });

    // Step 5: Initial prekeys + register
    const otpks = await this.manager.generateOneTimePreKeys(20);
    const bundle = await this.manager.createPreKeyBundle();

    try {
      await this.transport.register(
        this.address,
        this.manager.getPublicIdentity(),
        bundle.signedPreKey,
        otpks,
      );
    } catch (err) {
      console.warn(
        `[Shade] Failed to register with prekey server at ${this.config.prekeyServer}: ${(err as Error).message}. Will retry on next replenish.`,
      );
    }

    // Step 6: Background tasks
    this.background = new BackgroundTasks(
      this.manager,
      this.transport,
      this.address,
      this.config,
    );
    this.background.start();

    this.initialized = true;
  }

  /** Your identity's safety number (12 groups × 5 digits) */
  get fingerprint(): Promise<string> {
    if (!this.initialized) throw new Error('Not initialized');
    return this.manager.getIdentityFingerprint();
  }

  /** Your address on the prekey server */
  get myAddress(): string {
    if (!this.initialized) throw new Error('Not initialized');
    return this.address;
  }

  /** Access the underlying event emitter (for observer integration) */
  getEvents(): ShadeEventEmitter {
    return this.events;
  }

  /** Access the underlying session manager (for advanced usage) */
  getManager(): ShadeSessionManager {
    return this.manager;
  }

  /** Access the underlying transport (for advanced usage) */
  getTransport(): ShadeFetchTransport {
    return this.transport;
  }

  /**
   * Encrypt a message to a peer. Auto-establishes a session if none exists.
   * Returns the ShadeEnvelope ready to send over any transport.
   */
  async send(address: string, plaintext: string): Promise<ShadeEnvelope> {
    if (!this.initialized) throw new Error('Not initialized');

    // Serialize all sends to the same peer: the SessionManager mutates
    // ratchet state in place, and interleaved mutations corrupt it.
    const previous = this.encryptChains.get(address) ?? Promise.resolve();
    const next = previous
      .catch(() => {}) // don't propagate upstream failures to later sends
      .then(async () => {
        try {
          return await this.manager.encrypt(address, plaintext);
        } catch (err) {
          if (!(err instanceof NoSessionError)) throw err;
          await this.ensureSession(address);
          return this.manager.encrypt(address, plaintext);
        }
      });
    this.encryptChains.set(address, next);
    return next as Promise<ShadeEnvelope>;
  }

  /**
   * Decrypt an incoming envelope and notify registered message handlers.
   * Returns the plaintext.
   *
   * The caller provides the `from` address because the envelope itself
   * doesn't authenticate the sender — that's determined by your transport
   * layer (auth header, WebSocket peer, push notification metadata, etc.).
   */
  async receive(from: string, envelope: ShadeEnvelope): Promise<string> {
    if (!this.initialized) throw new Error('Not initialized');
    const plaintext = await this.manager.decrypt(from, envelope);
    for (const handler of this.messageHandlers) {
      try {
        handler(from, plaintext);
      } catch (err) {
        console.error('[Shade] Message handler threw:', err);
      }
    }
    return plaintext;
  }

  /** Register a handler for incoming messages */
  onMessage(handler: (from: string, plaintext: string) => void): () => void {
    this.messageHandlers.push(handler);
    return () => {
      this.messageHandlers = this.messageHandlers.filter((h) => h !== handler);
    };
  }

  /** Get a peer's fingerprint (requires an existing session) */
  async getFingerprintFor(address: string): Promise<string> {
    if (!this.initialized) throw new Error('Not initialized');
    return this.manager.getRemoteFingerprint(address);
  }

  /** Verify a fingerprint matches the pinned identity for an address */
  async verify(address: string, fingerprint: string): Promise<boolean> {
    const remote = await this.getFingerprintFor(address);
    return normalize(remote) === normalize(fingerprint);
  }

  /** Manually rotate the identity (destructive — see docs) */
  async rotate(): Promise<void> {
    if (!this.initialized) throw new Error('Not initialized');

    // Rotate locally first
    const newBundle = await this.manager.rotateIdentity();

    // Rebuild the transport with the new signing key so subsequent
    // signed operations (replenish, delete, register) work
    const identity = await this.storage.getIdentityKeyPair();
    if (!identity) throw new Error('Identity missing after rotate');
    this.transport = new ShadeFetchTransport({
      baseUrl: this.config.prekeyServer,
      crypto: this.crypto,
      signingPrivateKey: identity.signingPrivateKey,
    });

    // Re-upload the new bundle
    await this.transport.register(
      this.address,
      this.manager.getPublicIdentity(),
      newBundle.signedPreKey,
      [],
    );

    // Rebuild background tasks so they use the new transport
    if (this.background) {
      this.background.stop();
      this.background = new BackgroundTasks(
        this.manager,
        this.transport,
        this.address,
        this.config,
      );
      this.background.start();
    }
  }

  /** Manually trigger replenishment (normally background task handles this) */
  async replenish(): Promise<number> {
    if (!this.initialized) throw new Error('Not initialized');
    if (!this.background) return 0;
    return this.background.runReplenish();
  }

  /**
   * Export an encrypted backup blob that can be restored to a new device.
   *
   * @param passphrase User passphrase (minimum 12 characters)
   * @param knownAddresses Peer addresses whose sessions should be included
   */
  async exportBackup(passphrase: string, knownAddresses: string[] = []): Promise<string> {
    if (!this.initialized) throw new Error('Not initialized');
    const blob = await exportBackup(this.crypto, this.storage, passphrase, knownAddresses);
    return backupToString(blob);
  }

  /**
   * Restore state from a backup string. Overwrites existing state.
   * Call this BEFORE initialize() on a fresh device, or after shutdown() + re-init.
   */
  async importBackup(backupString: string, passphrase: string): Promise<void> {
    if (!this.initialized) throw new Error('Not initialized');
    const blob = backupFromString(backupString);
    await importBackup(this.crypto, this.storage, blob, passphrase);
    // Reload identity after restore
    const restored = await this.storage.getIdentityKeyPair();
    if (restored) {
      // Rebuild the manager and transport with the restored identity
      this.manager = new ShadeSessionManager(this.crypto, this.storage, { events: this.events });
      await this.manager.initialize();
      this.transport = new ShadeFetchTransport({
        baseUrl: this.config.prekeyServer,
        crypto: this.crypto,
        signingPrivateKey: restored.signingPrivateKey,
      });
    }
  }

  /** Clean shutdown: stop timers, close storage if it supports it */
  async shutdown(): Promise<void> {
    this.background?.stop();
    // Close storage if it has a close method (SQLite)
    const closable = this.storage as unknown as { close?: () => void | Promise<void> };
    if (typeof closable.close === 'function') {
      await closable.close();
    }
    this.initialized = false;
  }

  // ─── Internals ─────────────────────────────────────────────

  private async ensureSession(address: string): Promise<void> {
    // Deduplicate concurrent establishment requests
    const existing = this.establishing.get(address);
    if (existing) {
      await existing;
      return;
    }

    const promise = (async () => {
      const bundle = await this.transport.fetchBundle(address);
      await this.manager.initSessionFromBundle(address, bundle);
    })();

    this.establishing.set(address, promise);
    try {
      await promise;
    } finally {
      this.establishing.delete(address);
    }
  }
}

// ─── Helpers ─────────────────────────────────────────────────

async function resolveStorage(
  spec: string | StorageProvider | { type: 'postgres'; url: string },
): Promise<StorageProvider> {
  if (typeof spec === 'object' && 'getIdentityKeyPair' in spec) {
    return spec;
  }

  if (spec === 'memory') {
    return new MemoryStorage();
  }

  if (typeof spec === 'string' && spec.startsWith('sqlite:')) {
    const path = spec.slice('sqlite:'.length);
    const { SQLiteStorage } = await import('@shade/storage-sqlite');
    return new SQLiteStorage(path);
  }

  if (typeof spec === 'object' && spec.type === 'postgres') {
    const { PostgresStorage } = await import('@shade/storage-postgres');
    return PostgresStorage.create(spec.url);
  }

  throw new Error(`Unsupported storage spec: ${JSON.stringify(spec)}`);
}

async function resolveAddress(storage: StorageProvider): Promise<string> {
  // Try to load a persisted address, else generate a random one and save it.
  // We reuse the config table by storing a special key.
  // Since StorageProvider doesn't expose a generic key-value, we just use
  // the local registration ID as a deterministic fallback.
  const id = await storage.getLocalRegistrationId();
  return `device:${id}`;
}

function normalize(fp: string): string {
  return fp.replace(/\s+/g, ' ').trim();
}