import crypto from 'crypto';
import { PrismaClient } from '@prisma/client';

const prisma = new PrismaClient();

function decrypt(ciphertext: string, keyHex: string): string {
  const key = Buffer.from(keyHex, 'hex');
  const [ivHex, encHex, tagHex] = ciphertext.split(':');
  const iv = Buffer.from(ivHex, 'hex');
  const encrypted = Buffer.from(encHex, 'hex');
  const tag = Buffer.from(tagHex, 'hex');
  const decipher = crypto.createDecipheriv('aes-256-gcm', key, iv);
  decipher.setAuthTag(tag);
  return decipher.update(encrypted, undefined, 'utf8') + decipher.final('utf8');
}

function encrypt(plaintext: string, keyHex: string): string {
  const key = Buffer.from(keyHex, 'hex');
  const iv = crypto.randomBytes(12);
  const cipher = crypto.createCipheriv('aes-256-gcm', key, iv);
  const encrypted = Buffer.concat([cipher.update(plaintext, 'utf8'), cipher.final()]);
  const tag = cipher.getAuthTag();
  return `${iv.toString('hex')}:${encrypted.toString('hex')}:${tag.toString('hex')}`;
}

async function main() {
  const oldKey = process.argv[2];
  const newKey = process.argv[3];
  const dryRun = process.argv.includes('--dry-run');

  if (!oldKey || !newKey) {
    console.error('Usage: tsx scripts/rotate-totp-key.ts <old-key-hex> <new-key-hex> [--dry-run]');
    process.exit(1);
  }

  const users = await prisma.users.findMany({
    where: { totp_enabled: true, totp_secret: { not: null } },
    select: { id: true, username: true, totp_secret: true },
  });

  console.log(`Found ${users.length} users with TOTP enabled`);

  if (dryRun) {
    for (const user of users) {
      try {
        const decrypted = decrypt(user.totp_secret!, oldKey);
        const reEncrypted = encrypt(decrypted, newKey);
        decrypt(reEncrypted, newKey); // verify round-trip
        console.log(`  [OK] ${user.username} (id=${user.id}) — decryption and re-encryption verified`);
      } catch (e) {
        console.error(`  [FAIL] ${user.username} (id=${user.id}) — ${e}`);
      }
    }
    console.log('\nDry run complete. No changes made.');
    return;
  }

  await prisma.$transaction(async (tx) => {
    for (const user of users) {
      const decrypted = decrypt(user.totp_secret!, oldKey);
      const reEncrypted = encrypt(decrypted, newKey);
      await tx.users.update({ where: { id: user.id }, data: { totp_secret: reEncrypted } });
      console.log(`  Re-encrypted TOTP for ${user.username} (id=${user.id})`);
    }
  });

  console.log('\nAll TOTP secrets re-encrypted successfully.');
  await prisma.$disconnect();
}

main().catch((e) => { console.error(e); process.exit(1); });