WizardConnect/packages/react/src/hooks/useWizardConnect.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 { useState, useRef, useEffect, useCallback } from "react";
import {
  initiateDappRelay,
  type DappRelayResult,
  type RelayUpdatePayload,
  type PathXpub,
  binToHex,
} from "@wizardconnect/core";
import { DappConnectionManager } from "@wizardconnect/dapp";
import type {
  UseWizardConnectOptions,
  UseWizardConnectResult,
  WizardConnectState,
} from "../types.js";

const DEFAULT_SESSION_KEY = "wizardconnect-session";

interface StoredSession {
  privateKey: string;
  secret: string;
  walletPublicKey?: string;
  /** Raw xpub paths from wallet_ready, persisted for offline pubkey derivation. */
  paths?: PathXpub[];
}

function loadSession(key: string): StoredSession | null {
  if (typeof localStorage === "undefined") return null;
  const stored = localStorage.getItem(key);
  if (!stored) return null;
  try {
    const parsed = JSON.parse(stored) as StoredSession;
    if (!parsed.privateKey || !parsed.secret) return null;
    return parsed;
  } catch {
    return null;
  }
}

function saveSession(key: string, session: StoredSession): void {
  if (typeof localStorage === "undefined") return;
  localStorage.setItem(key, JSON.stringify(session));
}

function clearSession(key: string): void {
  if (typeof localStorage === "undefined") return;
  localStorage.removeItem(key);
}

/**
 * React hook that encapsulates the WizardConnect relay lifecycle.
 *
 * Manages: relay initiation, DappConnectionManager, key exchange events,
 * session persistence, and auto-reconnect on page refresh.
 *
 * The returned `manager` can be used to build an app-specific wallet adapter.
 */
export function useWizardConnect(
  options: UseWizardConnectOptions = {},
): UseWizardConnectResult {
  const {
    dappName,
    dappIcon,
    relayUrls,
    sessionKey = DEFAULT_SESSION_KEY,
    persistSession = true,
  } = options;

  const [state, setState] = useState<WizardConnectState>("idle");
  const [manager, setManager] = useState<DappConnectionManager | null>(null);
  const [uri, setUri] = useState<string | null>(null);
  const [qrUri, setQrUri] = useState<string | null>(null);
  const [walletName, setWalletName] = useState<string | null>(null);
  const [walletIcon, setWalletIcon] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);

  const relayRef = useRef<DappRelayResult | null>(null);
  const managerRef = useRef<DappConnectionManager | null>(null);
  const autoReconnectAttempted = useRef(false);

  const startRelay = useCallback(
    (existingCredentials?: {
      privateKey: string;
      secret: string;
      walletPublicKey: string;
    }): boolean => {
      if (state === "connecting" || state === "connected") return false;

      setError(null);
      setState("connecting");

      const mgr = new DappConnectionManager(dappName, dappIcon);
      managerRef.current = mgr;
      setManager(mgr);

      mgr.on("walletready", () => {
        setWalletName(mgr.walletName);
        setWalletIcon(mgr.walletIcon);
        setState("connected");

        // Persist xpub paths so getPubkey works on next page load
        if (persistSession) {
          const paths = mgr.getSessionPaths();
          if (paths.length > 0) {
            const stored = loadSession(sessionKey);
            if (stored) {
              stored.paths = paths;
              saveSession(sessionKey, stored);
            }
          }
        }
      });

      mgr.on("disconnect", () => {
        setState("disconnected");
        setWalletName(null);
        setWalletIcon(null);
      });

      try {
        const relay = initiateDappRelay(
          (payload: RelayUpdatePayload) => {
            mgr.updateConnection(payload.client, payload.status);
          },
          {
            existingCredentials,
            explicitRelayUrls: relayUrls,
          },
        );

        relayRef.current = relay;
        setUri(relay.uri);
        setQrUri(relay.qrUri);

        if (persistSession) {
          // Merge with existing session to preserve walletPublicKey and paths
          const existing = loadSession(sessionKey);
          saveSession(sessionKey, {
            ...existing,
            privateKey: relay.credentials.privateKey,
            secret: relay.credentials.secret,
          });
        }

        relay.events.on(
          "keyexchangecomplete",
          (walletPublicKey: Uint8Array) => {
            if (persistSession) {
              const stored = loadSession(sessionKey);
              if (stored) {
                stored.walletPublicKey = binToHex(walletPublicKey);
                saveSession(sessionKey, stored);
              }
            }
          },
        );

        return true;
      } catch (e) {
        const message =
          e instanceof Error ? e.message : "Failed to start relay";
        setError(message);
        setState("idle");
        return false;
      }
    },
    [state, dappName, dappIcon, relayUrls, sessionKey, persistSession],
  );

  const connect = useCallback((): boolean => {
    return startRelay();
  }, [startRelay]);

  const disconnect = useCallback(async (): Promise<void> => {
    try {
      if (managerRef.current) {
        await managerRef.current.sendDisconnect().catch(() => {});
      }
    } finally {
      relayRef.current?.cleanup();
      relayRef.current = null;
      managerRef.current = null;
      setManager(null);
      setUri(null);
      setQrUri(null);
      setWalletName(null);
      setWalletIcon(null);
      setState("idle");
      if (persistSession) {
        clearSession(sessionKey);
      }
    }
  }, [persistSession, sessionKey]);

  // Auto-reconnect on mount if a stored session exists
  useEffect(() => {
    if (autoReconnectAttempted.current) return;
    if (!persistSession) return;
    autoReconnectAttempted.current = true;

    const stored = loadSession(sessionKey);
    console.log(
      "[wizardconnect/react] auto-reconnect: stored session:",
      stored
        ? `walletPublicKey=${!!stored.walletPublicKey}, paths=${stored.paths?.length ?? 0}`
        : "null",
    );
    if (!stored || !stored.walletPublicKey) return;

    const started = startRelay({
      privateKey: stored.privateKey,
      secret: stored.secret,
      walletPublicKey: stored.walletPublicKey,
    });

    // Restore cached xpub paths so getPubkey works before wallet_ready
    if (started && stored.paths?.length && managerRef.current) {
      try {
        managerRef.current.restoreSessionPaths(stored.paths);
      } catch (e) {
        // Corrupt cached paths — clear them but don't block reconnect
        console.warn(
          "[wizardconnect/react] Failed to restore cached xpub paths:",
          e,
        );
        const refreshed = loadSession(sessionKey);
        if (refreshed) {
          delete refreshed.paths;
          saveSession(sessionKey, refreshed);
        }
      }
    }
  }, [persistSession, sessionKey, startRelay]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      relayRef.current?.cleanup();
    };
  }, []);

  return {
    state,
    manager,
    uri,
    qrUri,
    walletName,
    walletIcon,
    connect,
    disconnect,
    error,
  };
}