import assert from "node:assert/strict";
import { mkdtempSync, readFileSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import test from "node:test";
import { privateKeyToAccount } from "viem/accounts";
import type { AgentConfig } from "../src/config.js";
import { AgentState } from "../src/state.js";
import { decryptPrivateKey, encryptPrivateKey, provisionWallet, resolveExecutionConfig } from "../src/wallet-vault.js";

const masterKey = Buffer.alloc(32, 7).toString("base64url");
const privateKey = `0x${"12".repeat(32)}`;

test("wallet private keys round-trip through authenticated encryption", () => {
  const encrypted = encryptPrivateKey(privateKey, masterKey);
  assert.notEqual(encrypted, privateKey);
  assert.equal(decryptPrivateKey(encrypted, masterKey), privateKey);
  const tampered = Buffer.from(encrypted, "base64url");
  tampered[tampered.length - 1] ^= 1;
  assert.throws(() => decryptPrivateKey(tampered.toString("base64url"), masterKey));
});

test("wallet provisioning stores ciphertext and becomes the active execution account", () => {
  const directory = mkdtempSync(join(tmpdir(), "zoi-agent-vault-"));
  const stateDbPath = join(directory, "state.db");
  const config = {
    stateDbPath,
    account: null,
    feeWallet: null,
    control: { walletMasterKey: masterKey },
  } as unknown as AgentConfig;
  try {
    const address = provisionWallet(config, privateKey);
    assert.equal(address, privateKeyToAccount(privateKey).address);
    const resolved = resolveExecutionConfig(config);
    assert.equal(resolved.account?.address, address);
    assert.equal(resolved.feeWallet, address);
    assert.equal(readFileSync(stateDbPath).includes(Buffer.from(privateKey)), false);
  } finally {
    rmSync(directory, { recursive: true, force: true });
  }
});

test("a connected-wallet mismatch is rejected before the vault changes", () => {
  const directory = mkdtempSync(join(tmpdir(), "zoi-agent-vault-mismatch-"));
  const stateDbPath = join(directory, "state.db");
  const config = {
    stateDbPath,
    account: null,
    feeWallet: null,
    control: { walletMasterKey: masterKey },
  } as unknown as AgentConfig;
  try {
    assert.throws(() => provisionWallet(config, privateKey, `0x${"ab".repeat(20)}`), /does not match/);
    const state = new AgentState(stateDbPath);
    assert.equal(state.getActiveWalletProfile(), undefined);
    state.close();
  } finally {
    rmSync(directory, { recursive: true, force: true });
  }
});
