pulse/lib/auth/context.tsx

'use client'

import React, { createContext, useContext, useEffect, useState, useCallback } from 'react'
import { useRouter, usePathname } from 'next/navigation'
import apiRequest from '@/lib/api/client'
import { LoadingOverlay, useSessionSync, SessionExpiryWarning } from '@ciphera-net/ui'
import { logoutAction, getSessionAction, setSessionAction } from '@/app/actions/auth'
import { getUserOrganizations, switchContext } from '@/lib/api/organization'
import { logger } from '@/lib/utils/logger'

interface User {
  id: string
  email: string
  display_name?: string
  totp_enabled: boolean
  org_id?: string
  role?: string
  preferences?: {
    email_notifications?: {
      new_file_received: boolean
      file_downloaded: boolean
      login_alerts: boolean
      password_alerts: boolean
      two_factor_alerts: boolean
    }
  }
}

interface AuthContextType {
  user: User | null
  loading: boolean
  login: (user: User) => void
  logout: () => void
  refresh: () => Promise<void>
  refreshSession: () => Promise<void>
}

const AuthContext = createContext<AuthContextType>({
  user: null,
  loading: true,
  login: () => {},
  logout: () => {},
  refresh: async () => {},
  refreshSession: async () => {},
})

export function AuthProvider({ children }: { children: React.ReactNode }) {
  const [user, setUser] = useState<User | null>(null)
  const [loading, setLoading] = useState(true)
  const [isLoggingOut, setIsLoggingOut] = useState(false)
  const router = useRouter()
  const pathname = usePathname()

  const refreshToken = useCallback(async (): Promise<boolean> => {
    try {
      const res = await fetch('/api/auth/refresh', {
        method: 'POST',
        credentials: 'include',
      })
      if (res.ok) {
        localStorage.setItem('ciphera_token_refreshed_at', Date.now().toString())
      }
      return res.ok
    } catch {
      return false
    }
  }, [])

  const login = (userData: User) => {
    // * We still store user profile in localStorage for optimistic UI, but NOT the token
    localStorage.setItem('user', JSON.stringify(userData))
    localStorage.setItem('ciphera_token_refreshed_at', Date.now().toString())
    setUser(userData)
    router.refresh()
    // * Fetch full profile (including display_name) so header shows correct name without page refresh
    apiRequest<User>('/auth/user/me')
      .then((fullProfile) => {
        setUser((prev) => {
          const merged = {
            ...fullProfile,
            org_id: prev?.org_id ?? fullProfile.org_id,
            role: prev?.role ?? fullProfile.role,
          }
          localStorage.setItem('user', JSON.stringify(merged))
          return merged
        })
      })
      .catch((e) => logger.error('Failed to fetch full profile after login', e))
  }

  const logout = useCallback(async () => {
    setIsLoggingOut(true)
    try { await logoutAction() } catch { /* stale build — continue with client-side cleanup */ }
    localStorage.removeItem('user')
    localStorage.removeItem('ciphera_token_refreshed_at')
    localStorage.removeItem('ciphera_last_activity')
    // * Broadcast logout to other tabs (BroadcastChannel will handle if available)
    if (typeof window !== 'undefined' && 'BroadcastChannel' in window) {
      const channel = new BroadcastChannel('ciphera_session')
      channel.postMessage({ type: 'LOGOUT' })
      channel.close()
    }
    setTimeout(() => {
      window.location.href = '/'
    }, 500)
  }, [])

  const refresh = useCallback(async () => {
    try {
      const userData = await apiRequest<User>('/auth/user/me')
      
      setUser(prev => {
        const merged = {
          ...userData,
          org_id: prev?.org_id,
          role: prev?.role
        }
        localStorage.setItem('user', JSON.stringify(merged))
        return merged
      })
    } catch (e) {
      logger.error('Failed to refresh user data', e)
    }
    router.refresh()
  }, [router])

  const refreshSession = useCallback(async () => {
      await refresh()
  }, [refresh])

  // Initial load
  useEffect(() => {
    const init = async () => {
        // * 1. Check server-side session (cookies)
        let session: Awaited<ReturnType<typeof getSessionAction>> = null
        try {
          session = await getSessionAction()
          sessionStorage.removeItem('pulse_reload_for_stale_build')
        } catch {
          // * Stale build — browser has cached JS with old Server Action hashes.
          // * Force a hard reload once to fetch fresh bundles. Guard prevents infinite loop.
          const key = 'pulse_reload_for_stale_build'
          if (!sessionStorage.getItem(key)) {
            sessionStorage.setItem(key, '1')
            window.location.reload()
            return
          }
          sessionStorage.removeItem(key)
          // * Reload didn't fix it — treat as no session
          setLoading(false)
          return
        }

        // * 2. If no access_token but refresh_token may exist, try refresh (fixes 15-min inactivity logout)
        if (!session && typeof window !== 'undefined') {
          const refreshRes = await fetch('/api/auth/refresh', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            credentials: 'include',
          })
          if (refreshRes.ok) {
            try {
              session = await getSessionAction()
            } catch {
              const key = 'pulse_reload_for_stale_build'
              if (!sessionStorage.getItem(key)) {
                sessionStorage.setItem(key, '1')
                window.location.reload()
                return
              }
              sessionStorage.removeItem(key)
              setLoading(false)
              return
            }
          }
        }

        if (session) {
            setUser(session)
            localStorage.setItem('user', JSON.stringify(session))
            localStorage.setItem('ciphera_token_refreshed_at', Date.now().toString())
            // * Fetch full profile (including display_name) from API; preserve org_id/role from session
            try {
              const userData = await apiRequest<User>('/auth/user/me')
              const merged = { ...userData, org_id: session.org_id, role: session.role }
              setUser(merged)
              localStorage.setItem('user', JSON.stringify(merged))
            } catch (e) {
              logger.error('Failed to fetch full profile', e)
            }
        } else {
            // * Session invalid/expired
            localStorage.removeItem('user')
            setUser(null)
        }

        setLoading(false)
    }
    init()
  }, [])

  // * Sync session across browser tabs using BroadcastChannel
  useSessionSync({
    onLogout: () => {
      localStorage.removeItem('user')
      localStorage.removeItem('ciphera_token_refreshed_at')
      localStorage.removeItem('ciphera_last_activity')
      window.location.href = '/'
    },
    onLogin: (userData) => {
      setUser(userData as User)
      router.refresh()
    },
    onRefresh: () => {
      refresh()
    },
  })

  // * Stable primitives for the effect dependency array — avoids re-running
  // * on every render when the `user` object reference changes.
  const isAuthenticated = !!user
  const userOrgId = user?.org_id

  // * Organization Wall & Auto-Switch
  useEffect(() => {
    const checkOrg = async () => {
      if (!loading && isAuthenticated) {
        if (pathname?.startsWith('/onboarding')) return
        if (pathname?.startsWith('/auth/callback')) return

        try {
          const organizations = await getUserOrganizations()

          if (organizations.length === 0) {
            if (pathname?.startsWith('/welcome')) return
            router.push('/welcome')
            return
          }

          // * If user has organizations but no context (org_id), switch to the first one
          if (!userOrgId && organizations.length > 0) {
             const firstOrg = organizations[0]
             
             try {
                 const { access_token } = await switchContext(firstOrg.organization_id)
                 
                 // * Update session cookie
                 const result = await setSessionAction(access_token)
                 if (result.success && result.user) {
                     try {
                       const fullProfile = await apiRequest<{ id: string; email: string; display_name?: string; totp_enabled: boolean; org_id?: string; role?: string }>('/auth/user/me')
                       const merged = { ...fullProfile, org_id: result.user.org_id ?? fullProfile.org_id, role: result.user.role ?? fullProfile.role }
                       setUser(merged)
                       localStorage.setItem('user', JSON.stringify(merged))
                     } catch {
                       setUser(result.user)
                       localStorage.setItem('user', JSON.stringify(result.user))
                     }
                     router.refresh()
                 }
             } catch (e) {
                 logger.error('Failed to auto-switch context', e)
             }
          }
        } catch (e) {
          logger.error("Failed to fetch organizations", e)
        }
      }
    }
    
    checkOrg()
  }, [loading, isAuthenticated, userOrgId, pathname, router])

  return (
    <AuthContext.Provider value={{ user, loading, login, logout, refresh, refreshSession }}>
      {isLoggingOut && <LoadingOverlay logoSrc="/pulse_icon_no_margins.png" title="Pulse" />}
      <SessionExpiryWarning
        isAuthenticated={!!user}
        onRefreshToken={refreshToken}
        onExpired={logout}
      />
      {children}
    </AuthContext.Provider>
  )
}

export const useAuth = () => useContext(AuthContext)