import { Injectable, Logger } from '@nestjs/common';
import {
  decryptMfaSecret,
  encryptMfaSecret,
  mfaSecretNeedsRekey,
} from '../auth/mfa/mfa-crypto.util';
import { loadKeyRing, SecretPurpose } from '../common/crypto/key-ring';
import { loggableError } from '../common/loggable-error';
import { PrismaService } from '../prisma/prisma.service';
import {
  isPlaintextWebhookSecret,
  openWebhookSecret,
  sealWebhookSecret,
  webhookSecretNeedsRekey,
} from '../webhooks/webhook-secret';
import { VoiceCaptureStatus } from '../../generated/prisma';
import {
  isVoiceSealedWithCurrentKey,
  openVoiceBytes,
  openVoiceText,
  sealVoiceBytes,
  sealVoiceText,
  voiceAad,
} from '../voice/core/voice-crypto';
import { Job, JobResult } from './job';
import { JobName } from './job-names';

/** `rekey --dry-run`: open and count every secret, write none. */
export const REKEY_DRY_RUN_ARG = '--dry-run';

/** Voice content of these states is already deleted (or about to be) and is not resealed. */
const FINISHED_VOICE_STATUSES: VoiceCaptureStatus[] = [
  VoiceCaptureStatus.CONFIRMED,
  VoiceCaptureStatus.DISCARDED,
  VoiceCaptureStatus.EXPIRED,
];

/** What one run did with the secrets of one purpose (mfa, webhook, voice). */
export interface RekeyCounts {
  /** Rows with a stored secret. */
  total: number;
  /** Sealed with the previous key or before the key ring, resealed. */
  rekeyed: number;
  /**
   * Already sealed with the current key, left as they are. Only a dry-run
   * opens them (a value that does not open counts as `failed`).
   */
  alreadyCurrent: number;
  /** Webhook secrets stored in plain text, now sealed. */
  plaintextSealed: number;
  /** Written by someone else since this run read them; not overwritten. */
  changedMeanwhile: number;
  /** Could not be opened or written; they keep their old value. */
  failed: number;
}

type RowOutcome = Exclude<keyof RekeyCounts, 'total'>;

interface StoredSecret {
  id: string;
  stored: string;
}

/** One column of sealed secrets, with the helpers of its purpose. */
interface SecretColumn {
  purpose: SecretPurpose;
  /** How log lines name a row: "<entity> <id>". */
  entity: string;
  read(): Promise<StoredSecret[]>;
  /**
   * Instead of read(): the ids first, then one row at a time with readOne(),
   * so large values (audio) are never held all at once. readOne gives null
   * for a row that is gone or no longer eligible.
   */
  readIds?(): Promise<string[]>;
  readOne?(id: string): Promise<StoredSecret | null>;
  /**
   * Writes `next` only while the row still holds `previous` - one UPDATE
   * with both in its WHERE, so a value written meanwhile (a new MFA setup,
   * another run) is never overwritten. False if nothing matched.
   */
  swap(id: string, previous: string, next: string): Promise<boolean>;
  needsRekey(stored: string): boolean;
  isPlaintext(stored: string): boolean;
  /** `id` binds voice content to its row (associated data); the other purposes ignore it. */
  open(stored: string, id: string): string;
  seal(plaintext: string, id: string): string;
}

/** Refuses anything but --dry-run: a typo must never start the real run. */
function isDryRun(args: readonly string[]): boolean {
  const refused = args.filter((arg) => arg !== REKEY_DRY_RUN_ARG);
  if (refused.length > 0) {
    throw new Error(
      `rekey does not accept ${refused.map((arg) => JSON.stringify(arg)).join(', ')}`,
    );
  }
  return args.length > 0;
}

function countOutcomes(outcomes: readonly RowOutcome[]): RekeyCounts {
  const count = (outcome: RowOutcome): number =>
    outcomes.filter((each) => each === outcome).length;
  return {
    total: outcomes.length,
    rekeyed: count('rekeyed'),
    alreadyCurrent: count('alreadyCurrent'),
    plaintextSealed: count('plaintextSealed'),
    changedMeanwhile: count('changedMeanwhile'),
    failed: count('failed'),
  };
}

/** Sums the counts of columns that share one purpose (voice has three). */
function sumCounts(all: readonly RekeyCounts[]): RekeyCounts {
  const sum = (key: keyof RekeyCounts): number =>
    all.reduce((total, each) => total + each[key], 0);
  return {
    total: sum('total'),
    rekeyed: sum('rekeyed'),
    alreadyCurrent: sum('alreadyCurrent'),
    plaintextSealed: sum('plaintextSealed'),
    changedMeanwhile: sum('changedMeanwhile'),
    failed: sum('failed'),
  };
}

/** `{ dryRun, 'mfa.total', …, 'webhook.failed' }` - counts only, never values. */
function toJobResult(
  dryRun: boolean,
  counts: readonly (readonly [SecretPurpose, RekeyCounts])[],
): JobResult {
  return {
    dryRun,
    ...Object.fromEntries(
      counts.flatMap(([purpose, perPurpose]) =>
        (Object.keys(perPurpose) as (keyof RekeyCounts)[]).map((key) => [
          `${purpose}.${key}`,
          perPurpose[key],
        ]),
      ),
    ),
  };
}

function describeResult(result: JobResult): string {
  return Object.entries(result)
    .map(([key, value]) => `${key}=${String(value)}`)
    .join(', ');
}

/**
 * Reseals every stored MFA (TOTP) and webhook signing secret, and the
 * content of open voice notes (transcript, suggestion, audio), with the
 * current key of its purpose (common/crypto/key-ring.ts): values sealed
 * with `*_PREVIOUS` or before the key ring (MFA without key id, webhook
 * `enc:v1:`, both under sha256(JWT_ACCESS_SECRET)), and webhook secrets
 * still stored in plain text. Run it after setting new keys, before
 * removing `*_PREVIOUS` and before rotating the JWT secrets.
 *
 * One row at a time, each write its own compare-and-swap statement, no
 * transaction around the run: an interrupted run keeps what it resealed and
 * the next run skips those rows. A row that cannot be opened or written
 * keeps its old value; the run goes on and fails at the end, with the
 * counts in its error. Logs name rows by id, never a secret or ciphertext.
 */
@Injectable()
export class RekeyJob implements Job {
  readonly name: JobName = 'rekey';

  private readonly logger = new Logger(RekeyJob.name);

  constructor(private readonly prisma: PrismaService) {}

  async run(args: readonly string[]): Promise<JobResult> {
    const dryRun = isDryRun(args);
    const columns = [
      this.mfaColumn(),
      this.webhookColumn(),
      ...this.voiceColumns(),
    ];
    // A key that is set but unusable fails the run before any row is read.
    for (const purpose of new Set(columns.map((column) => column.purpose))) {
      loadKeyRing(purpose);
    }
    const perColumn: (readonly [SecretPurpose, RekeyCounts])[] = [];
    for (const column of columns) {
      perColumn.push([column.purpose, await this.rekeyColumn(column, dryRun)]);
    }
    const counts = (['mfa', 'webhook', 'voice'] as const).map(
      (purpose) =>
        [
          purpose,
          sumCounts(
            perColumn
              .filter(([each]) => each === purpose)
              .map(([, each]) => each),
          ),
        ] as const,
    );
    const result = toJobResult(dryRun, counts);
    const failed = counts.reduce((sum, [, each]) => sum + each.failed, 0);
    if (failed > 0) {
      throw new Error(
        `rekey: ${failed} secret(s) could not be rekeyed - ${describeResult(result)}`,
      );
    }
    return result;
  }

  private async rekeyColumn(
    column: SecretColumn,
    dryRun: boolean,
  ): Promise<RekeyCounts> {
    const outcomes: RowOutcome[] = [];
    if (column.readIds && column.readOne) {
      for (const id of await column.readIds()) {
        const row = await column.readOne(id);
        if (row !== null) {
          outcomes.push(await this.rekeyRow(column, row, dryRun));
        }
      }
      return countOutcomes(outcomes);
    }
    for (const row of await column.read()) {
      outcomes.push(await this.rekeyRow(column, row, dryRun));
    }
    return countOutcomes(outcomes);
  }

  private async rekeyRow(
    column: SecretColumn,
    row: StoredSecret,
    dryRun: boolean,
  ): Promise<RowOutcome> {
    const current = !column.needsRekey(row.stored);
    // A real run skips current rows unopened; a dry-run opens them too, so
    // the closing all-zero check also catches a corrupted current value.
    if (current && !dryRun) {
      return 'alreadyCurrent';
    }
    const outcome: RowOutcome = current
      ? 'alreadyCurrent'
      : column.isPlaintext(row.stored)
        ? 'plaintextSealed'
        : 'rekeyed';
    const subject = `${column.entity} ${row.id}: ${column.purpose} secret`;
    try {
      // Opened in a dry-run too: it reports what a real run could not reseal.
      const plaintext = column.open(row.stored, row.id);
      if (dryRun) {
        return outcome;
      }
      const resealed = column.seal(plaintext, row.id);
      // The swap replaces the only copy of the old value.
      if (column.open(resealed, row.id) !== plaintext) {
        throw new Error('resealed secret does not open to the same value');
      }
      if (await column.swap(row.id, row.stored, resealed)) {
        return outcome;
      }
      this.logger.warn(
        `${subject} changed since it was read, left as it is (run rekey --dry-run to check)`,
      );
      return 'changedMeanwhile';
    } catch (error: unknown) {
      this.logger.error(
        `${subject} not rekeyed: ${loggableError(error).message}`,
      );
      return 'failed';
    }
  }

  private mfaColumn(): SecretColumn {
    return {
      purpose: 'mfa',
      entity: 'User',
      read: async () => {
        const users = await this.prisma.user.findMany({
          where: { mfaSecret: { not: null } },
          select: { id: true, mfaSecret: true },
          orderBy: { id: 'asc' },
        });
        return users.flatMap(({ id, mfaSecret }) =>
          mfaSecret === null ? [] : [{ id, stored: mfaSecret }],
        );
      },
      swap: async (id, previous, next) => {
        const { count } = await this.prisma.user.updateMany({
          where: { id, mfaSecret: previous },
          data: { mfaSecret: next },
        });
        return count === 1;
      },
      needsRekey: mfaSecretNeedsRekey,
      // TOTP secrets were never stored in plain text.
      isPlaintext: () => false,
      open: decryptMfaSecret,
      seal: encryptMfaSecret,
    };
  }

  private webhookColumn(): SecretColumn {
    return {
      purpose: 'webhook',
      entity: 'WebhookSubscription',
      read: async () => {
        const subscriptions = await this.prisma.webhookSubscription.findMany({
          select: { id: true, secret: true },
          orderBy: { id: 'asc' },
        });
        return subscriptions.map(({ id, secret }) => ({ id, stored: secret }));
      },
      swap: async (id, previous, next) => {
        const { count } = await this.prisma.webhookSubscription.updateMany({
          where: { id, secret: previous },
          data: { secret: next },
        });
        return count === 1;
      },
      needsRekey: webhookSecretNeedsRekey,
      isPlaintext: isPlaintextWebhookSecret,
      open: openWebhookSecret,
      seal: sealWebhookSecret,
    };
  }

  /**
   * Transcript, suggestion and audio of open voice notes. Audio is handled
   * as base64 text so it shares the row logic; it is read only for open
   * notes (at most a few small files).
   */
  private voiceColumns(): SecretColumn[] {
    const openStatus = { notIn: FINISHED_VOICE_STATUSES };
    const textColumn = (
      field: 'transcriptEnc' | 'suggestionEnc',
      kind: 'transcript' | 'suggestion',
    ): SecretColumn => ({
      purpose: 'voice',
      entity: 'VoiceCapture',
      read: async () => {
        const rows = await this.prisma.voiceCapture.findMany({
          where: { status: openStatus, [field]: { not: null } },
          select: { id: true, [field]: true },
          orderBy: { id: 'asc' },
        });
        return rows.flatMap((row) => {
          const stored = row[field];
          return typeof stored === 'string' ? [{ id: row.id, stored }] : [];
        });
      },
      swap: async (id, previous, next) => {
        const { count } = await this.prisma.voiceCapture.updateMany({
          where: { id, status: openStatus, [field]: previous },
          data: { [field]: next },
        });
        return count === 1;
      },
      needsRekey: (stored) => !isVoiceSealedWithCurrentKey(stored),
      isPlaintext: () => false,
      open: (stored, id) => openVoiceText(stored, voiceAad(id, kind)),
      seal: (plaintext, id) => sealVoiceText(plaintext, voiceAad(id, kind)),
    });
    const audioColumn: SecretColumn = {
      purpose: 'voice',
      entity: 'VoiceCapture',
      read: () => Promise.resolve([]),
      readIds: async () => {
        const rows = await this.prisma.voiceCaptureAudio.findMany({
          where: { capture: { status: openStatus } },
          select: { captureId: true },
          orderBy: { captureId: 'asc' },
        });
        return rows.map((row) => row.captureId);
      },
      readOne: async (id) => {
        const row = await this.prisma.voiceCaptureAudio.findFirst({
          where: { captureId: id, capture: { status: openStatus } },
          select: { captureId: true, dataEnc: true },
        });
        return row
          ? {
              id: row.captureId,
              stored: Buffer.from(row.dataEnc).toString('base64'),
            }
          : null;
      },
      swap: async (id, previous, next) => {
        const { count } = await this.prisma.voiceCaptureAudio.updateMany({
          where: {
            captureId: id,
            dataEnc: Buffer.from(previous, 'base64'),
            capture: { status: openStatus },
          },
          data: { dataEnc: Buffer.from(next, 'base64') },
        });
        return count === 1;
      },
      needsRekey: (stored) =>
        !isVoiceSealedWithCurrentKey(Buffer.from(stored, 'base64')),
      isPlaintext: () => false,
      open: (stored, id) =>
        openVoiceBytes(
          Buffer.from(stored, 'base64'),
          voiceAad(id, 'audio'),
        ).toString('base64'),
      seal: (plaintext, id) =>
        sealVoiceBytes(
          Buffer.from(plaintext, 'base64'),
          voiceAad(id, 'audio'),
        ).toString('base64'),
    };
    return [
      textColumn('transcriptEnc', 'transcript'),
      textColumn('suggestionEnc', 'suggestion'),
      audioColumn,
    ];
  }
}