import {
  type AfterCallbackAppRoute,
  type AppRouteHandlerFnContext,
  getSession,
  handleAuth,
  handleCallback,
  handleLogin,
  handleLogout,
  type Session,
} from "@auth0/nextjs-auth0"
import { type NextRequest, NextResponse } from "next/server"
import { createOrUpdateUser, hasCompletedOnboarding } from "@codeflash-ai/common"
import { trackUserLogin } from "@/lib/analytics/tracking"
import { cookies } from "next/headers"
import { APP_ROUTES } from "@/lib/types"

//In case we want to change some future variables to set redirect to marketing campaign
const LOGOUT_REDIRECT_URL =
  process.env.CODEFLASH_LOGOUT_REDIRECT_URL ??
  process.env.CODEFLASH_MARKETING_URL ??
  "https://codeflash.ai"

// THIS IS THE KEY CHANGE - Your afterCallback was empty!
const afterCallback: AfterCallbackAppRoute = async (req: NextRequest, session: Session) => {
  if (!session.user) {
    return session
  }

  const user = session.user
  console.log(`[Auth] Processing login for user: ${user.sub}`)

  if (!user.sub || !user.nickname) {
    console.error("[Auth] Missing required user fields")
    return session
  }

  try {
    // 1. SAVE TO DATABASE (moved from login page!)
    console.log("[Auth] Saving user to database...")
    await createOrUpdateUser(user.sub, user.nickname, user.email ?? null, user.name ?? null)
    console.log("[Auth] User saved successfully")

    // 2. TRACK LOGIN (moved from login page!)
    await trackUserLogin({
      userId: user.sub,
      username: user.nickname,
      email: user.email,
      name: user.name,
    })

    // 3. CHECK ONBOARDING (moved from login page!)
    const completedOnboarding = await hasCompletedOnboarding(user.sub)
    console.log(`[Auth] Onboarding completed: ${completedOnboarding}`)

    // 4. Decide where to redirect - Auth0 preserves returnTo from login
    let intendedDestination = APP_ROUTES.BASE

    // Try to get returnTo from multiple sources
    const url = new URL(req.url)

    // Method 1: From URL search params (direct from Auth0 redirect)
    const returnToParam = url.searchParams.get("returnTo")
    if (returnToParam) {
      intendedDestination = returnToParam
      console.log(`[Auth] Found returnTo in URL params: ${intendedDestination}`)
    }

    // Method 2: From state parameter (fallback)
    const stateParam = url.searchParams.get("state")
    if (stateParam && !returnToParam) {
      try {
        const state = JSON.parse(Buffer.from(stateParam, "base64").toString("utf-8"))
        if (state.returnTo) {
          intendedDestination = state.returnTo
          console.log(`[Auth] Found returnTo in state: ${intendedDestination}`)
        }
      } catch (e) {
        console.warn("[Auth] Failed to parse state:", e)
      }
    }
    // check if the path is codeflash/auth/[token]
    const isAuthPath =
      intendedDestination.startsWith("/codeflash/auth") ||
      intendedDestination.includes("/codeflash/auth")
    console.log(`[Auth] isAuthPath: ${isAuthPath}`)

    // Handle onboarding redirect
    if (!completedOnboarding && !isAuthPath) {
      session.returnTo = "/onboarding"
    } else {
      session.returnTo = intendedDestination
    }
  } catch (error) {
    console.error("[Auth] Error in afterCallback:", error)
    // Don't fail login even if our processing fails
    session.returnTo = APP_ROUTES.BASE
  }

  return session
}

// Rest of your file stays mostly the same...
export const GET = handleAuth({
  // Fixed login handler to preserve returnTo parameter
  login: async (request: any, response: any) => {
    console.log("Logging in")
    try {
      const req = request as NextRequest
      const url = new URL(req.url)
      const returnTo = url.searchParams.get("returnTo") || APP_ROUTES.BASE

      console.log(`[Auth] Login with returnTo: ${returnTo}`)

      return await handleLogin(req, response as AppRouteHandlerFnContext, {
        returnTo,
        authorizationParams: {
          scope: "openid profile email offline_access",
        },
      })
    } catch (error) {
      console.error("Error logging in:", error)
      return NextResponse.json({ error: "Failed to initiate login" }, { status: 500 })
    }
  },

  // Your existing logout handler...
  logout: async (request: any, response: any) => {
    console.log("Logging out")
    try {
      return await handleLogout(request as NextRequest, response as AppRouteHandlerFnContext, {
        returnTo: LOGOUT_REDIRECT_URL,
      })
    } catch (error) {
      console.error("Error logging out:", error)
      return NextResponse.redirect(LOGOUT_REDIRECT_URL)
    }
  },

  // Updated callback handler
  callback: async (req: any, res: any) => {
    try {
      const response = (await handleCallback(req as NextRequest, res as AppRouteHandlerFnContext, {
        afterCallback, // NOW THIS DOES SOMETHING!
      })) as NextResponse

      const session = await getSession(req as NextRequest, response)

      if (session != null) {
        // Use the returnTo set by afterCallback
        const returnTo = session.returnTo || APP_ROUTES.BASE
        const isAbsolute = returnTo.includes(process.env.AUTH0_BASE_URL ?? "")
        const redirectUrl = isAbsolute ? returnTo : `${process.env.AUTH0_BASE_URL}${returnTo}`
        return NextResponse.redirect(redirectUrl, response)
      } else {
        return NextResponse.redirect(`${process.env.AUTH0_BASE_URL}/waitlist`, response)
      }
    } catch (error: any) {
      console.error("Error in callback:", error)
      // Your existing error handling...
      if (error.status === 400 && error.message.search("allowlist-fail") !== -1) {
        const re = /allowlist-fail\s(.*)\s(.*)\)/
        const match = error.message.match(re)
        if (match != null) {
          const userId = match[1]
          const userNickname = match[2]
          return NextResponse.redirect(
            `${process.env.AUTH0_BASE_URL}/waitlist?username=${userNickname}&userid=${userId}`,
          )
        }
      }
      // If error doesn't match any specific case, return error page
      return NextResponse.redirect(`${process.env.AUTH0_BASE_URL}/login?error=callback_failed`)
    }
  },
})
