import { randomBytes } from 'node:crypto';
import { Prisma } from '../generated/prisma';
import {
  decryptMfaSecret,
  encryptMfaSecret,
} from '../src/auth/mfa/mfa-crypto.util';
import { generateTotpSecret } from '../src/auth/mfa/totp.util';
import { gcmSeal, legacyJwtKey } from '../src/common/crypto/key-ring';
import { RekeyJob } from '../src/jobs/rekey.job';
import { runCli } from '../src/jobs/run-cli';
import { PrismaService } from '../src/prisma/prisma.service';
import {
  openWebhookSecret,
  sealWebhookSecret,
} from '../src/webhooks/webhook-secret';

type Form = 'legacy' | 'previous' | 'current' | 'plaintext';
type MfaForm = Exclude<Form, 'plaintext'>;

const MFA_FORMS: readonly MfaForm[] = ['legacy', 'previous', 'current'];
const WEBHOOK_FORMS: readonly Form[] = [
  'legacy',
  'plaintext',
  'previous',
  'current',
];

const PREVIOUS_KEYS = {
  MFA_ENCRYPTION_KEY_PREVIOUS: `rekey-e2e-previous-mfa-${'p'.repeat(32)}`,
  WEBHOOK_ENCRYPTION_KEY_PREVIOUS: `rekey-e2e-previous-webhook-${'q'.repeat(32)}`,
} as const;

/** How the code before the key ring sealed: AES-256-GCM under sha256(JWT_ACCESS_SECRET). */
function legacyCiphertext(plaintext: string): string {
  const { iv, tag, ciphertext } = gcmSeal(
    legacyJwtKey(),
    Buffer.from(plaintext, 'utf8'),
  );
  return [iv, tag, ciphertext].map((part) => part.toString('base64')).join('.');
}

/** Seals with the key in `<variable>_PREVIOUS`, as if before the rotation. */
function sealedWithPrevious(variable: string, seal: () => string): string {
  const current = process.env[variable];
  process.env[variable] = process.env[`${variable}_PREVIOUS`];
  try {
    return seal();
  } finally {
    if (current === undefined) {
      delete process.env[variable];
    } else {
      process.env[variable] = current;
    }
  }
}

function sealMfa(form: MfaForm, plaintext: string): string {
  if (form === 'legacy') {
    return legacyCiphertext(plaintext);
  }
  if (form === 'previous') {
    return sealedWithPrevious('MFA_ENCRYPTION_KEY', () =>
      encryptMfaSecret(plaintext),
    );
  }
  return encryptMfaSecret(plaintext);
}

function sealWebhook(form: Form, plaintext: string): string {
  if (form === 'legacy') {
    return `enc:v1:${legacyCiphertext(plaintext)}`;
  }
  if (form === 'plaintext') {
    return plaintext;
  }
  if (form === 'previous') {
    return sealedWithPrevious('WEBHOOK_ENCRYPTION_KEY', () =>
      sealWebhookSecret(plaintext),
    );
  }
  return sealWebhookSecret(plaintext);
}

const keyIdOf = (stored: string): string => stored.split(':')[0];

/**
 * The job as `node dist/cli.js rekey [--dry-run]` runs it: runCli with the
 * real JobsModule against the shared test database. The job visits every
 * row there, also those of e2e files running in parallel. So this file
 * sets only *_PREVIOUS: the current keys stay the ones every e2e process
 * has (the development keys from JWT_ACCESS_SECRET in .env.test, unless the
 * shell sets real ones), and whatever this job reseals still opens there.
 * Assertions on single rows use this file's rows; counts over the whole
 * database are lower bounds.
 */
describe('rekey job (e2e)', () => {
  const suffix = `${Date.now()}-${Math.round(Math.random() * 1e6)}`;
  const testStart = new Date(Date.now() - 1000);
  const savedEnv = Object.fromEntries(
    Object.keys(PREVIOUS_KEYS).map((name) => [name, process.env[name]]),
  );
  let prisma: PrismaService;
  const mfa = {} as Record<MfaForm, { id: string; plaintext: string }>;
  const webhook = {} as Record<Form, { id: string; plaintext: string }>;
  const created = { mfa: {}, webhook: {} } as {
    mfa: Record<MfaForm, string>;
    webhook: Record<Form, string>;
  };
  const extraUserIds: string[] = [];

  const createUser = async (name: string, mfaSecret: string) => {
    const user = await prisma.user.create({
      data: {
        email: `rekey-e2e-${name}-${suffix}@example.com`,
        passwordHash: 'not-used',
        name: `Rekey ${name}`,
        mfaEnabled: true,
        mfaSecret,
        role: { connect: { key: 'READONLY' } },
      },
    });
    return user.id;
  };

  /** This file's rows as stored now. */
  const storedValues = async () => {
    const users = await prisma.user.findMany({
      where: { id: { in: MFA_FORMS.map((form) => mfa[form].id) } },
      select: { id: true, mfaSecret: true },
    });
    const subscriptions = await prisma.webhookSubscription.findMany({
      where: { id: { in: WEBHOOK_FORMS.map((form) => webhook[form].id) } },
      select: { id: true, secret: true },
    });
    const userById = new Map(users.map((row) => [row.id, row.mfaSecret]));
    const subscriptionById = new Map(
      subscriptions.map((row) => [row.id, row.secret]),
    );
    return {
      mfa: Object.fromEntries(
        MFA_FORMS.map((form) => [form, userById.get(mfa[form].id)]),
      ) as Record<MfaForm, string>,
      webhook: Object.fromEntries(
        WEBHOOK_FORMS.map((form) => [
          form,
          subscriptionById.get(webhook[form].id),
        ]),
      ) as Record<Form, string>,
    };
  };

  const rekeyRuns = () =>
    prisma.jobRun.findMany({
      where: { job: 'rekey', startedAt: { gte: testStart } },
      orderBy: { startedAt: 'asc' },
    });

  beforeAll(async () => {
    Object.assign(process.env, PREVIOUS_KEYS);
    prisma = new PrismaService();
    await prisma.$connect();

    for (const form of MFA_FORMS) {
      const plaintext = generateTotpSecret();
      created.mfa[form] = sealMfa(form, plaintext);
      mfa[form] = {
        id: await createUser(form, created.mfa[form]),
        plaintext,
      };
    }
    for (const form of WEBHOOK_FORMS) {
      const plaintext = randomBytes(32).toString('hex');
      created.webhook[form] = sealWebhook(form, plaintext);
      const subscription = await prisma.webhookSubscription.create({
        data: {
          url: 'https://receiver.example.com/rekey-e2e',
          events: ['task.created'],
          secret: created.webhook[form],
          active: false,
        },
      });
      webhook[form] = { id: subscription.id, plaintext };
    }
  });

  afterAll(async () => {
    const runs = await rekeyRuns();
    const runIds = runs.map((run) => run.id);
    await prisma.notification.deleteMany({
      where: { relatedType: 'JobRun', relatedId: { in: runIds } },
    });
    await prisma.jobRun.deleteMany({ where: { id: { in: runIds } } });
    const userIds = [
      ...MFA_FORMS.flatMap((form) => (mfa[form] ? [mfa[form].id] : [])),
      ...extraUserIds,
    ];
    await prisma.notification.deleteMany({
      where: { userId: { in: userIds } },
    });
    await prisma.user.deleteMany({ where: { id: { in: userIds } } });
    await prisma.webhookSubscription.deleteMany({
      where: {
        id: {
          in: WEBHOOK_FORMS.flatMap((form) =>
            webhook[form] ? [webhook[form].id] : [],
          ),
        },
      },
    });
    await prisma.$disconnect();
    for (const [name, value] of Object.entries(savedEnv)) {
      if (value === undefined) {
        delete process.env[name];
      } else {
        process.env[name] = value;
      }
    }
  });

  it('a second run rekeys nothing; dry-run writes nothing', async () => {
    const currentMfaKeyId = keyIdOf(encryptMfaSecret('key-id-probe'));
    const currentWebhookKeyId = keyIdOf(sealWebhookSecret('key-id-probe'));

    // Dry-run: counts, opens every value, writes none.
    expect(await runCli(['rekey', '--dry-run'])).toBe(0);
    expect(await storedValues()).toEqual(created);

    // Run: everything not sealed with the current key is resealed.
    expect(await runCli(['rekey'])).toBe(0);
    const afterFirst = await storedValues();
    for (const form of ['legacy', 'previous'] as const) {
      expect(keyIdOf(afterFirst.mfa[form])).toBe(currentMfaKeyId);
      expect(afterFirst.mfa[form]).not.toBe(created.mfa[form]);
    }
    for (const form of ['legacy', 'plaintext', 'previous'] as const) {
      expect(keyIdOf(afterFirst.webhook[form])).toBe(currentWebhookKeyId);
      expect(afterFirst.webhook[form]).not.toBe(created.webhook[form]);
    }
    expect(afterFirst.mfa.current).toBe(created.mfa.current);
    expect(afterFirst.webhook.current).toBe(created.webhook.current);

    // Second run: a reseal would draw a fresh IV, so unchanged bytes prove
    // that nothing was written.
    expect(await runCli(['rekey'])).toBe(0);
    expect(await storedValues()).toEqual(afterFirst);

    const [dryRun, first, second, ...others] = await rekeyRuns();
    expect(others).toHaveLength(0);
    for (const run of [dryRun, first, second]) {
      expect(run).toMatchObject({ status: 'SUCCEEDED', error: null });
    }
    const counts = (run: typeof first) => run.result as Record<string, number>;
    expect(counts(dryRun)).toMatchObject({ dryRun: true, 'mfa.failed': 0 });
    expect(counts(first)).toMatchObject({ dryRun: false, 'mfa.failed': 0 });
    for (const run of [dryRun, first]) {
      expect(counts(run)['mfa.rekeyed']).toBeGreaterThanOrEqual(2);
      expect(counts(run)['webhook.rekeyed']).toBeGreaterThanOrEqual(2);
      expect(counts(run)['webhook.plaintextSealed']).toBeGreaterThanOrEqual(1);
    }
    // No e2e file writes old-format MFA or sealed webhook secrets; others
    // may add plaintext webhook secrets meanwhile, so that count stays open.
    expect(counts(second)).toMatchObject({
      dryRun: false,
      'mfa.rekeyed': 0,
      'mfa.plaintextSealed': 0,
      'mfa.changedMeanwhile': 0,
      'mfa.failed': 0,
      'webhook.rekeyed': 0,
      'webhook.failed': 0,
    });
    expect(counts(second)['mfa.alreadyCurrent']).toBeGreaterThanOrEqual(3);
    expect(counts(second)['webhook.alreadyCurrent']).toBeGreaterThanOrEqual(4);
  });

  it('rekeyed values open to the original plaintext', async () => {
    // The previous keys can go now. (JWT_ACCESS_SECRET cannot change here:
    // it is also the source of the development keys.)
    for (const name of Object.keys(PREVIOUS_KEYS)) {
      delete process.env[name];
    }
    const stored = await storedValues();

    for (const form of MFA_FORMS) {
      expect(decryptMfaSecret(stored.mfa[form])).toBe(mfa[form].plaintext);
    }
    for (const form of WEBHOOK_FORMS) {
      expect(openWebhookSecret(stored.webhook[form])).toBe(
        webhook[form].plaintext,
      );
    }
  });

  it('skips a row changed since it was read and keeps the new secret', async () => {
    const userId = await createUser(
      'raced',
      legacyCiphertext(generateTotpSecret()),
    );
    extraUserIds.push(userId);
    const reenrolled = encryptMfaSecret(generateTotpSecret());
    // Between the job's read and its write, the user sets MFA up anew -
    // through the real database, so the SQL compare-and-swap is what has
    // to keep the new secret.
    const racing = {
      user: {
        findMany: (args: Prisma.UserFindManyArgs) => prisma.user.findMany(args),
        updateMany: async (args: Prisma.UserUpdateManyArgs) => {
          if (args.where?.id === userId) {
            await prisma.user.update({
              where: { id: userId },
              data: { mfaSecret: reenrolled },
            });
          }
          return prisma.user.updateMany(args);
        },
      },
      webhookSubscription: prisma.webhookSubscription,
    } as unknown as PrismaService;

    const result = await new RekeyJob(racing).run([]);

    expect(result['mfa.changedMeanwhile']).toBeGreaterThanOrEqual(1);
    expect(result['mfa.failed']).toBe(0);
    const after = await prisma.user.findUniqueOrThrow({
      where: { id: userId },
      select: { mfaSecret: true },
    });
    expect(after.mfaSecret).toBe(reenrolled);
  });
});