import "server-only";
import { db, findById, withLock, type RecordModel } from "./db";
import { hashRecoveryCode, matchTotp, newRecoveryCodes, newTotpSecret, normalizeCode } from "./totp";

// Two-factor authentication with an authenticator app. With it on, signing in in the browser takes
// a code from the app (or a recovery code) after the password. Git, the registry and agents don't
// use passwords (they use SSH keys, API tokens and OAuth), so they aren't affected.

export type TwoFactorStatus = {
  enabled: boolean;
  enabledAt: number | null;
  recoveryCodesLeft: number;
  /** A secret being set up, not confirmed with a code yet. */
  pendingSecret: string | null;
};

export class TwoFactorError extends Error {}

const recoveryHashes = (r: RecordModel) => (Array.isArray(r.recovery_codes) ? (r.recovery_codes as string[]) : []);

async function userRecord(userId: string) {
  const record = await findById("users", userId);
  if (!record) throw new TwoFactorError("Account not found.");
  return record;
}

export async function twoFactorStatus(userId: string): Promise<TwoFactorStatus> {
  const r = await userRecord(userId);
  return {
    enabled: !!r.totp_secret,
    enabledAt: r.totp_enabled_at || null,
    recoveryCodesLeft: recoveryHashes(r).length,
    pendingSecret: (!r.totp_secret && r.totp_pending) || null,
  };
}

export async function hasTwoFactor(userId: string) {
  return !!(await findById("users", userId))?.totp_secret;
}

const update = async (userId: string, data: Record<string, unknown>) => (await db()).collection("users").update(userId, data);
const lock = <T>(userId: string, fn: () => Promise<T>) => withLock(`2fa:${userId}`, fn);

/** Starts setting up an app: a new secret, kept aside until a code from the app confirms it. */
export function startTwoFactorSetup(userId: string) {
  return lock(userId, async () => {
    if ((await userRecord(userId)).totp_secret) throw new TwoFactorError("Two-factor authentication is already on.");
    await update(userId, { totp_pending: newTotpSecret() });
  });
}

export function cancelTwoFactorSetup(userId: string) {
  return lock(userId, () => update(userId, { totp_pending: "" }));
}

/** Turns 2FA on once a code from the app matches the secret being set up. Returns the recovery codes. */
export function enableTwoFactor(userId: string, code: string) {
  return lock(userId, async () => {
    const r = await userRecord(userId);
    if (r.totp_secret) throw new TwoFactorError("Two-factor authentication is already on.");
    if (!r.totp_pending) throw new TwoFactorError("Start the setup again.");
    const step = matchTotp(r.totp_pending, normalizeCode(code));
    if (step === null) throw new TwoFactorError("That code didn't match. Check your device's time and try the newest code.");
    const codes = newRecoveryCodes();
    await update(userId, {
      totp_secret: r.totp_pending,
      totp_pending: "",
      totp_enabled_at: Date.now(),
      totp_last_step: step,
      recovery_codes: codes.map(hashRecoveryCode),
    });
    return codes;
  });
}

/**
 * Checks a code from the app or a recovery code, and uses it up: an app code can't be used again
 * (nor an older one), a recovery code is deleted.
 */
export function verifySecondFactor(userId: string, input: string) {
  return lock(userId, async () => {
    const r = await findById("users", userId);
    if (!r?.totp_secret) return false;
    const code = normalizeCode(input);
    if (/^\d{6}$/.test(code)) {
      const step = matchTotp(r.totp_secret, code);
      if (step === null || step <= (r.totp_last_step || 0)) return false;
      await update(userId, { totp_last_step: step });
      return true;
    }
    const hashes = recoveryHashes(r);
    const hash = hashRecoveryCode(code);
    if (!code || !hashes.includes(hash)) return false;
    await update(userId, { recovery_codes: hashes.filter((h) => h !== hash) });
    return true;
  });
}

export async function disableTwoFactor(userId: string, code: string) {
  if (!(await verifySecondFactor(userId, code))) throw new TwoFactorError("That code didn't match.");
  await lock(userId, () => update(userId, { totp_secret: "", totp_pending: "", totp_enabled_at: 0, totp_last_step: 0, recovery_codes: [] }));
}

/** New recovery codes in place of the old ones, after a code from the app. */
export async function regenerateRecoveryCodes(userId: string, code: string) {
  if (!/^\d{6}$/.test(normalizeCode(code))) throw new TwoFactorError("Enter a code from your authenticator app.");
  if (!(await verifySecondFactor(userId, code))) throw new TwoFactorError("That code didn't match.");
  const codes = newRecoveryCodes();
  await lock(userId, () => update(userId, { recovery_codes: codes.map(hashRecoveryCode) }));
  return codes;
}