import { env } from "cloudflare:workers";
import { deleteCookie, getCookie, setCookie } from "@tanstack/react-start/server";
import { first, run } from "./db";
import { accountDb, stmt } from "./scoped-db";
import { decryptSecret, encryptSecret, isLegacySecret, randomToken, sha256Hex, toHex } from "./crypto";
import type { SealContext } from "./crypto";
import { attemptAuthToken, deleteAuthToken, peekAuthToken, storeAuthToken } from "./auth-tokens";
import { badRequest, HttpError } from "./errors";
import { notifySecurityEvent } from "./account-security";
import { cookieAttributes, cookieName, cookieSecure } from "./auth";
import { generateTotpSecret, totpUri, verifyTotp } from "./totp";
import { userChanged } from "./user-row";

const PENDING_COOKIE_BASE = "gitfor_2fa";
const PENDING_TTL_SECONDS = 5 * 60;
const SETUP_TTL_SECONDS = 10 * 60;
const MAX_PENDING_ATTEMPTS = 5;
const RECOVERY_CODE_COUNT = 10;
const RECOVERY_HALF_BYTES = 3;

function masterKey(): string {
  if (!env.SECRETS_KEY) throw badRequest("SECRETS_KEY must be configured to use two-factor authentication");
  return env.SECRETS_KEY;
}

/** TOTP seeds use their own derived key and are bound to the user row, so a seed cannot be moved between users (AUTH-A10). */
function totpContext(userId: string): SealContext {
  return { purpose: "totp", aad: `user:${userId}` };
}

export async function secondFactorMethods(userId: string): Promise<{ totp: boolean; passkey: boolean }> {
  const [row, keys] = await Promise.all([
    first<{ totp: number | null }>("SELECT totp_enabled_at AS totp FROM users WHERE account_id = ?", userId),
    accountDb(userId).first<{ passkeys: number }>(
      "SELECT COUNT(*) AS passkeys FROM passkeys WHERE user_id = ?",
      userId,
    ),
  ]);
  return { totp: typeof row?.totp === "number", passkey: (keys?.passkeys ?? 0) > 0 };
}

export async function requiresSecondFactor(userId: string): Promise<boolean> {
  const methods = await secondFactorMethods(userId);
  return methods.totp || methods.passkey;
}

async function storePending(userId: string): Promise<string> {
  const token = randomToken("mfa");
  await storeAuthToken("pending-login", token, userId, "{}", PENDING_TTL_SECONDS);
  return token;
}

/** Records a password-verified login awaiting its second factor, bound to this browser by cookie. */
export async function startPendingLogin(userId: string): Promise<void> {
  const token = await storePending(userId);
  setCookie(cookieName(PENDING_COOKIE_BASE), token, {
    httpOnly: true,
    secure: cookieSecure(),
    sameSite: "lax",
    path: "/",
    maxAge: PENDING_TTL_SECONDS,
  });
}

export async function pendingLoginCookie(userId: string): Promise<string> {
  const token = await storePending(userId);
  return `${cookieName(PENDING_COOKIE_BASE)}=${token}; ${cookieAttributes(PENDING_TTL_SECONDS)}`;
}

async function stillValid(userId: string, createdAt: number): Promise<boolean> {
  const user = await first<{ sessions_revoked_at: number | null }>(
    "SELECT sessions_revoked_at FROM users WHERE account_id = ?",
    userId,
  );
  return user !== null && (user.sessions_revoked_at === null || createdAt >= user.sessions_revoked_at);
}

/** Read-only lookup of the pending login (for rendering the challenge page). */
export async function peekPendingLogin(): Promise<string | null> {
  const token = getCookie(cookieName(PENDING_COOKIE_BASE));
  if (!token) return null;
  const pending = await peekAuthToken("pending-login", token, MAX_PENDING_ATTEMPTS);
  if (!pending?.userId || !(await stillValid(pending.userId, pending.createdAt))) return null;
  return pending.userId;
}

/** Atomically counts a verification attempt; returns the user id while attempts remain. */
export async function attemptPendingLogin(): Promise<string | null> {
  const token = getCookie(cookieName(PENDING_COOKIE_BASE));
  if (!token) return null;
  const pending = await attemptAuthToken("pending-login", token, MAX_PENDING_ATTEMPTS);
  if (!pending?.userId || !(await stillValid(pending.userId, pending.createdAt))) return null;
  return pending.userId;
}

export async function clearPendingLogin(): Promise<void> {
  const token = getCookie(cookieName(PENDING_COOKIE_BASE));
  if (token) await deleteAuthToken("pending-login", token);
  deleteCookie(cookieName(PENDING_COOKIE_BASE), { path: "/", secure: cookieSecure() });
}

export async function startTotpSetup(userId: string, accountName: string): Promise<{ secret: string; uri: string }> {
  masterKey();
  const secret = generateTotpSecret();
  await env.KV.put(`totp-setup:${userId}`, secret, { expirationTtl: SETUP_TTL_SECONDS });
  return { secret, uri: totpUri(env.FORGE_NAME || "gitfor", accountName, secret) };
}

/** Confirms the pending secret with a code from the authenticator and returns fresh recovery codes. */
export async function enableTotp(userId: string, code: string): Promise<string[] | null> {
  const secret = await env.KV.get(`totp-setup:${userId}`);
  if (!secret) throw badRequest("Setup expired; start again");
  const step = await verifyTotp(secret, code, null);
  if (step === null) throw badRequest("That code is not valid");
  const encrypted = await encryptSecret(masterKey(), secret, totpContext(userId));
  await run(
    "UPDATE users SET totp_secret = ?, totp_iv = ?, totp_enabled_at = ?, totp_last_step = ?, totp_failures = 0, totp_locked_until = NULL WHERE account_id = ?",
    encrypted.ciphertext,
    encrypted.iv,
    Date.now(),
    step,
    userId,
  );
  await userChanged(userId);
  await env.KV.delete(`totp-setup:${userId}`);
  return (await remainingRecoveryCodes(userId)) > 0 ? null : regenerateRecoveryCodes(userId);
}

/** Recovery codes stay while passkeys remain, since they still act as a second factor. */
export async function disableTotp(userId: string): Promise<void> {
  await run(
    "UPDATE users SET totp_secret = NULL, totp_iv = NULL, totp_enabled_at = NULL, totp_last_step = NULL, totp_failures = 0, totp_locked_until = NULL WHERE account_id = ?",
    userId,
  );
  await userChanged(userId);
  await accountDb(userId).run(
    "DELETE FROM recovery_codes WHERE user_id = ?1 AND NOT EXISTS (SELECT 1 FROM passkeys WHERE user_id = ?1)",
    userId,
  );
}

function formatRecoveryCode(): string {
  const bytes = crypto.getRandomValues(new Uint8Array(RECOVERY_HALF_BYTES * 2));
  const hex = toHex(bytes);
  return `${hex.slice(0, RECOVERY_HALF_BYTES * 2)}-${hex.slice(RECOVERY_HALF_BYTES * 2)}`;
}

function normalizeRecoveryCode(code: string): string {
  const compact = code.toLowerCase().replace(/[^0-9a-f]/g, "");
  return `${compact.slice(0, RECOVERY_HALF_BYTES * 2)}-${compact.slice(RECOVERY_HALF_BYTES * 2)}`;
}

export async function regenerateRecoveryCodes(userId: string): Promise<string[]> {
  const codes = Array.from({ length: RECOVERY_CODE_COUNT }, formatRecoveryCode);
  const hashes = await Promise.all(codes.map((code) => sha256Hex(code)));
  await accountDb(userId).batch([
    stmt("DELETE FROM recovery_codes WHERE user_id = ?", userId),
    ...hashes.map((hash) => stmt("INSERT INTO recovery_codes (user_id, code_hash) VALUES (?, ?)", userId, hash)),
  ]);
  return codes;
}

export async function remainingRecoveryCodes(userId: string): Promise<number> {
  const rows = await accountDb(userId).all<{ count: number }>(
    "SELECT COUNT(*) AS count FROM recovery_codes WHERE user_id = ? AND used_at IS NULL",
    userId,
  );
  return rows[0]?.count ?? 0;
}

export const TOTP_LOCKOUT_THRESHOLD = 10;
const TOTP_LOCKOUT_BASE_MS = 5 * 60 * 1000;
const TOTP_LOCKOUT_MAX_MS = 24 * 60 * 60 * 1000;
const HTTP_TOO_MANY = 429;

/**
 * Wrong codes are counted on the user row, so the five-attempt limit of a pending login cannot be reset by
 * signing in with the password again (AUTH-P4). From the threshold on, every wrong code locks TOTP for a
 * period that doubles each time (capped at a day); count and lock are one statement so parallel guesses
 * cannot slip between them.
 */
export async function recordTotpFailure(userId: string): Promise<void> {
  const now = Date.now();
  const row = await first<{ totp_failures: number; totp_locked_until: number | null }>(
    `UPDATE users SET totp_failures = totp_failures + 1,
       totp_locked_until = CASE WHEN totp_failures + 1 >= ?1
         THEN ?2 + MIN(?3 << MIN(totp_failures + 1 - ?1, 20), ?4) ELSE totp_locked_until END
     WHERE account_id = ?5 RETURNING totp_failures, totp_locked_until`,
    TOTP_LOCKOUT_THRESHOLD,
    now,
    TOTP_LOCKOUT_BASE_MS,
    TOTP_LOCKOUT_MAX_MS,
    userId,
  );
  if (!row || row.totp_failures < TOTP_LOCKOUT_THRESHOLD || row.totp_locked_until === null) return;
  await notifySecurityEvent(
    userId,
    `two-factor sign-in was locked for ${Math.round((row.totp_locked_until - now) / 60000)} minutes after ${row.totp_failures} wrong authenticator codes; if this was not you, change your password`,
  );
}

/** Any successful second factor ends the failure streak. */
export async function clearTotpFailures(userId: string): Promise<void> {
  await run(
    "UPDATE users SET totp_failures = 0, totp_locked_until = NULL WHERE account_id = ? AND totp_failures > 0",
    userId,
  );
}

/** Verifies a TOTP code with replay protection; the step update is conditional so a code works once. */
export async function verifyUserTotp(userId: string, code: string): Promise<boolean> {
  const row = await first<{
    totp_secret: string | null;
    totp_iv: string | null;
    totp_last_step: number | null;
    totp_locked_until: number | null;
  }>(
    "SELECT totp_secret, totp_iv, totp_last_step, totp_locked_until FROM users WHERE account_id = ? AND totp_enabled_at IS NOT NULL",
    userId,
  );
  if (!row?.totp_secret || !row.totp_iv) return false;
  if (row.totp_locked_until !== null && row.totp_locked_until > Date.now()) {
    throw new HttpError(HTTP_TOO_MANY, "Too many wrong codes; try again later or use a recovery code or passkey");
  }
  const secret = await decryptSecret(masterKey(), row.totp_secret, row.totp_iv, totpContext(userId));
  if (isLegacySecret(row.totp_secret)) {
    const sealed = await encryptSecret(masterKey(), secret, totpContext(userId));
    await run(
      "UPDATE users SET totp_secret = ?, totp_iv = ? WHERE account_id = ? AND totp_secret = ?",
      sealed.ciphertext,
      sealed.iv,
      userId,
      row.totp_secret,
    );
  }
  const step = await verifyTotp(secret, code, row.totp_last_step);
  if (step === null) {
    await recordTotpFailure(userId);
    return false;
  }
  const result = await run(
    `UPDATE users SET totp_last_step = ?1, totp_failures = 0, totp_locked_until = NULL
     WHERE account_id = ?2 AND (totp_last_step IS NULL OR totp_last_step < ?1)
       AND (totp_locked_until IS NULL OR totp_locked_until <= ?3)`,
    step,
    userId,
    Date.now(),
  );
  return result.meta.changes === 1;
}

export async function useRecoveryCode(userId: string, code: string): Promise<boolean> {
  const result = await accountDb(userId).run(
    "UPDATE recovery_codes SET used_at = ? WHERE user_id = ? AND code_hash = ? AND used_at IS NULL",
    Date.now(),
    userId,
    await sha256Hex(normalizeRecoveryCode(code)),
  );
  if (result.meta.changes !== 1) return false;
  await clearTotpFailures(userId);
  return true;
}