import { requestIdContext, userContext } from '@/context'
import { createSessionStorage } from '@/lib/session'

import { createAuth, type User } from '@/lib/auth'
import { decodeJwt } from 'jose'
import { redirect, type LoaderFunctionArgs } from 'react-router'
import type { OAuth2Strategy } from 'remix-auth-oauth2'

export const authMiddleware = async (
  { request, context }: LoaderFunctionArgs,
  next: () => Promise<Response>
): Promise<Response> => {
  const sessionStorage = createSessionStorage(context.cloudflare.env)
  const authenticator = createAuth(context.cloudflare.env)
  const strategy = authenticator.get<OAuth2Strategy<User>>('oidc')
  const session = await sessionStorage.getSession(request.headers.get('cookie'))
  const requestId = context.get(requestIdContext)
  let user = session.get('user') as User | null

  session.set('returnTo', new URL(request.url).toString())

  if (!user) {
    console.log('There is no user logged in')

    return redirect('/login', {
      headers: new Headers({
        'Set-Cookie': await sessionStorage.commitSession(session)
      })
    })
  }

  try {
    const accessToken = decodeJwt(user.accessToken) as { exp: number }
    const accessTokenExp = accessToken.exp * 1000
    const leeway = 120 * 1000 // 2 minutes

    if (Date.now() > accessTokenExp - leeway) {
      const tokens = await (strategy as any).refreshToken(user.refreshToken)

      user = {
        ...user,
        accessToken: tokens.accessToken(),
        refreshToken: tokens.refreshToken()
      }

      console.debug(`[${requestId}] Refresh token retrieved`, user)
      // Should replace the user in the session
      session.set('user', user)
    }
  } catch (error) {
    console.error(`[${requestId}]`, error?.stack)

    // If refreshing the token fails, remove the user from the current session
    // so the user is forced to sign in again
    session.unset('user')

    return redirect('/login', {
      headers: new Headers({
        'Set-Cookie': await sessionStorage.commitSession(session)
      })
    })
  }

  context.set(userContext, user)

  const response = await next()
  const sessionCookie = await sessionStorage.commitSession(session)
  response.headers.set('Set-Cookie', sessionCookie)
  return response
}