import { prettyPrintJSON } from '@aztec/cli/utils';
import { asyncPool } from '@aztec/foundation/async-pool';
import { deriveBlsPrivateKey } from '@aztec/foundation/crypto/bls';
import { createBn254Keystore } from '@aztec/foundation/crypto/bls/bn254_keystore';
import { computeBn254G1PublicKeyCompressed } from '@aztec/foundation/crypto/bn254';
import { EthAddress } from '@aztec/foundation/eth-address';
import type { LogFn } from '@aztec/foundation/log';
import type { EthAccount, EthPrivateKey, ValidatorKeyStore } from '@aztec/node-keystore/types';
import type { AztecAddress } from '@aztec/stdlib/aztec-address';

import { Wallet } from '@ethersproject/wallet';
import { constants as fsConstants, mkdirSync } from 'fs';
import { access, writeFile } from 'fs/promises';
import { availableParallelism, homedir } from 'os';
import { dirname, isAbsolute, join } from 'path';
import { mnemonicToAccount } from 'viem/accounts';
import { Worker } from 'worker_threads';

import { defaultBlsPath } from './utils.js';

type EthJsonV3WorkerResult = { json: string } | { error: { message: string; name?: string; stack?: string } };

// ethers' default scrypt parameters use substantial memory, so keep worker fan-out bounded.
const maxEthKeystoreWorkers = Math.max(1, Math.min(4, availableParallelism()));

function deserializeWorkerError(error: { message: string; name?: string; stack?: string }) {
  const result = new Error(error.message);
  result.name = error.name ?? result.name;
  result.stack = error.stack;
  return result;
}

function encryptEthJsonV3InWorker(privateKeyHex: string, password: string): Promise<string> {
  return new Promise<string>((resolve, reject) => {
    const worker = new Worker(new URL('./eth_json_v3_worker.js', import.meta.url), {
      workerData: { privateKeyHex, password },
    });
    let settled = false;

    worker.once('message', (result: EthJsonV3WorkerResult) => {
      settled = true;

      if ('json' in result) {
        resolve(result.json);
      } else {
        reject(deserializeWorkerError(result.error));
      }
    });
    worker.once('error', error => {
      settled = true;
      reject(error);
    });
    worker.once('exit', code => {
      if (!settled) {
        reject(new Error(`ETH JSON V3 worker stopped with exit code ${code}`));
      }
    });
  });
}

async function encryptEthJsonV3(privateKeyHex: string, password: string): Promise<string> {
  if (process.env.JEST_WORKER_ID !== undefined) {
    return await new Wallet(privateKeyHex).encrypt(password);
  }
  return await encryptEthJsonV3InWorker(privateKeyHex, password);
}

export type ValidatorSummary = { attesterEth?: string; attesterBls?: string; publisherEth?: string[] };

export type BuildValidatorsInput = {
  validatorCount: number;
  publisherCount?: number;
  publishers?: string[];
  accountIndex: number;
  baseAddressIndex: number;
  mnemonic: string;
  ikm?: string;
  blsPath?: string;
  feeRecipient: AztecAddress;
  coinbase?: EthAddress;
  remoteSigner?: string;
};

export function withValidatorIndex(path: string, accountIndex: number = 0, addressIndex: number = 0) {
  // NOTE: The legacy BLS CLI is to allow users who generated keys in 2.1.4 to be able to use the same command
  // to re-generate their keys. In 2.1.5 we switched how we append addresses to the path so this is to maintain backwards compatibility.
  const useLegacyBlsCli = ['true', '1', 'yes', 'y'].includes(process.env.LEGACY_BLS_CLI ?? '');

  const defaultBlsPathParts = defaultBlsPath.split('/');

  const parts = path.split('/');
  if (parts.length == defaultBlsPathParts.length && parts.every((part, index) => part === defaultBlsPathParts[index])) {
    if (useLegacyBlsCli) {
      // In 2.1.4, we were using address-index in parts[3] and did NOT use account-index, check lines 32 & 84
      // https://github.com/AztecProtocol/aztec-packages/blob/v2.1.4/yarn-project/cli/src/cmds/validator_keys/shared.ts

      parts[3] = String(addressIndex);
    } else {
      parts[3] = String(accountIndex);
      parts[5] = String(addressIndex);
    }
    return parts.join('/');
  }
  return path;
}

/**
 * Compute a compressed BN254 G1 public key from a private key.
 * @param privateKeyHex - Private key as 0x-prefixed hex string
 * @returns Compressed G1 point (32 bytes with sign bit in MSB)
 */
export async function computeBlsPublicKeyCompressed(privateKeyHex: string): Promise<string> {
  return await computeBn254G1PublicKeyCompressed(privateKeyHex);
}

export function deriveEthAttester(
  mnemonic: string,
  baseAccountIndex: number,
  addressIndex: number,
  remoteSigner?: string,
): EthAccount | EthPrivateKey {
  const acct = mnemonicToAccount(mnemonic, { accountIndex: baseAccountIndex, addressIndex });
  return remoteSigner
    ? ({ address: acct.address as unknown as EthAddress, remoteSignerUrl: remoteSigner } as EthAccount)
    : (('0x' + Buffer.from(acct.getHdKey().privateKey!).toString('hex')) as EthPrivateKey);
}

/**
 * Resolve a `--funding-account` value into a keystore `EthAccount`. A 66-char value is a private key
 * (used verbatim). A 42-char value is an address: with an explicit `remoteSigner` URL it becomes an
 * `{ address, remoteSignerUrl }` pair; without one it is stored as a bare address that falls back to
 * the keystore-level remote signer at runtime. Callers must validate the value first (see
 * `validateFundingAccountOptions`).
 */
export function resolveFundingAccount(fundingAccount: string, remoteSigner?: string): EthAccount {
  if (fundingAccount.length === 66) {
    return fundingAccount as EthPrivateKey;
  }
  const address = EthAddress.fromString(fundingAccount);
  return remoteSigner ? ({ address, remoteSignerUrl: remoteSigner } as EthAccount) : address;
}

export async function buildValidatorEntries(input: BuildValidatorsInput) {
  const {
    validatorCount,
    publisherCount = 0,
    publishers,
    accountIndex,
    baseAddressIndex,
    mnemonic,
    ikm,
    blsPath,
    feeRecipient,
    coinbase,
    remoteSigner,
  } = input;

  const summaries: ValidatorSummary[] = [];

  const validators = await Promise.all(
    Array.from({ length: validatorCount }, async (_unused, i) => {
      const addressIndex = baseAddressIndex + i;
      const basePath = blsPath ?? defaultBlsPath;
      const perValidatorPath = withValidatorIndex(basePath, accountIndex, addressIndex);

      const blsPrivKey = ikm || mnemonic ? deriveBlsPrivateKey(mnemonic, ikm, perValidatorPath) : undefined;
      const blsPubCompressed = blsPrivKey ? await computeBlsPublicKeyCompressed(blsPrivKey) : undefined;

      const ethAttester = deriveEthAttester(mnemonic, accountIndex, addressIndex, remoteSigner);
      const attester = blsPrivKey ? { eth: ethAttester, bls: blsPrivKey } : ethAttester;

      let publisherField: EthAccount | EthPrivateKey | (EthAccount | EthPrivateKey)[] | undefined;
      const publisherAddresses: string[] = [];
      if (publishers && publishers.length > 0) {
        publisherAddresses.push(...publishers);
        publisherField = publishers.length === 1 ? (publishers[0] as EthPrivateKey) : (publishers as EthPrivateKey[]);
      } else if (publisherCount > 0) {
        const publishersBaseIndex = baseAddressIndex + validatorCount + i * publisherCount;
        const publisherAccounts = Array.from({ length: publisherCount }, (_unused2, j) => {
          const publisherAddressIndex = publishersBaseIndex + j;
          const pubAcct = mnemonicToAccount(mnemonic, {
            accountIndex,
            addressIndex: publisherAddressIndex,
          });
          publisherAddresses.push(pubAcct.address as unknown as string);
          return remoteSigner
            ? ({ address: pubAcct.address as unknown as EthAddress, remoteSignerUrl: remoteSigner } as EthAccount)
            : (('0x' + Buffer.from(pubAcct.getHdKey().privateKey!).toString('hex')) as EthPrivateKey);
        });
        publisherField = publisherCount === 1 ? publisherAccounts[0] : publisherAccounts;
      }

      const acct = mnemonicToAccount(mnemonic, {
        accountIndex,
        addressIndex,
      });
      const attesterEthAddress = acct.address as unknown as string;
      summaries.push({
        attesterEth: attesterEthAddress,
        attesterBls: blsPubCompressed,
        publisherEth: publisherAddresses.length > 0 ? publisherAddresses : undefined,
      });

      return {
        attester,
        ...(publisherField !== undefined ? { publisher: publisherField } : {}),
        feeRecipient,
        coinbase: coinbase ?? attesterEthAddress,
      } as ValidatorKeyStore;
    }),
  );

  return { validators, summaries };
}

export async function resolveKeystoreOutputPath(dataDir?: string, file?: string) {
  const defaultDataDir = join(homedir(), '.aztec', 'keystore');
  const resolvedDir = dataDir && dataDir.length > 0 ? dataDir : defaultDataDir;
  let outputPath: string;
  if (file && file.length > 0) {
    outputPath = isAbsolute(file) ? file : join(resolvedDir, file);
  } else {
    let index = 1;
    while (true) {
      const candidate = join(resolvedDir, `key${index}.json`);
      try {
        await access(candidate, fsConstants.F_OK);
        index += 1;
      } catch {
        outputPath = candidate;
        break;
      }
    }
  }
  return { resolvedDir, outputPath: outputPath! };
}

export async function writeKeystoreFile(path: string, keystore: unknown) {
  mkdirSync(dirname(path), { recursive: true });
  await writeFile(path, JSON.stringify(keystore, null, 2), { encoding: 'utf-8' });
}

export function logValidatorSummaries(log: LogFn, summaries: ValidatorSummary[]) {
  const lines: string[] = [];
  for (let i = 0; i < summaries.length; i++) {
    const v = summaries[i];
    lines.push(`acc${i + 1}:`);
    lines.push(`  attester:`);
    if (v.attesterEth) {
      lines.push(`    eth: ${v.attesterEth}`);
    }
    if (v.attesterBls) {
      lines.push(`    bls: ${v.attesterBls}`);
    }
    if (v.publisherEth && v.publisherEth.length > 0) {
      lines.push(`  publisher:`);
      for (const addr of v.publisherEth) {
        lines.push(`    - ${addr}`);
      }
    }
  }
  if (lines.length > 0) {
    log(lines.join('\n'));
  }
}

export function maybePrintJson(log: LogFn, jsonFlag: boolean | undefined, obj: unknown) {
  if (jsonFlag) {
    log(prettyPrintJSON(obj as Record<string, any>));
  }
}

/**
 * Writes a BN254 keystore file for a BN254 BLS private key.
 * Returns the absolute path to the written file.
 *
 * @param outDir - Directory to write the keystore file to
 * @param fileNameBase - Base name for the keystore file (will be sanitized)
 * @param password - Password for encrypting the private key
 * @param privateKeyHex - Private key as 0x-prefixed hex string (32 bytes)
 * @param pubkeyHex - Public key as hex string
 * @param derivationPath - BIP-44 style derivation path
 * @returns Absolute path to the written keystore file
 */
export async function writeBn254BlsKeystore(
  outDir: string,
  fileNameBase: string,
  password: string,
  privateKeyHex: string,
  pubkeyHex: string,
  derivationPath: string,
): Promise<string> {
  mkdirSync(outDir, { recursive: true });

  const keystore = await createBn254Keystore(password, privateKeyHex, pubkeyHex, derivationPath);

  const safeBase = fileNameBase.replace(/[^a-zA-Z0-9_-]/g, '_');
  const outPath = join(outDir, `keystore-${safeBase}.json`);
  await writeFile(outPath, JSON.stringify(keystore, null, 2), { encoding: 'utf-8' });
  return outPath;
}

/** Replace plaintext BLS keys in validators with { path, password } pointing to BN254 keystore files. */
export async function writeBlsBn254ToFile(
  validators: ValidatorKeyStore[],
  options: { outDir: string; password: string; blsPath?: string },
): Promise<void> {
  await Promise.all(
    validators.map(async (v, i) => {
      if (!v || typeof v !== 'object' || !('attester' in v)) {
        return;
      }
      const att = (v as any).attester;

      // Shapes: { bls: <hex> } or { eth: <ethAccount>, bls?: <hex> } or plain EthAccount
      const blsKey: string | undefined = typeof att === 'object' && 'bls' in att ? (att as any).bls : undefined;
      if (!blsKey || typeof blsKey !== 'string') {
        return;
      }

      const pub = await computeBlsPublicKeyCompressed(blsKey);
      const path = options.blsPath ?? defaultBlsPath;
      const fileBase = `${String(i + 1)}_${pub.slice(2, 18)}`;
      const keystorePath = await writeBn254BlsKeystore(options.outDir, fileBase, options.password, blsKey, pub, path);

      if (typeof att === 'object') {
        (att as any).bls = { path: keystorePath, password: options.password };
      }
    }),
  );
}

/** Writes an Ethereum JSON V3 keystore using ethers, returns absolute path */
export async function writeEthJsonV3Keystore(
  outDir: string,
  fileNameBase: string,
  password: string,
  privateKeyHex: string,
): Promise<string> {
  const safeBase = fileNameBase.replace(/[^a-zA-Z0-9_-]/g, '_');
  mkdirSync(outDir, { recursive: true });
  const json = await encryptEthJsonV3(privateKeyHex, password);
  const outPath = join(outDir, `keystore-eth-${safeBase}.json`);
  await writeFile(outPath, json, { encoding: 'utf-8' });
  return outPath;
}

/**
 * If `account` is a plaintext ETH private key, encrypt it to a JSON V3 file and return a
 * { path, password } reference; otherwise return it unchanged.
 */
async function maybeEncryptEthAccount(account: any, label: string, options: { outDir: string; password: string }) {
  if (typeof account === 'string' && account.startsWith('0x') && account.length === 66) {
    const fileBase = `${label}_${account.slice(2, 10)}`;
    const p = await writeEthJsonV3Keystore(options.outDir, fileBase, options.password, account);
    return { path: p, password: options.password };
  }
  return account;
}

/** Encrypt a plaintext funding-account key to a JSON V3 file, replacing it with a { path, password } reference. */
export async function encryptFundingAccountToFile(
  account: EthAccount,
  options: { outDir: string; password: string },
): Promise<EthAccount> {
  return (await maybeEncryptEthAccount(account, 'funding', options)) as EthAccount;
}

/** Replace plaintext ETH keys in validators with { path, password } pointing to JSON V3 files. */
export async function writeEthJsonV3ToFile(
  validators: ValidatorKeyStore[],
  options: { outDir: string; password: string },
): Promise<void> {
  const tasks: (() => Promise<void>)[] = [];

  const maybeQueueEncryptEth = (account: any, label: string, setEncryptedAccount: (account: any) => void) => {
    if (typeof account === 'string' && account.startsWith('0x') && account.length === 66) {
      tasks.push(async () => {
        const fileBase = `${label}_${account.slice(2, 10)}`;
        const path = await writeEthJsonV3Keystore(options.outDir, fileBase, options.password, account);
        setEncryptedAccount({ path, password: options.password });
      });
    }
  };

  for (let i = 0; i < validators.length; i++) {
    const v = validators[i];
    if (!v || typeof v !== 'object') {
      continue;
    }

    // attester may be string (eth), object with eth, or remote signer
    const att = (v as any).attester;
    if (typeof att === 'string') {
      maybeQueueEncryptEth(att, `attester_${i + 1}`, account => ((v as any).attester = account));
    } else if (att && typeof att === 'object' && 'eth' in att) {
      maybeQueueEncryptEth((att as any).eth, `attester_${i + 1}`, account => ((att as any).eth = account));
    }

    // publisher can be single or array
    if ('publisher' in v) {
      const pub = (v as any).publisher;
      if (Array.isArray(pub)) {
        for (let j = 0; j < pub.length; j++) {
          maybeQueueEncryptEth(pub[j], `publisher_${i + 1}_${j + 1}`, account => (pub[j] = account));
        }
      } else if (pub !== undefined) {
        maybeQueueEncryptEth(pub, `publisher_${i + 1}`, account => ((v as any).publisher = account));
      }
    }
  }

  await asyncPool(maxEthKeystoreWorkers, tasks, task => task());
}