AutistMask/src/popup/views/addWallet.js

const { $, showView, showFlash, goBack, clearViewStack } = require("./helpers");
const {
    generateMnemonic,
    hdWalletFromMnemonic,
    isValidMnemonic,
    addressFromPrivateKey,
    hdWalletFromXprv,
    isValidXprv,
} = require("../../shared/wallet");
const { encryptWithPassword } = require("../../shared/vault");
const { state, saveState } = require("../../shared/state");
const { scanForAddresses } = require("../../shared/balances");

/**
 * Check if an address already exists in ANY wallet (hd, xprv, or key).
 * Returns the wallet object if found, or undefined.
 */
function findWalletByAddress(addr) {
    const lower = addr.toLowerCase();
    return state.wallets.find((w) =>
        w.addresses.some((a) => a.address.toLowerCase() === lower),
    );
}

/**
 * Check if an xpub already exists in any HD-type wallet (hd or xprv).
 * Returns the wallet object if found, or undefined.
 */
function findWalletByXpub(xpub) {
    return state.wallets.find((w) => w.xpub && w.xpub === xpub);
}

let currentMode = "mnemonic";

const MODES = ["mnemonic", "privkey", "xprv"];

const PASSWORD_HINTS = {
    mnemonic:
        "This password encrypts your recovery phrase on this device. You will need it to send funds.",
    privkey:
        "This password encrypts your private key on this device. You will need it to send funds.",
    xprv: "This password encrypts your key on this device. You will need it to send funds.",
};

function switchMode(mode) {
    currentMode = mode;
    for (const m of MODES) {
        $("add-wallet-section-" + m).classList.toggle("hidden", m !== mode);
        const tab = $("tab-" + m);
        const isActive = m === mode;
        // Active: bold, solid border on top/sides, no bottom border (connects to content)
        tab.classList.toggle("font-bold", isActive);
        tab.classList.toggle("border-solid", isActive);
        tab.classList.toggle("border-border", isActive);
        tab.classList.toggle("border-b-bg", isActive);
        tab.classList.toggle("bg-bg", isActive);
        // Inactive: muted text, dashed border on top/sides, transparent bottom, hover invert
        tab.classList.toggle("text-muted", !isActive);
        tab.classList.toggle("border-dashed", !isActive);
        tab.classList.toggle("border-border-light", !isActive);
        tab.classList.toggle("border-b-transparent", !isActive);
        tab.classList.toggle("hover:bg-fg", !isActive);
        tab.classList.toggle("hover:text-bg", !isActive);
    }
    $("add-wallet-password-hint").textContent = PASSWORD_HINTS[mode];
}

function show() {
    $("wallet-mnemonic").value = "";
    $("import-private-key").value = "";
    $("import-xprv-key").value = "";
    $("add-wallet-password").value = "";
    $("add-wallet-password-confirm").value = "";
    $("add-wallet-phrase-warning").style.visibility = "hidden";
    switchMode("mnemonic");
    showView("add-wallet");
}

function validatePassword() {
    const pw = $("add-wallet-password").value;
    const pw2 = $("add-wallet-password-confirm").value;
    if (!pw) {
        showFlash("Please choose a password.");
        return null;
    }
    if (pw.length < 12) {
        showFlash("Password must be at least 12 characters.");
        return null;
    }
    if (pw !== pw2) {
        showFlash("Passwords do not match.");
        return null;
    }
    return pw;
}

async function importMnemonic(ctx) {
    const mnemonic = $("wallet-mnemonic").value.trim();
    if (!mnemonic) {
        showFlash("Enter a recovery phrase or press the die to generate one.");
        return;
    }
    const words = mnemonic.split(/\s+/);
    if (words.length !== 12 && words.length !== 24) {
        showFlash(
            "Recovery phrase must be 12 or 24 words. You entered " +
                words.length +
                ".",
        );
        return;
    }
    if (!isValidMnemonic(mnemonic)) {
        showFlash("Invalid recovery phrase. Check for typos.");
        return;
    }
    const pw = validatePassword();
    if (!pw) return;
    const { xpub, firstAddress } = hdWalletFromMnemonic(mnemonic);
    const xpubDup = findWalletByXpub(xpub);
    if (xpubDup) {
        showFlash(
            "This recovery phrase is already added (" + xpubDup.name + ").",
        );
        return;
    }
    const addrDup = findWalletByAddress(firstAddress);
    if (addrDup) {
        showFlash("Address already exists in wallet (" + addrDup.name + ").");
        return;
    }
    const encrypted = await encryptWithPassword(mnemonic, pw);
    const walletNum = state.wallets.length + 1;
    const wallet = {
        type: "hd",
        name: "Wallet " + walletNum,
        xpub: xpub,
        encryptedSecret: encrypted,
        nextIndex: 1,
        addresses: [
            { address: firstAddress, balance: "0.0000", tokenBalances: [] },
        ],
    };
    state.wallets.push(wallet);
    state.hasWallet = true;
    await saveState();
    clearViewStack();
    ctx.renderWalletList();
    showView("main");

    // Scan for used HD addresses beyond index 0.
    showFlash("Scanning for addresses...", 30000);
    const scan = await scanForAddresses(xpub, state.rpcUrl);
    if (scan.addresses.length > 1) {
        wallet.addresses = scan.addresses.map((a) => ({
            address: a.address,
            balance: "0.0000",
            tokenBalances: [],
        }));
        wallet.nextIndex = scan.nextIndex;
        await saveState();
        ctx.renderWalletList();
        showFlash("Found " + scan.addresses.length + " addresses.");
    } else {
        showFlash("Ready.", 1000);
    }

    ctx.doRefreshAndRender();
}

async function importPrivateKey(ctx) {
    const key = $("import-private-key").value.trim();
    if (!key) {
        showFlash("Please enter your private key.");
        return;
    }
    let addr;
    try {
        addr = addressFromPrivateKey(key);
    } catch (e) {
        showFlash("Invalid private key.");
        return;
    }
    const pw = validatePassword();
    if (!pw) return;
    const duplicate = findWalletByAddress(addr);
    if (duplicate) {
        showFlash(
            "This address already exists in wallet (" + duplicate.name + ").",
        );
        return;
    }
    const encrypted = await encryptWithPassword(key, pw);
    const walletNum = state.wallets.length + 1;
    state.wallets.push({
        type: "key",
        name: "Wallet " + walletNum,
        encryptedSecret: encrypted,
        addresses: [{ address: addr, balance: "0.0000", tokenBalances: [] }],
    });
    state.hasWallet = true;
    await saveState();
    clearViewStack();
    ctx.renderWalletList();
    showView("main");

    ctx.doRefreshAndRender();
}

async function importXprvKey(ctx) {
    const xprv = $("import-xprv-key").value.trim();
    if (!xprv) {
        showFlash("Please enter your extended private key.");
        return;
    }
    if (!isValidXprv(xprv)) {
        showFlash("Invalid extended private key.");
        return;
    }
    let result;
    try {
        result = hdWalletFromXprv(xprv);
    } catch (e) {
        showFlash("Invalid extended private key.");
        return;
    }
    const { xpub, firstAddress } = result;
    const xpubDup = findWalletByXpub(xpub);
    if (xpubDup) {
        showFlash("This key is already added (" + xpubDup.name + ").");
        return;
    }
    const addrDup = findWalletByAddress(firstAddress);
    if (addrDup) {
        showFlash("Address already exists in wallet (" + addrDup.name + ").");
        return;
    }
    const pw = validatePassword();
    if (!pw) return;
    const encrypted = await encryptWithPassword(xprv, pw);
    const walletNum = state.wallets.length + 1;
    const wallet = {
        type: "xprv",
        name: "Wallet " + walletNum,
        xpub: xpub,
        encryptedSecret: encrypted,
        nextIndex: 1,
        addresses: [
            { address: firstAddress, balance: "0.0000", tokenBalances: [] },
        ],
    };
    state.wallets.push(wallet);
    state.hasWallet = true;
    await saveState();
    clearViewStack();
    ctx.renderWalletList();
    showView("main");

    // Scan for used HD addresses beyond index 0.
    showFlash("Scanning for addresses...", 30000);
    const scan = await scanForAddresses(xpub, state.rpcUrl);
    if (scan.addresses.length > 1) {
        wallet.addresses = scan.addresses.map((a) => ({
            address: a.address,
            balance: "0.0000",
            tokenBalances: [],
        }));
        wallet.nextIndex = scan.nextIndex;
        await saveState();
        ctx.renderWalletList();
        showFlash("Found " + scan.addresses.length + " addresses.");
    } else {
        showFlash("Ready.", 1000);
    }

    ctx.doRefreshAndRender();
}

function init(ctx) {
    // Tab click handlers
    $("tab-mnemonic").addEventListener("click", () => switchMode("mnemonic"));
    $("tab-privkey").addEventListener("click", () => switchMode("privkey"));
    $("tab-xprv").addEventListener("click", () => switchMode("xprv"));

    // Generate mnemonic
    $("btn-generate-phrase").addEventListener("click", () => {
        $("wallet-mnemonic").value = generateMnemonic();
        $("add-wallet-phrase-warning").style.visibility = "visible";
    });

    // Import / confirm
    $("btn-add-wallet-confirm").addEventListener("click", async () => {
        if (currentMode === "mnemonic") {
            await importMnemonic(ctx);
        } else if (currentMode === "privkey") {
            await importPrivateKey(ctx);
        } else if (currentMode === "xprv") {
            await importXprvKey(ctx);
        }
    });

    // Back button
    $("btn-add-wallet-back").addEventListener("click", () => {
        goBack();
    });
}

module.exports = { init, show };