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 { encodeEnvelope, decodeEnvelope, inspectEnvelopeType } from '@shade/proto';
import { ShadeFetchTransport } from '@shade/transport';
import {
  TransferEngine,
  ShadeTransferHttpTransport,
  type ITransferTransport,
  type IncomingTransfer,
  type TransferHandle,
  type TransferOptions,
  type TransferSummary,
} from '@shade/transfer';
import type { Hono } from 'hono';
import { BackgroundTasks } from './background.js';
import { exportBackup, importBackup, backupToString, backupFromString } from './backup.js';
import type { ResolvedConfig } from './config.js';
import {
  ShadeControlChannel,
  ShadeTransferAuthenticator,
  type ControlEnvelopeTransport,
} from './streams-bridge.js';
import { createFilesNamespace, type FilesNamespace } from '@shade/files';

/**
 * 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 — may be sync or async; receive() awaits each.
  private messageHandlers: Array<
    (from: string, plaintext: string) => void | Promise<void>
  > = [];

  // Stream-transfer engine, lazily constructed on first use.
  private transferEngine: TransferEngine | null = null;
  private controlChannel: ShadeControlChannel | null = null;
  private peerBaseUrlResolver: ((peerAddress: string) => Promise<string>) | null = null;
  private envelopeOutboxes: ControlEnvelopeTransport | null = null;

  // `@shade/files` namespace, lazy + memoized.
  private filesNamespace: FilesNamespace | null = null;

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

  /**
   * `@shade/files` namespace — high-level entry point for E2EE filesystem
   * RPC. Lazily creates the underlying channel + streams bridges on first
   * access; subsequent accesses return the same instance.
   *
   * ```ts
   * const files = shade.files;
   * const stop = await files.serve({ list: ..., write: ..., ... });
   * const fs = await files.client('bob');
   * await fs.list('/');
   * ```
   *
   * Requires `configureTransfers({ resolveBaseUrl })` to be called first
   * (same as `upload`/`onIncomingTransfer`).
   */
  get files(): FilesNamespace {
    if (!this.initialized) throw new Error('Not initialized');
    if (this.filesNamespace !== null) return this.filesNamespace;
    // `@shade/files` only imports `Shade` as a type, so the cyclic ESM
    // import is type-only at the value layer — safe to bind synchronously.
    this.filesNamespace = createFilesNamespace(this);
    return this.filesNamespace;
  }

  /** Internal — exposes the BackgroundTasks for `@shade/files` to wire prune. */
  get background(): BackgroundTasks | null {
    return this._background;
  }

  /** 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 {
        await handler(from, plaintext);
      } catch (err) {
        console.error('[Shade] Message handler threw:', err);
      }
    }
    return plaintext;
  }

  /** Register a handler for incoming messages. Async handlers are awaited. */
  onMessage(
    handler: (from: string, plaintext: string) => void | Promise<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();
    if (this.transferEngine !== null) this.transferEngine.destroy();
    if (this.controlChannel !== null) this.controlChannel.destroy();
    // 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;
  }

  // ─── Stream transfers (v0.2.0) ─────────────────────────────

  /**
   * Configure how stream-transfer chunks reach peers. Provide a resolver
   * that returns the peer's HTTP base URL (e.g. by looking up a
   * `transfer.baseUrl` field in your prekey-bundle metadata or a static
   * directory map). If unset, `upload()` rejects with a clear error.
   *
   * Optionally also override the control-envelope transport (defaults to
   * HTTP POSTs to `<base>/v1/transfer/control`).
   */
  configureTransfers(opts: {
    resolveBaseUrl: (peerAddress: string) => Promise<string>;
    envelopeTransport?: ControlEnvelopeTransport;
  }): void {
    this.peerBaseUrlResolver = opts.resolveBaseUrl;
    this.envelopeOutboxes =
      opts.envelopeTransport ?? new HttpEnvelopeTransport(opts.resolveBaseUrl, this.address);
  }

  /**
   * Deliver a freshly-encrypted ratchet envelope to a peer using the
   * configured envelope transport (HTTP POST to `/v1/transfer/control` by
   * default). Used by `@shade/files` for RPC plaintext delivery.
   */
  async deliverControlEnvelope(peerAddress: string, envelope: ShadeEnvelope): Promise<void> {
    if (this.envelopeOutboxes === null) {
      throw new Error(
        'Call shade.configureTransfers({ resolveBaseUrl }) before deliverControlEnvelope()',
      );
    }
    await this.envelopeOutboxes.send(peerAddress, envelope);
  }

  /**
   * Upload bytes to a peer. Returns a `TransferHandle` that can be paused/
   * aborted and awaited. Requires `configureTransfers` to be called first.
   */
  async upload(opts: TransferOptions): Promise<TransferHandle> {
    return (await this.engine()).upload(opts);
  }

  /**
   * Subscribe to incoming transfers from peers. Handler is invoked when a
   * `stream-init` arrives; the handler MUST call `incoming.accept(...)` to
   * begin receiving (or `incoming.decline(...)` to reject).
   */
  async onIncomingTransfer(
    handler: (incoming: IncomingTransfer) => void | Promise<void>,
  ): Promise<() => void> {
    return (await this.engine()).onIncomingTransfer(handler);
  }

  /**
   * Mount the receiver-side HTTP routes on a Hono app. Mount under any
   * base path: `app.route('/shade', await shade.transferRoute())`.
   *
   * Routes:
   *   POST /v1/transfer/:streamId/chunk  — wire-encoded 0x11 chunks
   *   GET  /v1/transfer/:streamId/state  — resume-state lookup
   *   POST /v1/transfer/control          — wire-encoded 0x02 control envelopes
   *   GET  /v1/transfer/health           — peer reachability probe
   */
  async transferRoute(): Promise<Hono> {
    const engine = await this.engine();
    const { createTransferRoutes, PermissiveAuthenticator } = await import('@shade/transfer');
    const app = await createTransferRoutes(engine, {
      authenticator: PermissiveAuthenticator,
    });
    // Add the control-envelope POST route on top.
    app.post('/v1/transfer/control', async (c) => {
      const senderAddress = c.req.header('X-Shade-Sender-Address');
      if (senderAddress === undefined || senderAddress === '') {
        return c.json({ error: 'missing X-Shade-Sender-Address' }, 400);
      }
      const ab = await c.req.arrayBuffer();
      const bytes = new Uint8Array(ab);
      try {
        await this.acceptTransferEnvelope(senderAddress, bytes);
      } catch (err) {
        return c.json({ error: (err as Error).message }, 400);
      }
      return c.json({ ok: true });
    });
    return app;
  }

  /**
   * Low-level entry for custom transports: hand a `0x02` ratchet envelope
   * (control-plane) or a `0x11` stream-chunk envelope to the engine.
   * Used internally by `transferRoute()`.
   */
  async acceptTransferEnvelope(from: string, env: ShadeEnvelope | Uint8Array): Promise<void> {
    const engine = await this.engine();
    if (env instanceof Uint8Array) {
      const kind = inspectEnvelopeType(env);
      if (kind === 'stream-chunk') {
        // Engine extracts laneId/seq from the wire bytes via decodeStreamChunk.
        const headers = parseChunkHeader(env);
        await engine.receiveChunk(from, headers.streamId, headers.laneId, headers.seq, env);
        return;
      }
      if (kind === 'ratchet' || kind === 'prekey') {
        const decoded = decodeEnvelope(env);
        await this.controlChannel!.acceptEnvelope(from, decoded);
        return;
      }
      throw new Error(`Unknown envelope type ${kind}`);
    }
    // Already-decoded envelope (ratchet or prekey)
    await this.controlChannel!.acceptEnvelope(from, env);
  }

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

  private async engine(): Promise<TransferEngine> {
    if (this.transferEngine !== null) return this.transferEngine;
    if (this.peerBaseUrlResolver === null || this.envelopeOutboxes === null) {
      throw new Error(
        'Call shade.configureTransfers({ resolveBaseUrl }) before using upload()/onIncomingTransfer()',
      );
    }
    this.controlChannel = new ShadeControlChannel(this, this.envelopeOutboxes);
    const transport: ITransferTransport = new ShadeTransferHttpTransport({
      resolveBaseUrl: this.peerBaseUrlResolver,
      authenticator: await this.makeAuthenticator(),
    });
    this.transferEngine = new TransferEngine({
      crypto: this.crypto,
      controlChannel: this.controlChannel,
      transport,
      myAddress: this.address,
    });
    return this.transferEngine;
  }

  private async makeAuthenticator(): Promise<ShadeTransferAuthenticator> {
    const identity = await this.storage.getIdentityKeyPair();
    if (identity === null) throw new Error('Identity not initialized');
    return new ShadeTransferAuthenticator(this.crypto, this.address, identity.signingPrivateKey);
  }

  /** Returns a list of in-flight stream transfers from storage (resume support). */
  async listTransfers(filter?: {
    direction?: 'send' | 'receive';
  }): Promise<TransferSummary[]> {
    if (this.storage.listActiveStreamStates === undefined) return [];
    const rows = await this.storage.listActiveStreamStates(filter?.direction);
    return rows.map((s) => ({
      streamId: s.streamId,
      direction: s.direction,
      peerAddress: s.peerAddress,
      status: s.status,
      bytesProcessed: 0, // computed from laneState
      createdAt: s.createdAt,
      updatedAt: s.updatedAt,
      metadata: tryParseMetadata(s.metadataJson),
    }));
  }

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

function tryParseMetadata(json: string): import('@shade/streams').StreamMetadata | null {
  try {
    return JSON.parse(json) as import('@shade/streams').StreamMetadata;
  } catch {
    return null;
  }
}

function parseChunkHeader(bytes: Uint8Array): {
  streamId: string;
  laneId: number;
  seq: bigint;
} {
  // [0]=ver [1]=type [2..18]=streamId(16) [18..22]=laneId u32 [22..30]=seq u64
  if (bytes.length < 30) throw new Error('truncated stream-chunk header');
  const view = new DataView(bytes.buffer, bytes.byteOffset);
  const sidBytes = bytes.slice(2, 18);
  const laneId = view.getUint32(18, false);
  const seq = view.getBigUint64(22, false);
  // Encode streamId as base64url
  let bin = '';
  for (let i = 0; i < sidBytes.length; i++) bin += String.fromCharCode(sidBytes[i]!);
  const streamId = btoa(bin).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');
  return { streamId, laneId, seq };
}

// ─── Default HTTP envelope transport ──────────────────────────

class HttpEnvelopeTransport implements ControlEnvelopeTransport {
  constructor(
    private readonly resolveBaseUrl: (peerAddress: string) => Promise<string>,
    private readonly myAddress: string,
  ) {}
  async send(peerAddress: string, envelope: ShadeEnvelope): Promise<void> {
    const base = (await this.resolveBaseUrl(peerAddress)).replace(/\/$/, '');
    const url = `${base}/v1/transfer/control`;
    const bytes = encodeEnvelope(envelope);
    const res = await fetch(url, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/octet-stream',
        'X-Shade-Sender-Address': this.myAddress,
      },
      body: bytes as unknown as never,
    });
    if (!res.ok) {
      throw new Error(`control envelope POST failed: ${res.status} ${await res.text()}`);
    }
  }
}

// ─── 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') {
    // Dynamic import keeps @shade/storage-postgres optional — consumers that
    // never use postgres don't need to install it. The string-form import
    // path makes the resolver lazy at type-check time too.
    const moduleId = '@shade/storage-postgres';
    const mod = (await import(moduleId)) as {
      PostgresStorage: { create(url: string): Promise<StorageProvider> };
    };
    return mod.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();
}