stable-diffusion.cpp-rest/webui/lib/auth-context.tsx

"use client"

import React, { createContext, useContext, useEffect, useState, ReactNode } from 'react'
import { authApi } from './api'

interface User {
  id: string
  username: string
  email?: string
  role: 'admin' | 'user'
  createdAt: string
  lastLogin?: string
}

interface AuthState {
  user: User | null
  token: string | null
  isAuthenticated: boolean
  isLoading: boolean
  error: string | null
}

interface AuthContextType extends AuthState {
  login: (username: string, password?: string) => Promise<void>
  logout: () => void
  refreshToken: () => Promise<void>
  clearError: () => void
  authMethod: 'none' | 'unix' | 'jwt' | 'shared-key'
  authEnabled: boolean
}

const AuthContext = createContext<AuthContextType | undefined>(undefined)

export function AuthProvider({ children }: { children: ReactNode }) {
  const [state, setState] = useState<AuthState>({
    user: null,
    token: null,
    isAuthenticated: false,
    isLoading: true,
    error: null
  })

  // Get authentication method from server config
  const authMethod = typeof window !== 'undefined' && window.__SERVER_CONFIG__
    ? window.__SERVER_CONFIG__.authMethod
    : 'jwt';
  const authEnabled = typeof window !== 'undefined' && window.__SERVER_CONFIG__
    ? window.__SERVER_CONFIG__.authEnabled
    : false;

  useEffect(() => {
    // Check for existing authentication on mount
    if (authMethod === 'shared-key') {
      // For shared-key auth, automatically authenticate
      validateSharedKey()
    } else if (authMethod === 'unix') {
      // For Unix auth, check for unix_user in localStorage
      const unixUser = localStorage.getItem('unix_user')
      if (unixUser) {
        validateUnixUser(unixUser)
      } else {
        setState(prev => ({ ...prev, isLoading: false }))
      }
    } else {
      // For JWT auth, check for token in localStorage
      const token = localStorage.getItem('auth_token')
      if (token) {
        validateToken(token)
      } else {
        setState(prev => ({ ...prev, isLoading: false }))
      }
    }
  }, [authMethod])

  const validateToken = async (token: string) => {
    try {
      const data = await authApi.validateToken(token)
      setState({
        user: data.user,
        token,
        isAuthenticated: true,
        isLoading: false,
        error: null
      })
      localStorage.setItem('auth_token', token)
    } catch (error) {
      console.error('Token validation error:', error)
      localStorage.removeItem('auth_token')
      setState({
        user: null,
        token: null,
        isAuthenticated: false,
        isLoading: false,
        error: 'Session expired. Please log in again.'
      })
    }
  }

  const validateSharedKey = async () => {
    try {
      // For shared-key auth, we can validate by making a simple API call
      const data = await authApi.getCurrentUser()
      setState({
        user: data.user,
        token: 'shared-key-auth', // Store a pseudo-token for consistency
        isAuthenticated: true,
        isLoading: false,
        error: null
      })
      localStorage.setItem('auth_token', 'shared-key-auth') // Store pseudo-token for API calls
    } catch (error) {
      console.error('Shared key validation error:', error)
      localStorage.removeItem('auth_token')
      setState({
        user: null,
        token: null,
        isAuthenticated: false,
        isLoading: false,
        error: 'Shared key authentication failed.'
      })
    }
  }

  const validateUnixUser = async (unixUser: string) => {
    try {
      // For Unix auth, we need to validate by attempting to get current user
      const data = await authApi.getCurrentUser()
      setState({
        user: data.user,
        token: `unix_${unixUser}`, // Store a pseudo-token for consistency
        isAuthenticated: true,
        isLoading: false,
        error: null
      })
      localStorage.setItem('unix_user', unixUser)
      localStorage.setItem('auth_token', `unix_${unixUser}`) // Store pseudo-token for API calls
    } catch (error) {
      console.error('Unix user validation error:', error)
      localStorage.removeItem('unix_user')
      localStorage.removeItem('auth_token')
      setState({
        user: null,
        token: null,
        isAuthenticated: false,
        isLoading: false,
        error: 'Session expired. Please log in again.'
      })
    }
  }

  const login = async (username: string, password?: string) => {
    setState(prev => ({ ...prev, isLoading: true, error: null }))

    try {
      const data = await authApi.login(username, password)
      const { token, user } = data

      setState({
        user,
        token,
        isAuthenticated: true,
        isLoading: false,
        error: null
      })

      if (authMethod === 'shared-key') {
        localStorage.setItem('auth_token', token) // Store pseudo-token for API calls
      } else if (authMethod === 'unix') {
        localStorage.setItem('unix_user', username)
        localStorage.setItem('auth_token', token) // Store token for API calls
      } else {
        localStorage.setItem('auth_token', token)
      }
    } catch (error) {
      setState(prev => ({
        ...prev,
        isLoading: false,
        error: error instanceof Error ? error.message : 'Login failed'
      }))
    }
  }

  const logout = () => {
    if (authMethod === 'unix') {
      localStorage.removeItem('unix_user')
    }
    localStorage.removeItem('auth_token')
    setState({
      user: null,
      token: null,
      isAuthenticated: false,
      isLoading: false,
      error: null
    })
  }

  const refreshToken = async () => {
    const { token } = state
    if (!token) return

    try {
      const data = await authApi.refreshToken()
      const newToken = data.token

      setState(prev => ({ ...prev, token: newToken }))
      localStorage.setItem('auth_token', newToken)
    } catch (error) {
      console.error('Token refresh error:', error)
      logout()
    }
  }

  const clearError = () => {
    setState(prev => ({ ...prev, error: null }))
  }

  const value: AuthContextType = {
    ...state,
    login,
    logout,
    refreshToken,
    clearError,
    authMethod,
    authEnabled
  }

  return (
    <AuthContext.Provider value={value}>
      {children}
    </AuthContext.Provider>
  )
}

export function useAuth(): AuthContextType {
  const context = useContext(AuthContext)
  if (context === undefined) {
    throw new Error('useAuth must be used within an AuthProvider')
  }
  return context
}