WizardConnect/packages/dapp/src/dapp-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 { decodeHdPublicKey } from "@bitauth/libauth";
import type { HdPublicNodeValid } from "@bitauth/libauth";
import {
  RelayClient,
  RelayStatus,
  RelayMsgAction,
  DappReadyMessage,
  WalletReadyMessage,
  DisconnectMessage,
  DisconnectReason,
  SignTransactionRequest,
  SignTransactionResponse,
  SignCancelMessage,
  ProtocolMessage,
  PROTOCOL_NAME,
  PathName,
  childIndexOfPathName,
  isHdwalletv1Session,
} from "@wizardconnect/core";
import { DappPubkeyStateManager } from "./pubkey-state-manager.js";

export interface DappConnectionManagerEvents {
  /** Fired after wallet_ready is received and state is updated. */
  walletready: [msg: WalletReadyMessage];
  /** Fired for every protocol message sent by the dapp. */
  messagesent: [msg: ProtocolMessage];
  /** Fired for every protocol message received from the wallet. */
  messagereceived: [msg: ProtocolMessage];
  /** Fired on disconnect — either remote-initiated or protocol mismatch. */
  disconnect: [reason: DisconnectReason, message: string | undefined];
}

/**
 * Manages the dapp side of a single WizardConnect session.
 *
 * Protocol responsibilities:
 *  - Sends dapp_ready after key exchange completes (and on reconnect)
 *  - Receives wallet_ready, sign_transaction_response
 *  - Tracks pubkeys and address indices (via xpub on-demand derivation)
 *
 * No chain-specific logic — child indices are plain numbers:
 *   0 = Receive, 1 = Change, 7 = Cauldron (BCH convention)
 */
export class DappConnectionManager extends EventEmitter<DappConnectionManagerEvents> {
  private conn: RelayClient | null = null;
  private listenerAttached = false;

  /** Pubkey state — exposed for callers that need to query by index. */
  readonly pubkeyState: DappPubkeyStateManager;

  walletName: string | null = null;
  walletIcon: string | null = null;
  protocol: string | null = null;

  /** Protocols this dapp supports, in preference order. */
  private readonly supportedProtocols: string[] = [PROTOCOL_NAME];

  private walletDiscovered = false;
  private pendingSignatureRequests = new Map<
    number,
    {
      resolve: (r: SignTransactionResponse) => void;
      reject: (e: Error) => void;
    }
  >();

  /**
   * @param dappName  Optional display name of the dapp (sent in dapp_ready).
   * @param dappIcon  Optional icon URL/data-URI of the dapp (sent in dapp_ready).
   */
  constructor(
    private dappName?: string,
    private dappIcon?: string,
  ) {
    super();
    this.pubkeyState = new DappPubkeyStateManager();
  }

  /**
   * Call this from the RelayStatusCallback passed to `initiateDappRelay`.
   * Attaches the message listener exactly once and re-sends dapp_ready
   * each time the connection is established (handles reconnects).
   */
  updateConnection(
    client: RelayClient | null | undefined,
    status: RelayStatus,
  ): void {
    if (client) {
      // Attach message listener once (same RelayClient object is reused across reconnects)
      if (!this.listenerAttached) {
        this.listenerAttached = true;
        client.on("message", (msg: ProtocolMessage) => this.handleMessage(msg));
      }
      this.conn = client;
    }

    if (status.status === "connected" && this.conn) {
      this.onConnected();
    }
  }

  isWalletDiscovered(): boolean {
    return this.walletDiscovered;
  }

  /**
   * Get a sequence number from the relay client.
   * Use this to populate the `sequence` field of a SignTransactionRequest.
   */
  nextSequence(): number {
    if (!this.conn) throw new Error("[wizardconnect/dapp] Not connected");
    return this.conn.nextSequence();
  }

  /**
   * Send a sign transaction request and wait for the wallet's response.
   * The caller is responsible for creating the full SignTransactionRequest
   * (including sequence from `nextSequence()`).
   */
  async sendSignRequest(
    request: SignTransactionRequest,
  ): Promise<SignTransactionResponse> {
    if (!this.conn) throw new Error("[wizardconnect/dapp] Not connected");

    return new Promise<SignTransactionResponse>((resolve, reject) => {
      this.pendingSignatureRequests.set(request.sequence, { resolve, reject });

      this.conn!.relay(request)
        .then(() => {
          this.emit("messagesent", request);
        })
        .catch((err) => {
          this.pendingSignatureRequests.delete(request.sequence);
          reject(err instanceof Error ? err : new Error(String(err)));
        });
    });
  }

  /**
   * Cancel an in-flight sign request.
   * Immediately rejects the pending Promise and sends sign_cancel to the wallet.
   */
  async sendSignCancel(sequence: number, reason?: string): Promise<void> {
    // Reject pending promise immediately — no response will come
    const handlers = this.pendingSignatureRequests.get(sequence);
    if (handlers) {
      this.pendingSignatureRequests.delete(sequence);
      handlers.reject(new Error(reason ?? "Sign request cancelled"));
    }
    if (!this.conn) return;
    const msg: SignCancelMessage = {
      action: RelayMsgAction.SignCancel,
      sequence,
      ...(reason !== undefined && { reason }),
      time: Math.floor(Date.now() / 1000),
    };
    await this.conn.relay(msg);
    this.emit("messagesent", msg);
  }

  /**
   * Send a disconnect message to the wallet (courtesy notification).
   * The caller is responsible for calling dappRelay.cleanup() afterwards.
   */
  async sendDisconnect(message?: string): Promise<void> {
    if (!this.conn) return;
    const msg: DisconnectMessage = {
      action: RelayMsgAction.Disconnect,
      reason: DisconnectReason.UserDisconnect,
      time: Math.floor(Date.now() / 1000),
      ...(message !== undefined && { message }),
    };
    await this.conn.relay(msg);
    this.emit("messagesent", msg);
  }

  // --- Pubkey state delegation ------------------------------------------------
  // Convenience methods that forward to pubkeyState.

  getPubkey(childIndex: number, index: bigint): Uint8Array | undefined {
    return this.pubkeyState.getPubkey(childIndex, index);
  }

  /** Returns true if an xpub node is available for this child index. */
  hasPath(childIndex: number): boolean {
    return this.pubkeyState.hasPath(childIndex);
  }

  /**
   * Returns the stored xpub node for the given child index.
   * Available after wallet_ready is received.
   */
  getXpubNode(childIndex: number): HdPublicNodeValid | undefined {
    return this.pubkeyState.getXpubNode(childIndex);
  }

  // --- Private protocol handling -------------------------------------------

  private onConnected(): void {
    (async () => {
      // Wait until key exchange is complete before sending dapp_ready
      while (this.conn && !this.conn.isKeyExchangeComplete()) {
        await new Promise((r) => setTimeout(r, 100));
      }
      if (this.conn) {
        await this.pushDappReady();
      }
    })().catch((e) =>
      console.error("[wizardconnect/dapp] Error in onConnected:", e),
    );
  }

  private async pushDappReady(): Promise<void> {
    if (!this.conn) return;

    const msg: DappReadyMessage = {
      action: RelayMsgAction.DappReady,
      supported_protocols: this.supportedProtocols,
      wallet_discovered: this.walletDiscovered,
      time: Math.floor(Date.now() / 1000),
      // Include selected_protocol on the reactive send (after the dapp has seen the wallet)
      ...(this.walletDiscovered &&
        this.protocol && { selected_protocol: this.protocol }),
      ...(this.dappName !== undefined && { dapp_name: this.dappName }),
      ...(this.dappIcon !== undefined && { dapp_icon: this.dappIcon }),
    };

    await this.conn.relay(msg);
    this.emit("messagesent", msg);
  }

  private handleMessage(msg: ProtocolMessage): void {
    this.emit("messagereceived", msg);
    switch (msg.action) {
      case RelayMsgAction.WalletReady:
        this.handleWalletReady(msg as WalletReadyMessage);
        break;
      case RelayMsgAction.SignTransactionResponse:
        this.handleSignTransactionResponse(msg as SignTransactionResponse);
        break;
      case RelayMsgAction.Disconnect:
        this.handleRemoteDisconnect(msg as DisconnectMessage);
        break;
      case RelayMsgAction.DappReady:
        // Not expected on dapp side — silently ignore
        break;
      default:
        console.warn(
          "[wizardconnect/dapp] Unknown message action:",
          msg.action,
        );
    }
  }

  private handleRemoteDisconnect(msg: DisconnectMessage): void {
    this.emit("disconnect", msg.reason, msg.message);
  }

  private handleWalletReady(msg: WalletReadyMessage): void {
    this.walletDiscovered = true;
    this.walletName = msg.wallet_name;
    this.walletIcon = msg.wallet_icon;

    // Protocol selection: first match in dapp's preference order
    const agreed = this.supportedProtocols.find((p) =>
      msg.supported_protocols.includes(p),
    );
    if (!agreed) {
      const detail = `No protocol overlap. Wallet: [${msg.supported_protocols}], Dapp: [${this.supportedProtocols}]`;
      const disconnectMsg: DisconnectMessage = {
        action: RelayMsgAction.Disconnect,
        reason: DisconnectReason.ProtocolMismatch,
        message: detail,
        time: Math.floor(Date.now() / 1000),
      };
      this.conn?.relay(disconnectMsg).catch(() => {});
      this.emit("disconnect", DisconnectReason.ProtocolMismatch, detail);
      return;
    }
    this.protocol = agreed;

    // Extract and validate protocol-specific session data
    const sessionData = msg.session[agreed];
    if (!isHdwalletv1Session(sessionData)) {
      console.error(
        "[wizardconnect/dapp] Invalid hdwalletv1 session data:",
        sessionData,
      );
      return;
    }

    // Store xpub nodes — no eager derivation; consumer drives it
    for (const pathInfo of sessionData.paths) {
      const decoded = decodeHdPublicKey(pathInfo.xpub);
      if (typeof decoded === "string") {
        console.warn(
          "[wizardconnect/dapp] Bad xpub for path",
          pathInfo.name,
          decoded,
        );
        continue;
      }
      const ci = childIndexOfPathName(pathInfo.name as PathName);
      this.pubkeyState.setXpubNode(ci, decoded.node);
    }

    if (!msg.dapp_discovered) {
      this.pushDappReady().catch((e) =>
        console.error("[wizardconnect/dapp] Error pushing dapp_ready:", e),
      );
    }

    this.emit("walletready", msg);
  }

  private handleSignTransactionResponse(
    response: SignTransactionResponse,
  ): void {
    const handlers = this.pendingSignatureRequests.get(response.sequence);
    if (!handlers) {
      console.warn(
        "[wizardconnect/dapp] No pending request for sequence:",
        response.sequence,
      );
      return;
    }
    this.pendingSignatureRequests.delete(response.sequence);

    if (response.error) {
      handlers.reject(new Error(response.error));
    } else {
      handlers.resolve(response);
    }
  }
}