WizardConnect/packages/wallet/src/wallet-connection-manager.ts

// Copyright (C) 2026 Whiterun LLC,
// This software is licensed under the GNU Lesser General Public License (LGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/lgpl-3.0.html

import { EventEmitter } from "eventemitter3";
import {
  RelayClient,
  RelayStatus,
  ProtocolMessage,
  RelayUpdatePayload,
  RelayStatusCallback,
  initiateWalletRelay,
  SignTransactionRequest,
  SignTransactionResponse,
  SignCancelMessage,
  RelayMsgAction,
  DappReadyMessage,
  WalletReadyMessage,
  DisconnectMessage,
  DisconnectReason,
  PathXpub,
  Hdwalletv1Session,
  PROTOCOL_NAME,
  binToHex,
} from "@wizardconnect/core";
import { WalletAdapter } from "./wallet-adapter.js";
import { DerivationPath } from "./derivation-path.js";

export interface RelayConnectionState {
  id: string;
  uri: string;
  status: RelayStatus;
  label: string;
  dappName: string | null;
  dappIcon: string | null;
  connectedAt: number;
}

export interface PendingSignRequest {
  connectionId: string;
  request: SignTransactionRequest;
}

interface ActiveConnection {
  id: string;
  uri: string;
  cleanup: () => void;
  client: RelayClient | null;
  status: RelayStatus;
  label: string;
  dappName: string | null;
  dappIcon: string | null;
  connectedAt: number;
  dappDiscovered: boolean;
  /// Prevents duplicate wallet_ready messages within a single connection cycle.
  /// Reset to false on each new connect/reconnect; set to true after sending.
  walletReadySentThisCycle: boolean;
  notificationQueue: ProtocolMessage[];
  notificationProcessor: ReturnType<typeof setInterval> | null;
  /// Key exchange data embedded in wallet_ready
  walletPublicKeyHex: string;
  keyExchangeSecret: string;
}

export type WalletConnectionManagerEvents = {
  connectionStatusChanged: [connectionId: string, status: RelayStatus];
  pendingSignRequest: [request: PendingSignRequest];
  connectionsChanged: [];
  remoteDisconnect: [
    connectionId: string,
    reason: DisconnectReason,
    message: string | undefined,
  ];
  signCancelled: [
    connectionId: string,
    sequence: number,
    reason: string | undefined,
  ];
  /** Fired for protocol messages not handled by the core protocol (extension actions). */
  message: [connectionId: string, message: ProtocolMessage];
};

/**
 * Manages multiple simultaneous dapp connections for a wallet.
 * Handles sign request queuing and reconnection.
 */
export class WalletConnectionManager extends EventEmitter<WalletConnectionManagerEvents> {
  private connections: Map<string, ActiveConnection> = new Map();
  private activeSignSequences = new Set<number>();
  private adapter: WalletAdapter;

  constructor(adapter: WalletAdapter) {
    super();
    this.adapter = adapter;
  }

  /**
   * Connect to a dapp using a wiz:// URI.
   * Returns the connection ID.
   */
  connect(uri: string): string {
    // Return existing connection if one for this URI is already active
    for (const conn of this.connections.values()) {
      if (conn.uri === uri) {
        return conn.id;
      }
    }

    const id = `rc_${Date.now()}_${Math.random().toString(36).slice(2, 8)}`;

    const conn: ActiveConnection = {
      id,
      uri,
      cleanup: () => {},
      client: null,
      status: RelayStatus.disconnected(),
      label: "Connecting...",
      dappName: null,
      dappIcon: null,
      connectedAt: Date.now(),
      dappDiscovered: false,
      walletReadySentThisCycle: false,
      notificationQueue: [],
      notificationProcessor: null,
      walletPublicKeyHex: "",
      keyExchangeSecret: "",
    };

    this.connections.set(id, conn);

    const statusCallback: RelayStatusCallback = (
      payload: RelayUpdatePayload,
    ) => {
      // Attach message listener exactly once, the first time we receive the client.
      if (!conn.client && payload.client) {
        payload.client.on("message", (message: ProtocolMessage) => {
          this.handleMessage(conn, message).catch((err) => {
            console.error(
              "[wizardconnect/wallet] Error handling message:",
              err,
            );
          });
        });
      }

      conn.client = payload.client;
      conn.status = payload.status;

      if (payload.status.status === "connected") {
        this.onConnected(conn);
      } else if (
        payload.status.status === "disconnected" ||
        payload.status.status === "reconnecting"
      ) {
        if (conn.notificationProcessor) {
          clearInterval(conn.notificationProcessor);
          conn.notificationProcessor = null;
        }
      }

      this.emit("connectionStatusChanged", id, payload.status);
      this.emit("connectionsChanged");
    };

    const result = initiateWalletRelay(statusCallback, {
      uri,
      walletPrivateKey: this.adapter.getRelayPrivateKey(uri),
    });

    conn.cleanup = result.cleanup;
    conn.walletPublicKeyHex = binToHex(result.walletPublicKey);
    conn.keyExchangeSecret = result.secret;

    this.emit("connectionsChanged");
    return id;
  }

  /**
   * Disconnect a specific dapp connection, sending a courtesy UserDisconnect message.
   */
  disconnect(connectionId: string): void {
    this.doDisconnect(connectionId, true);
  }

  /**
   * Disconnect all connections.
   */
  disconnectAll(): void {
    for (const id of [...this.connections.keys()]) {
      this.doDisconnect(id, true);
    }
  }

  private doDisconnect(connectionId: string, sendMessage: boolean): void {
    const conn = this.connections.get(connectionId);
    if (!conn) return;

    if (sendMessage && conn.client) {
      const disconnectMsg: DisconnectMessage = {
        action: RelayMsgAction.Disconnect,
        reason: DisconnectReason.UserDisconnect,
        time: Math.floor(Date.now() / 1000),
      };
      conn.client.relay(disconnectMsg).catch(() => {});
    }

    clearInterval(conn.notificationProcessor ?? undefined);
    conn.cleanup();
    this.connections.delete(connectionId);
    this.emit("connectionsChanged");
  }

  /**
   * Get all connection states (for UI/Redux).
   */
  getConnections(): Record<string, RelayConnectionState> {
    const result: Record<string, RelayConnectionState> = {};
    for (const [id, conn] of this.connections) {
      result[id] = {
        id: conn.id,
        uri: conn.uri,
        status: conn.status,
        label: conn.label,
        dappName: conn.dappName,
        dappIcon: conn.dappIcon,
        connectedAt: conn.connectedAt,
      };
    }
    return result;
  }

  /**
   * Send a sign transaction response back to the dapp.
   */
  async sendSignResponse(
    connectionId: string,
    sequence: number,
    signedTransactionHex: string,
  ): Promise<void> {
    const conn = this.connections.get(connectionId);
    if (!conn?.client) {
      throw new Error(`Connection ${connectionId} not found or not connected`);
    }

    this.activeSignSequences.delete(sequence);

    const response: SignTransactionResponse = {
      action: RelayMsgAction.SignTransactionResponse,
      sequence,
      signedTransaction: signedTransactionHex,
      time: Math.floor(Date.now() / 1000),
    };

    await conn.client.relay(response);
  }

  /**
   * Send a sign error response back to the dapp.
   */
  async sendSignError(
    connectionId: string,
    sequence: number,
    errorMessage: string,
  ): Promise<void> {
    const conn = this.connections.get(connectionId);
    if (!conn?.client) {
      return; // Already disconnected, nothing to do
    }

    this.activeSignSequences.delete(sequence);

    const response: SignTransactionResponse = {
      action: RelayMsgAction.SignTransactionResponse,
      sequence,
      signedTransaction: "",
      error: errorMessage,
      time: Math.floor(Date.now() / 1000),
    };

    await conn.client.relay(response);
  }

  // --- Private connection lifecycle ---

  /** Immediately attempt to flush the notification queue (fire-and-forget). */
  private flushNotificationQueue(conn: ActiveConnection): void {
    this.processNotificationQueue(conn).catch((err) => {
      console.error("[wizardconnect/wallet] Notification queue error:", err);
    });
  }

  private onConnected(conn: ActiveConnection): void {
    // New connection cycle: reset dedup flags so wallet_ready is sent fresh
    // and dapp_discovered is false — the dapp must re-send dapp_ready to
    // re-establish the session (matches "Wallet reconnects" scenario in protocol.md).
    conn.walletReadySentThisCycle = false;
    conn.dappDiscovered = false;

    // Start notification processor as fallback for retries after send errors
    if (conn.notificationProcessor) {
      clearInterval(conn.notificationProcessor);
    }
    conn.notificationProcessor = setInterval(() => {
      this.processNotificationQueue(conn).catch((err) => {
        console.error("[wizardconnect/wallet] Notification queue error:", err);
      });
    }, 1000);

    // Wait for key exchange, then send wallet_ready
    (async () => {
      while (conn.client && !conn.client.isKeyExchangeComplete()) {
        await new Promise((resolve) => setTimeout(resolve, 100));
      }
      if (conn.client) {
        this.pushWalletReady(conn).catch((err) => {
          console.error(
            "[wizardconnect/wallet] Failed to send wallet_ready:",
            err,
          );
        });
      }
    })();
  }

  private async handleMessage(
    conn: ActiveConnection,
    message: ProtocolMessage,
  ): Promise<void> {
    switch (message.action) {
      case RelayMsgAction.DappReady:
        await this.handleDappReady(conn, message as DappReadyMessage);
        break;
      case RelayMsgAction.SignTransactionRequest:
        this.handleSignRequest(conn, message as SignTransactionRequest);
        break;
      case RelayMsgAction.WalletReady:
        console.warn(
          "[wizardconnect/wallet] Got wallet_ready as wallet, ignoring",
        );
        break;
      case RelayMsgAction.Disconnect:
        this.handleRemoteDisconnect(conn, message as DisconnectMessage);
        break;
      case RelayMsgAction.SignCancel:
        this.handleSignCancel(conn, message as SignCancelMessage);
        break;
      default:
        this.emit("message", conn.id, message);
    }
  }

  private handleRemoteDisconnect(
    conn: ActiveConnection,
    msg: DisconnectMessage,
  ): void {
    this.emit("remoteDisconnect", conn.id, msg.reason, msg.message);
    this.doDisconnect(conn.id, false);
  }

  private handleSignCancel(
    conn: ActiveConnection,
    msg: SignCancelMessage,
  ): void {
    this.activeSignSequences.delete(msg.sequence);
    this.emit("signCancelled", conn.id, msg.sequence, msg.reason);
  }

  private async handleDappReady(
    conn: ActiveConnection,
    msg: DappReadyMessage,
  ): Promise<void> {
    // Log the dapp's protocol selection for diagnostics
    if (msg.selected_protocol) {
      console.debug(
        "[wizardconnect/wallet] Dapp selected protocol:",
        msg.selected_protocol,
      );
    }

    // Update dapp metadata from the first dapp_ready that carries it
    if (msg.dapp_name && !conn.dappName) {
      conn.dappName = msg.dapp_name;
      conn.dappIcon = msg.dapp_icon ?? null;
      conn.label = msg.dapp_name;
      this.emit("connectionStatusChanged", conn.id, conn.status);
    }

    if (!msg.wallet_discovered) {
      // Dapp hasn't seen us yet (or has reset, e.g. browser refresh) — force
      // re-introduction even if we already sent wallet_ready this cycle.
      conn.walletReadySentThisCycle = false;
      await this.pushWalletReady(conn);
      return;
    }

    conn.dappDiscovered = true;
  }

  private async pushWalletReady(conn: ActiveConnection): Promise<void> {
    if (conn.walletReadySentThisCycle) return;
    conn.walletReadySentThisCycle = true;

    const paths: PathXpub[] = [
      { name: "receive", xpub: this.adapter.getXpub(DerivationPath.Receive) },
      { name: "change", xpub: this.adapter.getXpub(DerivationPath.Change) },
      { name: "defi", xpub: this.adapter.getXpub(DerivationPath.Cauldron) },
      ...(this.adapter.getAdditionalPaths?.() ?? []),
    ];

    const extensions = this.adapter.getExtensions?.();
    const hdwv1Session: Hdwalletv1Session = {
      paths,
      ...(extensions ? { extensions } : {}),
    };

    const msg: WalletReadyMessage = {
      action: RelayMsgAction.WalletReady,
      supported_protocols: [PROTOCOL_NAME],
      wallet_name: this.adapter.walletName,
      wallet_icon: this.adapter.walletIcon,
      time: Math.floor(Date.now() / 1000),
      dapp_discovered: conn.dappDiscovered,
      session: {
        [PROTOCOL_NAME]: hdwv1Session,
      },
      public_key: conn.walletPublicKeyHex,
      secret: conn.keyExchangeSecret,
    };

    conn.notificationQueue.push(msg);
    this.flushNotificationQueue(conn);
  }

  private handleSignRequest(
    conn: ActiveConnection,
    msg: SignTransactionRequest,
  ): void {
    // Deduplicate: the dapp re-sends pending requests after reconnect, so the
    // wallet may receive the same sequence twice while it's still awaiting
    // user approval.
    if (this.activeSignSequences.has(msg.sequence)) {
      return;
    }
    this.activeSignSequences.add(msg.sequence);

    // Emit to host app for queuing/approval
    const pendingRequest: PendingSignRequest = {
      connectionId: conn.id,
      request: msg,
    };
    this.emit("pendingSignRequest", pendingRequest);
  }

  private async processNotificationQueue(
    conn: ActiveConnection,
  ): Promise<void> {
    if (!conn.client || conn.notificationQueue.length === 0) {
      return;
    }

    const notifications = [...conn.notificationQueue];
    conn.notificationQueue = [];

    for (const notification of notifications) {
      try {
        await conn.client.relay(notification);
      } catch (error) {
        console.error(
          "[wizardconnect/wallet] Failed to send notification:",
          error,
        );
        conn.notificationQueue.push(notification);
      }
    }
  }
}