/**
 * Layer 4: Sync (Server) — Multi-document WebSocket sync relay.
 *
 * Acts as a sync peer for multiple Automerge documents. Can operate in two modes:
 *
 * 1. **Relay mode** (encrypted docs): Forwards encrypted sync messages between peers.
 *    Server can't read document content — just routes by docId.
 *
 * 2. **Participant mode** (shared/unencrypted docs): Server maintains its own copy
 *    of each document, can read/index/validate. Used for Intent/Claim pattern.
 *
 * Wire protocol matches the client (shared/local-first/sync.ts).
 */

import * as Automerge from '@automerge/automerge';

// ============================================================================
// TYPES
// ============================================================================

/** Minimal WebSocket interface — works with both browser WebSocket and Bun ServerWebSocket */
export interface SyncWebSocket {
  send(data: string): void;
  readyState: number;
}

interface SyncMessage {
  type: 'sync';
  docId: string;
  data: number[];
}

interface SubscribeMessage {
  type: 'subscribe';
  docIds: string[];
}

interface UnsubscribeMessage {
  type: 'unsubscribe';
  docIds: string[];
}

interface AwarenessMessage {
  type: 'awareness';
  docId: string;
  peer: string;
  cursor?: { x: number; y: number };
  selection?: string;
  username?: string;
  color?: string;
}

interface RelayBackupMessage {
  type: 'relay-backup';
  docId: string;
  data: number[];
}

interface RelayRestoreMessage {
  type: 'relay-restore';
  docId: string;
  data: number[];
}

interface DocListRequestMessage {
  type: 'doc-list-request';
  prefix: string;
  requestId: string;
}

interface DocListResponseMessage {
  type: 'doc-list-response';
  prefix: string;
  docIds: string[];
  requestId: string;
}

type WireMessage =
  | SyncMessage
  | SubscribeMessage
  | UnsubscribeMessage
  | AwarenessMessage
  | RelayBackupMessage
  | RelayRestoreMessage
  | DocListRequestMessage
  | DocListResponseMessage
  | { type: 'ping' }
  | { type: 'pong' };

export interface Peer {
  id: string;
  ws: SyncWebSocket;
  subscribedDocs: Set<string>;
  syncStates: Map<string, Automerge.SyncState>;
  claims?: Record<string, unknown>; // Auth claims
}

export interface SyncServerOptions {
  /** If true, server maintains Automerge docs and participates in sync (default: true) */
  participantMode?: boolean;
  /** Called when a document changes (participant mode only) */
  onDocChange?: (docId: string, doc: Automerge.Doc<any>) => void;
  /** Called when a relay-backup message is received (opaque blob storage) */
  onRelayBackup?: (docId: string, blob: Uint8Array) => void;
  /** Called to load a relay blob for restore on subscribe */
  onRelayLoad?: (docId: string) => Promise<Uint8Array | null>;
  /** Called before evicting a doc from memory (persist to disk) */
  onDocEvict?: (docId: string, doc: Automerge.Doc<any>) => void;
  /** Max docs to keep in memory; 0 = unlimited (default: 500) */
  maxDocs?: number;
}

// ============================================================================
// SyncServer
// ============================================================================

export class SyncServer {
  #peers = new Map<string, Peer>();
  #docs = new Map<string, Automerge.Doc<any>>();
  #docSubscribers = new Map<string, Set<string>>(); // docId → Set<peerId>
  #docLastAccess = new Map<string, number>(); // docId → timestamp
  #participantMode: boolean;
  #relayOnlyDocs = new Set<string>(); // docIds forced to relay mode (encrypted spaces)
  #watchers: Array<{ prefix: string; cb: (docId: string, doc: Automerge.Doc<any>) => void }> = [];
  #onDocChange?: (docId: string, doc: Automerge.Doc<any>) => void;
  #onRelayBackup?: (docId: string, blob: Uint8Array) => void;
  #onRelayLoad?: (docId: string) => Promise<Uint8Array | null>;
  /** Called before evicting a doc so it can be persisted to disk. */
  #onDocEvict?: (docId: string, doc: Automerge.Doc<any>) => void;
  /** Max docs to keep in memory (0 = unlimited). */
  #maxDocs: number;

  constructor(opts: SyncServerOptions = {}) {
    this.#participantMode = opts.participantMode ?? true;
    this.#onDocChange = opts.onDocChange;
    this.#onRelayBackup = opts.onRelayBackup;
    this.#onRelayLoad = opts.onRelayLoad;
    this.#onDocEvict = opts.onDocEvict;
    this.#maxDocs = opts.maxDocs ?? 500;
  }

  /**
   * Mark a doc (or all docs matching a prefix) as relay-only.
   * Relay-only docs are forwarded between peers without server reading them.
   * Used for encrypted spaces where server can't decrypt content.
   */
  setRelayOnly(docIdOrPrefix: string, relay: boolean): void {
    if (relay) {
      this.#relayOnlyDocs.add(docIdOrPrefix);
    } else {
      this.#relayOnlyDocs.delete(docIdOrPrefix);
    }
  }

  /**
   * Check if a docId should use relay mode (either explicitly set or prefix-matched).
   */
  isRelayOnly(docId: string): boolean {
    if (this.#relayOnlyDocs.has(docId)) return true;
    for (const prefix of this.#relayOnlyDocs) {
      if (docId.startsWith(prefix + ':')) return true;
    }
    return false;
  }

  /**
   * Register a new WebSocket peer.
   */
  addPeer(peerId: string, ws: SyncWebSocket, claims?: Record<string, unknown>): Peer {
    const peer: Peer = {
      id: peerId,
      ws,
      subscribedDocs: new Set(),
      syncStates: new Map(),
      claims,
    };
    this.#peers.set(peerId, peer);
    return peer;
  }

  /**
   * Remove a peer and clean up subscriptions.
   */
  removePeer(peerId: string): void {
    const peer = this.#peers.get(peerId);
    if (!peer) return;

    // Remove from all document subscriber sets
    for (const docId of peer.subscribedDocs) {
      const subs = this.#docSubscribers.get(docId);
      if (subs) {
        subs.delete(peerId);
        if (subs.size === 0) {
          this.#docSubscribers.delete(docId);
        }
      }
    }

    this.#peers.delete(peerId);
  }

  /**
   * Handle an incoming message from a peer.
   */
  handleMessage(peerId: string, raw: string | ArrayBuffer): void {
    const peer = this.#peers.get(peerId);
    if (!peer) return;

    try {
      const data = typeof raw === 'string' ? raw : new TextDecoder().decode(raw);
      const msg: WireMessage = JSON.parse(data);

      switch (msg.type) {
        case 'subscribe':
          this.#handleSubscribe(peer, msg as SubscribeMessage);
          break;
        case 'unsubscribe':
          this.#handleUnsubscribe(peer, msg as UnsubscribeMessage);
          break;
        case 'sync':
          this.#handleSync(peer, msg as SyncMessage);
          break;
        case 'awareness':
          this.#handleAwareness(peer, msg as AwarenessMessage);
          break;
        case 'relay-backup':
          this.#handleRelayBackup(peer, msg as RelayBackupMessage);
          break;
        case 'doc-list-request':
          this.#handleDocListRequest(peer, msg as DocListRequestMessage);
          break;
        case 'ping':
          this.#sendToPeer(peer, { type: 'pong' });
          break;
      }
    } catch (e) {
      console.error(`[SyncServer] Error handling message from ${peerId}:`, e);
    }
  }

  /**
   * Get a server-side document (participant mode).
   */
  getDoc<T>(docId: string): Automerge.Doc<T> | undefined {
    if (this.#docs.has(docId)) this.#touchDoc(docId);
    return this.#docs.get(docId);
  }

  /**
   * List all document IDs held by the server.
   */
  listDocs(): string[] {
    return Array.from(this.#docs.keys());
  }

  /**
   * Set/replace a server-side document and sync to all subscribed peers.
   */
  setDoc(docId: string, doc: Automerge.Doc<any>): void {
    this.#docs.set(docId, doc);
    this.#touchDoc(docId);
    this.#syncDocToAllPeers(docId);
    if (this.#onDocChange) {
      this.#onDocChange(docId, doc);
    }
    this.#evictIfNeeded();
  }

  /**
   * Apply a change to a server-side document (e.g. for Intent/Claim validation).
   */
  changeDoc<T>(docId: string, message: string, fn: (doc: T) => void): Automerge.Doc<T> | null {
    let doc = this.#docs.get(docId);
    if (!doc) return null;

    doc = Automerge.change(doc, message, fn as any);
    this.#docs.set(docId, doc);
    this.#touchDoc(docId);
    this.#syncDocToAllPeers(docId);

    if (this.#onDocChange) {
      this.#onDocChange(docId, doc);
    }
    for (const w of this.#watchers) if (docId.includes(w.prefix)) w.cb(docId, doc);

    return doc;
  }

  /**
   * Check if any documents exist whose ID starts with the given prefix.
   */
  hasDocsWithPrefix(prefix: string): boolean {
    for (const key of this.#docs.keys()) {
      if (key.startsWith(prefix)) return true;
    }
    return false;
  }

  /**
   * Get all document IDs held by the server.
   */
  getDocIds(): string[] {
    return Array.from(this.#docs.keys());
  }

  /**
   * Register a callback to fire when any doc whose ID contains `prefix` changes.
   * Returns an unsubscribe function.
   */
  registerWatcher(prefix: string, cb: (docId: string, doc: Automerge.Doc<any>) => void): () => void {
    const entry = { prefix, cb };
    this.#watchers.push(entry);
    return () => {
      const idx = this.#watchers.indexOf(entry);
      if (idx >= 0) this.#watchers.splice(idx, 1);
    };
  }

  /**
   * Get list of connected peer IDs.
   */
  getPeerIds(): string[] {
    return Array.from(this.#peers.keys());
  }

  /**
   * Get subscribers for a document.
   */
  getDocSubscribers(docId: string): string[] {
    return Array.from(this.#docSubscribers.get(docId) ?? []);
  }

  /**
   * Iterate all in-memory docs and call the callback for each.
   * Uses Promise.allSettled so one failure doesn't block others.
   */
  async flushAll(cb: (docId: string, doc: Automerge.Doc<any>) => Promise<void>): Promise<void> {
    const tasks: Promise<void>[] = [];
    for (const [docId, doc] of this.#docs) {
      tasks.push(cb(docId, doc).catch(e => {
        console.error(`[SyncServer] flushAll failed for ${docId}:`, e);
      }));
    }
    await Promise.allSettled(tasks);
  }

  // ---------- Private ----------

  #handleSubscribe(peer: Peer, msg: SubscribeMessage): void {
    for (const docId of msg.docIds) {
      peer.subscribedDocs.add(docId);

      let subs = this.#docSubscribers.get(docId);
      if (!subs) {
        subs = new Set();
        this.#docSubscribers.set(docId, subs);
      }
      subs.add(peer.id);

      // Initialize sync state for this peer-doc pair
      if (!peer.syncStates.has(docId)) {
        peer.syncStates.set(docId, Automerge.initSyncState());
      }

      if (this.isRelayOnly(docId)) {
        // Relay mode: try to send stored encrypted blob
        if (this.#onRelayLoad) {
          this.#onRelayLoad(docId).then((blob) => {
            if (blob) {
              this.#sendToPeer(peer, {
                type: 'relay-restore',
                docId,
                data: Array.from(blob),
              });
            }
          }).catch((e) => {
            console.error(`[SyncServer] Failed to load relay blob for ${docId}:`, e);
          });
        }
      } else if (this.#participantMode && this.#docs.has(docId)) {
        // Participant mode: send initial sync
        this.#sendSyncToPeer(peer, docId);
      }
    }
  }

  #handleUnsubscribe(peer: Peer, msg: UnsubscribeMessage): void {
    for (const docId of msg.docIds) {
      peer.subscribedDocs.delete(docId);
      peer.syncStates.delete(docId);

      const subs = this.#docSubscribers.get(docId);
      if (subs) {
        subs.delete(peer.id);
        if (subs.size === 0) {
          this.#docSubscribers.delete(docId);
        }
      }
    }
  }

  #handleSync(peer: Peer, msg: SyncMessage): void {
    const { docId, data } = msg;
    const syncMsg = new Uint8Array(data);

    if (this.#participantMode && !this.isRelayOnly(docId)) {
      // Server participates: apply sync message to server's doc
      let doc = this.#docs.get(docId);
      if (!doc) {
        // Create an empty doc if this is the first time we see this docId
        doc = Automerge.init();
        this.#docs.set(docId, doc);
        this.#evictIfNeeded();
      }

      let syncState = peer.syncStates.get(docId) ?? Automerge.initSyncState();
      const [newDoc, newSyncState] = Automerge.receiveSyncMessage(doc, syncState, syncMsg);

      const changed = newDoc !== doc;
      this.#docs.set(docId, newDoc);
      this.#touchDoc(docId);
      peer.syncStates.set(docId, newSyncState);

      // Send response sync message back to this peer
      this.#sendSyncToPeer(peer, docId);

      // If doc changed, sync to other subscribers
      if (changed) {
        this.#syncDocToOtherPeers(docId, peer.id);
        if (this.#onDocChange) {
          this.#onDocChange(docId, newDoc);
        }
        for (const w of this.#watchers) if (docId.includes(w.prefix)) w.cb(docId, newDoc);
      }
    } else {
      // Relay mode: forward sync message to all other subscribers
      const subs = this.#docSubscribers.get(docId);
      if (!subs) return;

      for (const subPeerId of subs) {
        if (subPeerId === peer.id) continue;
        const subPeer = this.#peers.get(subPeerId);
        if (subPeer) {
          this.#sendToPeer(subPeer, msg);
        }
      }
    }
  }

  #handleAwareness(peer: Peer, msg: AwarenessMessage): void {
    // Forward awareness to all other subscribers of this document
    const subs = this.#docSubscribers.get(msg.docId);
    if (!subs) return;

    for (const subPeerId of subs) {
      if (subPeerId === peer.id) continue;
      const subPeer = this.#peers.get(subPeerId);
      if (subPeer) {
        this.#sendToPeer(subPeer, msg);
      }
    }
  }

  #handleRelayBackup(_peer: Peer, msg: RelayBackupMessage): void {
    const blob = new Uint8Array(msg.data);
    if (this.#onRelayBackup) {
      this.#onRelayBackup(msg.docId, blob);
    }
  }

  #handleDocListRequest(peer: Peer, msg: DocListRequestMessage): void {
    const docIds = this.getDocIds().filter(id => id.startsWith(msg.prefix));
    this.#sendToPeer(peer, {
      type: 'doc-list-response',
      prefix: msg.prefix,
      docIds,
      requestId: msg.requestId,
    });
  }

  #sendSyncToPeer(peer: Peer, docId: string): void {
    const doc = this.#docs.get(docId);
    if (!doc) return;

    let syncState = peer.syncStates.get(docId) ?? Automerge.initSyncState();
    const [newSyncState, syncMessage] = Automerge.generateSyncMessage(doc, syncState);

    peer.syncStates.set(docId, newSyncState);

    if (syncMessage) {
      this.#sendToPeer(peer, {
        type: 'sync',
        docId,
        data: Array.from(syncMessage),
      });
    }
  }

  #syncDocToAllPeers(docId: string): void {
    const subs = this.#docSubscribers.get(docId);
    if (!subs) return;

    for (const peerId of subs) {
      const peer = this.#peers.get(peerId);
      if (peer) {
        this.#sendSyncToPeer(peer, docId);
      }
    }
  }

  #syncDocToOtherPeers(docId: string, excludePeerId: string): void {
    const subs = this.#docSubscribers.get(docId);
    if (!subs) return;

    for (const peerId of subs) {
      if (peerId === excludePeerId) continue;
      const peer = this.#peers.get(peerId);
      if (peer) {
        this.#sendSyncToPeer(peer, docId);
      }
    }
  }

  #sendToPeer(peer: Peer, msg: object): void {
    try {
      if (peer.ws.readyState === 1) {
        peer.ws.send(JSON.stringify(msg));
      }
    } catch (e) {
      console.error(`[SyncServer] Failed to send to peer ${peer.id}:`, e);
    }
  }

  /** Mark a doc as recently accessed for LRU tracking. */
  #touchDoc(docId: string): void {
    this.#docLastAccess.set(docId, Date.now());
  }

  /**
   * Evict least-recently-used docs that have no active subscribers.
   * Called after adding new docs to keep memory bounded.
   */
  #evictIfNeeded(): void {
    if (this.#maxDocs <= 0 || this.#docs.size <= this.#maxDocs) return;

    // Build list of evictable docs (no active subscribers)
    const evictable: Array<{ docId: string; lastAccess: number }> = [];
    for (const docId of this.#docs.keys()) {
      const subs = this.#docSubscribers.get(docId);
      if (subs && subs.size > 0) continue; // has active subscribers, skip
      evictable.push({
        docId,
        lastAccess: this.#docLastAccess.get(docId) ?? 0,
      });
    }

    // Sort by oldest access first
    evictable.sort((a, b) => a.lastAccess - b.lastAccess);

    // Evict until under the cap
    const toEvict = this.#docs.size - this.#maxDocs;
    let evicted = 0;
    for (const { docId } of evictable) {
      if (evicted >= toEvict) break;

      const doc = this.#docs.get(docId);
      if (doc && this.#onDocEvict) {
        this.#onDocEvict(docId, doc);
      }
      this.#docs.delete(docId);
      this.#docLastAccess.delete(docId);
      evicted++;
    }

    if (evicted > 0) {
      console.log(`[SyncServer] Evicted ${evicted} idle docs (${this.#docs.size} remaining)`);
    }
  }
}
