import { createCipheriv, createDecipheriv, createHash, randomBytes } from "node:crypto";
import type { Hex } from "viem";
import { privateKeyToAccount } from "viem/accounts";
import type { AgentConfig } from "./config.js";
import { AgentState } from "./state.js";

const keyPattern = /^0x[a-fA-F0-9]{64}$/;

const decodeMasterKey = (value: string) => {
  const buffer = /^[a-fA-F0-9]{64}$/.test(value)
    ? Buffer.from(value, "hex")
    : Buffer.from(value, "base64url");
  if (buffer.length !== 32) throw new Error("WALLET_MASTER_KEY must decode to exactly 32 bytes");
  return buffer;
};

export const encryptPrivateKey = (privateKey: string, masterKey: string) => {
  if (!keyPattern.test(privateKey)) throw new Error("Invalid EVM private key");
  const iv = randomBytes(12);
  const cipher = createCipheriv("aes-256-gcm", decodeMasterKey(masterKey), iv);
  const ciphertext = Buffer.concat([cipher.update(privateKey, "utf8"), cipher.final()]);
  const tag = cipher.getAuthTag();
  return Buffer.concat([iv, tag, ciphertext]).toString("base64url");
};

export const decryptPrivateKey = (encrypted: string, masterKey: string) => {
  const payload = Buffer.from(encrypted, "base64url");
  if (payload.length < 29) throw new Error("Encrypted wallet payload is invalid");
  const iv = payload.subarray(0, 12);
  const tag = payload.subarray(12, 28);
  const ciphertext = payload.subarray(28);
  const decipher = createDecipheriv("aes-256-gcm", decodeMasterKey(masterKey), iv);
  decipher.setAuthTag(tag);
  const privateKey = Buffer.concat([decipher.update(ciphertext), decipher.final()]).toString("utf8");
  if (!keyPattern.test(privateKey)) throw new Error("Decrypted wallet key is invalid");
  return privateKey as Hex;
};

export const provisionWallet = (config: AgentConfig, privateKey: string, expectedAddress?: string) => {
  const masterKey = config.control.walletMasterKey;
  if (!masterKey) throw new Error("Wallet encryption is not configured");
  const account = privateKeyToAccount(privateKey as Hex);
  if (expectedAddress && account.address.toLowerCase() !== expectedAddress.toLowerCase()) {
    throw new Error("Private key does not match the connected wallet");
  }
  const state = new AgentState(config.stateDbPath);
  const previousAddress = state.getActiveWalletProfile()?.address ?? config.account?.address;
  if (previousAddress) state.scopeLegacyWalletData(previousAddress);
  state.upsertWalletProfile({
    id: createHash("sha256").update(account.address.toLowerCase()).digest("hex").slice(0, 32),
    address: account.address,
    encryptedPrivateKey: encryptPrivateKey(privateKey, masterKey),
  });
  state.close();
  return account.address;
};

export const resolveExecutionConfig = (config: AgentConfig): AgentConfig => {
  const masterKey = config.control.walletMasterKey;
  if (!masterKey) return config;
  const state = new AgentState(config.stateDbPath);
  const profile = state.getActiveWalletProfile();
  if (profile) state.scopeLegacyWalletData(profile.address);
  else if (config.account) state.scopeLegacyWalletData(config.account.address);
  state.close();
  if (!profile) return config;
  const account = privateKeyToAccount(decryptPrivateKey(profile.encryptedPrivateKey, masterKey));
  if (account.address.toLowerCase() !== profile.address.toLowerCase()) {
    throw new Error("Encrypted wallet does not match its stored address");
  }
  return { ...config, account, feeWallet: account.address };
};
