import { env } from "cloudflare:workers";
import { deleteCookie, getCookie, setCookie } from "@tanstack/react-start/server";
import { all, batch, first, run, stmt } from "./db";
import { decryptSecret, encryptSecret, randomToken, sha256Hex, toHex } from "./crypto";
import { attemptAuthToken, deleteAuthToken, peekAuthToken, storeAuthToken } from "./auth-tokens";
import { badRequest } from "./errors";
import { cookieAttributes, cookieName, cookieSecure } from "./auth";
import { generateTotpSecret, totpUri, verifyTotp } from "./totp";

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;
}

export async function secondFactorMethods(userId: string): Promise<{ totp: boolean; passkey: boolean }> {
  const row = await first<{ totp: number | null; passkeys: number }>(
    `SELECT u.totp_enabled_at AS totp, (SELECT COUNT(*) FROM passkeys p WHERE p.user_id = u.account_id) AS passkeys
     FROM users u WHERE u.account_id = ?`,
    userId,
  );
  return { totp: typeof row?.totp === "number", passkey: (row?.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);
  await run(
    "UPDATE users SET totp_secret = ?, totp_iv = ?, totp_enabled_at = ?, totp_last_step = ? WHERE account_id = ?",
    encrypted.ciphertext,
    encrypted.iv,
    Date.now(),
    step,
    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 batch([
    stmt(
      "UPDATE users SET totp_secret = NULL, totp_iv = NULL, totp_enabled_at = NULL, totp_last_step = NULL WHERE account_id = ?",
      userId,
    ),
    stmt(
      "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 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 all<{ count: number }>(
    "SELECT COUNT(*) AS count FROM recovery_codes WHERE user_id = ? AND used_at IS NULL",
    userId,
  );
  return rows[0]?.count ?? 0;
}

/** 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 }>(
    "SELECT totp_secret, totp_iv, totp_last_step FROM users WHERE account_id = ? AND totp_enabled_at IS NOT NULL",
    userId,
  );
  if (!row?.totp_secret || !row.totp_iv) return false;
  const secret = await decryptSecret(masterKey(), row.totp_secret, row.totp_iv);
  const step = await verifyTotp(secret, code, row.totp_last_step);
  if (step === null) return false;
  const result = await run(
    "UPDATE users SET totp_last_step = ? WHERE account_id = ? AND (totp_last_step IS NULL OR totp_last_step < ?)",
    step,
    userId,
    step,
  );
  return result.meta.changes === 1;
}

export async function useRecoveryCode(userId: string, code: string): Promise<boolean> {
  const result = await 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)),
  );
  return result.meta.changes === 1;
}