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
}

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
  })

  useEffect(() => {
    // Check for existing token on mount
    const token = localStorage.getItem('auth_token')
    if (token) {
      validateToken(token)
    } else {
      setState(prev => ({ ...prev, isLoading: false }))
    }
  }, [])

  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 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
      })

      localStorage.setItem('auth_token', token)
    } catch (error) {
      setState(prev => ({
        ...prev,
        isLoading: false,
        error: error instanceof Error ? error.message : 'Login failed'
      }))
    }
  }

  const logout = () => {
    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
  }

  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
}