assessment-model/src/app/api/auth/[...nextauth]/authOptions.ts

import { NextAuthOptions } from "next-auth";
import GoogleProvider from "next-auth/providers/google";
import AzureADB2CProvider from "next-auth/providers/azure-ad-b2c";
import EmailProvider from "next-auth/providers/email";
import CredentialsProvider from "next-auth/providers/credentials";
import DrizzleEmailAdapter from "./DrizzleEmailAdapter";
import { MagicLinksEmail } from "@/app/email_templates/magic_link";
import {
  evaluateCodeAttempt,
  generateCode,
  hashCode,
} from "@/app/lib/verificationCode";
import { createHash } from "crypto";

import { db } from "@/app/db/db";
import {
  user as users,
  accounts,
  authRateLimits,
  verificationTokens,
} from "@/app/db/schema/users";
import { normaliseEmail } from "@/app/lib/email";
import { eq, and, ne } from "drizzle-orm";

// ------------------------------------------------------------------
// Environment variables
// ------------------------------------------------------------------
const {
  GOOGLE_CLIENT_ID = "",
  GOOGLE_CLIENT_SECRET = "",
  AZURE_AD_B2C_TENANT_NAME = "",
  AZURE_AD_B2C_CLIENT_ID = "",
  AZURE_AD_B2C_CLIENT_SECRET = "",
  AZURE_AD_B2C_PRIMARY_USER_FLOW = "",
  EMAIL_SERVER_HOST = "",
  EMAIL_SERVER_PORT = "",
  EMAIL_SERVER_USER = "",
  EMAIL_SERVER_PASSWORD = "",
  EMAIL_FROM = "",
} = process.env;

type OauthProvider = "google" | "azure-ad-b2c";

// ------------------------------------------------------------------
// NextAuth configuration
// ------------------------------------------------------------------
export const AuthOptions: NextAuthOptions = {
  adapter: DrizzleEmailAdapter(db, {
    user: users,
    accounts,
    verificationTokens,
  }),

  providers: [
    // ------------------ Google ------------------
    GoogleProvider({
      clientId: GOOGLE_CLIENT_ID,
      clientSecret: GOOGLE_CLIENT_SECRET,
      authorization: {
        params: {
          access_type: "offline",
          prompt: "consent",
          response_type: "code",
        },
      },
    }),

    // ------------------ Azure AD B2C ------------------
    AzureADB2CProvider({
      tenantId: AZURE_AD_B2C_TENANT_NAME,
      clientId: AZURE_AD_B2C_CLIENT_ID,
      clientSecret: AZURE_AD_B2C_CLIENT_SECRET,
      primaryUserFlow: AZURE_AD_B2C_PRIMARY_USER_FLOW,
      authorization: {
        params: {
          scope: "openid profile offline_access",
          prompt: "login",
        },
      },
    }),

    // ------------------ Email (SES Magic Link) ------------------
    EmailProvider({
      server: {
        host: EMAIL_SERVER_HOST,
        port: Number(EMAIL_SERVER_PORT),
        auth: {
          user: EMAIL_SERVER_USER,
          pass: EMAIL_SERVER_PASSWORD,
        },
      },
      from: EMAIL_FROM,
      maxAge: 60 * 10, // code and link share a 10-minute window
      sendVerificationRequest: async ({ identifier, url, provider, token }) => {
        const secret = process.env.NEXTAUTH_SECRET!;
        const hashedToken = createHash("sha256")
          .update(`${token}${secret}`)
          .digest("hex");
        const now = new Date();
        const oneHourMs = 60 * 60 * 1000;
        const SEND_LIMIT = 5;

        // Per-email send rate limit, fixed 1-hour window.
        const limited = await db.transaction(async (tx) => {
          const [existing] = await tx
            .select()
            .from(authRateLimits)
            .where(
              and(
                eq(authRateLimits.scope, "send"),
                eq(authRateLimits.key, identifier),
              ),
            );

          const inWindow =
            existing &&
            now.getTime() - existing.windowStart.getTime() < oneHourMs;

          if (inWindow) {
            if (existing.count >= SEND_LIMIT) return true;
            await tx
              .update(authRateLimits)
              .set({ count: existing.count + 1 })
              .where(
                and(
                  eq(authRateLimits.scope, "send"),
                  eq(authRateLimits.key, identifier),
                ),
              );
            return false;
          }

          await tx
            .insert(authRateLimits)
            .values({
              scope: "send",
              key: identifier,
              count: 1,
              windowStart: now,
            })
            .onConflictDoUpdate({
              target: [authRateLimits.scope, authRateLimits.key],
              set: { count: 1, windowStart: now },
            });
          return false;
        });

        if (limited) {
          await db
            .delete(verificationTokens)
            .where(
              and(
                eq(verificationTokens.identifier, identifier),
                eq(verificationTokens.token, hashedToken),
              ),
            );
          console.warn("EMAIL_RATE_LIMIT_EXCEEDED", {
            email: identifier,
            timestamp: now.toISOString(),
          });
          throw new Error(
            "Too many sign-in attempts. Please wait an hour before requesting another code.",
          );
        }

        // Generate code, attach to the just-created row, replace any older rows
        // for this identifier so only the latest send is live.
        const code = generateCode();
        const codeHash = hashCode(code, secret);

        await db.transaction(async (tx) => {
          await tx
            .update(verificationTokens)
            .set({ codeHash, attempts: 0 })
            .where(
              and(
                eq(verificationTokens.identifier, identifier),
                eq(verificationTokens.token, hashedToken),
              ),
            );
          await tx
            .delete(verificationTokens)
            .where(
              and(
                eq(verificationTokens.identifier, identifier),
                ne(verificationTokens.token, hashedToken),
              ),
            );
        });

        try {
          const { messageId } = await MagicLinksEmail({
            identifier,
            url,
            provider,
            code,
          });
          console.log("EMAIL_MAGIC_LINK_SUCCESS", {
            email: identifier,
            messageId,
            timestamp: new Date().toISOString(),
          });
        } catch (err) {
          console.error("EMAIL_MAGIC_LINK_FAILURE", {
            email: identifier,
            error: err instanceof Error ? err.message : String(err),
            timestamp: new Date().toISOString(),
          });
          throw err;
        }
      },
    }),

    // ------------------ Email code (typed at /auth/verify-code) ------------------
    CredentialsProvider({
      id: "email-code",
      name: "Email Code",
      credentials: {
        email: { label: "Email", type: "email" },
        code: { label: "Code", type: "text" },
      },
      async authorize(credentials) {
        if (!credentials?.email || !credentials?.code) return null;
        const email = credentials.email.trim().toLowerCase();
        const submitted = credentials.code.trim();
        if (!/^\d{6}$/.test(submitted)) return null;

        const secret = process.env.NEXTAUTH_SECRET!;
        const submittedCodeHash = hashCode(submitted, secret);
        const now = new Date();

        return await db.transaction(async (tx) => {
          const [row] = await tx
            .select()
            .from(verificationTokens)
            .where(eq(verificationTokens.identifier, email))
            .limit(1);

          const outcome = evaluateCodeAttempt({
            submittedCodeHash,
            row: row
              ? {
                  codeHash: row.codeHash,
                  attempts: row.attempts,
                  expires: row.expires,
                }
              : null,
            now,
          });

          if (outcome.outcome === "ok") {
            await tx
              .delete(verificationTokens)
              .where(eq(verificationTokens.identifier, email));

            let [dbUser] = await tx
              .select()
              .from(users)
              .where(eq(users.email, email));

            if (!dbUser) {
              [dbUser] = await tx
                .insert(users)
                .values({ email, emailVerified: now })
                .returning();
            } else if (!dbUser.emailVerified) {
              await tx
                .update(users)
                .set({ emailVerified: now })
                .where(eq(users.id, dbUser.id));
            }

            console.log("EMAIL_CODE_SIGN_IN_SUCCESS", {
              email,
              userId: String(dbUser.id),
              timestamp: now.toISOString(),
            });

            return {
              id: String(dbUser.id),
              email: dbUser.email,
              name: dbUser.firstName ?? null,
              dbId: String(dbUser.id),
              onboarded: dbUser.onboarded ?? false,
            };
          }

          if (outcome.outcome === "wrong") {
            await tx
              .update(verificationTokens)
              .set({ attempts: outcome.newAttempts })
              .where(eq(verificationTokens.identifier, email));
            return null;
          }

          if (
            outcome.outcome === "locked-out" ||
            outcome.outcome === "expired"
          ) {
            await tx
              .delete(verificationTokens)
              .where(eq(verificationTokens.identifier, email));
            return null;
          }

          return null;
        });
      },
    }),
  ],

  // ------------------------------------------------------------------
  // Pages
  // ------------------------------------------------------------------
  pages: {
    signIn: "/",
    verifyRequest: "/auth/verify-code",
    error: "/api/auth/error",
  },

  // ------------------------------------------------------------------
  // Callbacks
  // ------------------------------------------------------------------
  callbacks: {
    /**
     * Sign in callback — ensures user exists and links OAuth provider if needed
     */
    async signIn({ user, account, profile }) {
      try {
        if (!user?.email) return false;
        const normalisedEmail = user.email.toLowerCase();

        console.log("SignIn attempt:", {
          email: normalisedEmail,
          provider: account?.provider,
          type: account?.type,
          timestamp: new Date().toISOString(),
        });

        // Fetch the user (NextAuth will have created them already if new)
        const [dbUser] = await db
          .select()
          .from(users)
          .where(eq(users.email, normalisedEmail));

        // New user - next auth will handle
        if (!dbUser) {
          console.log("New user created through sign-in:", {
            email: normalisedEmail,
            provider: account?.provider,
          });
          return true;
        }

        // Auto-link provider if same verified email but account not linked yet
        if (account?.provider && account.type === "oauth") {
          const existingLink = await db
            .select()
            .from(accounts)
            .where(
              and(
                eq(accounts.userId, dbUser.id),
                eq(accounts.provider, account.provider),
              ),
            );

          const emailVerified =
            (profile as any)?.email_verified ?? account.provider === "google";

          if (existingLink.length === 0 && emailVerified) {
            // This handles the case where we had not set up accounts but
            // signed up users with oauth
            console.log(
              `Linking ${account.provider} account for user ${normalisedEmail}`,
            );

            await db
              .insert(accounts)
              .values({
                userId: dbUser.id,
                type: account.type,
                provider: account.provider,
                providerAccountId: account.providerAccountId,
                access_token: account.access_token,
                id_token: account.id_token,
                refresh_token: account.refresh_token,
                expires_at: account.expires_at,
              })
              .onConflictDoNothing();
          }
        }

        // Link OAuth ID if missing (helps for older accounts). Only for OAuth
        // providers — the email-code credentials flow has no provider identity.
        if (account?.type === "oauth" && !dbUser.oauthId) {
          console.log("Backfilling OAuth ID:", {
            email: normalisedEmail,
            provider: account.provider,
          });
          const provider = account.provider as OauthProvider;
          await db
            .update(users)
            .set({ oauthId: user.id, oauthProvider: provider })
            .where(eq(users.email, normalisedEmail));
        }

        // Always update last login timestamp
        await db
          .update(users)
          .set({ lastLogin: new Date() })
          .where(eq(users.id, dbUser.id));

        // Pending portfolio invitations are NOT auto-applied here anymore.
        // The invitee accepts/declines explicitly via the profile-menu
        // notifications panel (POST /api/user/invitations).

        // Pass bigint ID into NextAuth session/jwt
        user.dbId = dbUser.id.toString();
        user.onboarded = dbUser.onboarded ?? false;
        return true;
      } catch (error) {
        console.error("Error during sign-in:", error);
        return false;
      }
    },

    /**
     * Persist dbId in the JWT so it’s available in sessions
     */
    async jwt({ token, user: userFromLogin }) {
      // Initial sign-in: attach DB fields
      if (userFromLogin) {
        console.log("JWT initial login:", {
          email: userFromLogin.email,
          timestamp: new Date().toISOString(),
        });
        const existing = await db.query.user.findFirst({
          where: eq(users.email, userFromLogin.email!),
        });
        if (existing) {
          token.onboarded = existing.onboarded;
        }
      } else if (token.email) {
        // On subsequent calls, keep token synced from DB
        const existing = await db.query.user.findFirst({
          where: eq(users.email, token.email),
        });
        if (existing) {
          token.onboarded = existing.onboarded;
        }
      }

      if (userFromLogin?.dbId) {
        token.dbId = userFromLogin.dbId;
      }
      return token;
    },

    /**
     * Attach dbId to session.user, and normalise the email so downstream
     * lookups against `user.email` are case-insensitive without each call site
     * remembering to lowercase.
     */
    async session({ session, token }) {
      if (session.user) {
        if (session.user.email) {
          session.user.email = normaliseEmail(session.user.email);
        }
        if (token.dbId) {
          session.user.dbId = token.dbId;
        }
      }
      return session;
    },

    /**
     * Redirect users after login
     */
    async redirect({ url, baseUrl }) {
      // Respect internal callbackUrl so e.g. invitation emails can deep-link
      // to /portfolio/<id> after sign-in. Default to /home for bare sign-ins.
      if (url.startsWith("/")) return `${baseUrl}${url}`;
      if (url.startsWith(baseUrl)) return url;
      return `${baseUrl}/home`;
    },
  },

  // ------------------------------------------------------------------
  // General options
  // ------------------------------------------------------------------
  session: {
    strategy: "jwt", // lightweight, avoids DB session writes
  },
  jwt: {
    maxAge: 60 * 60 * 24 * 30, // 30 days
  },
  debug: false,
};