import { z } from "zod";
import { eq } from "drizzle-orm";
import { headers } from "next/headers";
import { generateSecret, generateURI, verifySync } from "otplib";
import QRCode from "qrcode";
import { db, schema } from "@/db";
import { getSession, markMfaPassed, rotateSession, setSessionCookie, clientIp, revokeOtherSessions } from "@/lib/auth/session";
import { assertSameOrigin, jsonError, HttpError, requireUser } from "@/lib/auth/guard";
import { rateLimit } from "@/lib/auth/ratelimit";
import { audit } from "@/lib/audit";

/** POST {step:"verify", code} — complete login MFA.  POST {step:"setup"} → secret+QR.  POST {step:"enable", code, secret}.  POST {step:"disable", code}. */
export async function POST(req: Request) {
  try {
    await assertSameOrigin();
    const ip = clientIp(await headers());
    if (!rateLimit(`totp:${ip}`, 10, 10 * 60_000)) throw new HttpError(429, "Too many attempts");
    const b = z.object({ step: z.enum(["verify", "setup", "enable", "disable"]), code: z.string().regex(/^\d{6}$/).optional(), secret: z.string().max(100).optional() }).parse(await req.json());
    if (b.step === "verify") {
      const s = await getSession(); if (!s) throw new HttpError(401, "Unauthenticated");
      const [u] = await db.select().from(schema.users).where(eq(schema.users.id, s.id));
      if (!u.totpEnabled || !u.totpSecret) throw new HttpError(400, "MFA not enabled");
      if (!b.code || !verifySync({ token: b.code, secret: u.totpSecret }).valid) { await audit({ id: u.id, name: u.email }, "auth.mfa_failed", undefined, {}, ip); throw new HttpError(401, "Invalid code"); }
      await markMfaPassed(s.sessionId); const t = await rotateSession(s.sessionId); await setSessionCookie(t);
      await audit({ id: u.id, name: u.email }, "auth.mfa_passed", undefined, {}, ip);
      return Response.json({ ok: true, next: u.mustChangePassword ? "change_password" : "done" });
    }
    const s = await requireUser(req);
    if (b.step === "setup") {
      const secret = generateSecret(); const uri = generateURI({ issuer: "tools.spboucher.ai", label: s.email, secret });
      return Response.json({ secret, uri, qr: await QRCode.toDataURL(uri, { margin: 1, width: 220 }) });
    }
    if (b.step === "enable") {
      if (!b.code || !b.secret || !verifySync({ token: b.code, secret: b.secret }).valid) throw new HttpError(400, "Code does not match the secret");
      await db.update(schema.users).set({ totpSecret: b.secret, totpEnabled: true, updatedAt: new Date() }).where(eq(schema.users.id, s.id));
      await markMfaPassed(s.sessionId); await revokeOtherSessions(s.id, s.sessionId);
      await audit({ id: s.id, name: s.email }, "auth.totp_enabled", undefined, {}, ip);
      return Response.json({ ok: true });
    }
    const [u] = await db.select().from(schema.users).where(eq(schema.users.id, s.id));
    if (!b.code || !u.totpSecret || !verifySync({ token: b.code, secret: u.totpSecret }).valid) throw new HttpError(400, "Invalid code");
    await db.update(schema.users).set({ totpSecret: null, totpEnabled: false, updatedAt: new Date() }).where(eq(schema.users.id, s.id));
    await audit({ id: s.id, name: s.email }, "auth.totp_disabled", undefined, {}, ip);
    return Response.json({ ok: true });
  } catch (e) { return jsonError(e); }
}