stable-diffusion.cpp-rest/webui/lib/api.ts

// API client for stable-diffusion REST API

// Type for server config injected by the server
declare global {
  interface Window {
    __SERVER_CONFIG__?: {
      apiUrl: string;
      apiBasePath: string;
      host: string;
      port: number;
      authMethod: 'none' | 'unix' | 'jwt';
      authEnabled: boolean;
    };
  }
}

// Request throttling to prevent excessive API calls
class RequestThrottler {
  private requests: Map<string, { count: number; resetTime: number }> = new Map();
  private maxRequests: number = 10; // Max requests per time window
  private timeWindow: number = 1000; // Time window in milliseconds

  canMakeRequest(key: string): boolean {
    const now = Date.now();
    const request = this.requests.get(key);

    if (!request || now >= request.resetTime) {
      this.requests.set(key, { count: 1, resetTime: now + this.timeWindow });
      return true;
    }

    if (request.count >= this.maxRequests) {
      return false;
    }

    request.count++;
    return true;
  }

  getWaitTime(key: string): number {
    const request = this.requests.get(key);
    if (!request) return 0;

    const now = Date.now();
    if (now >= request.resetTime) return 0;

    return request.resetTime - now;
  }
}

// Global throttler instance
const throttler = new RequestThrottler();

// Debounce utility for frequent calls
function debounce<T extends (...args: any[]) => any>(
  func: T,
  wait: number,
  immediate?: boolean
): (...args: Parameters<T>) => void {
  let timeout: NodeJS.Timeout | null = null;

  return function executedFunction(...args: Parameters<T>) {
    const later = () => {
      timeout = null;
      if (!immediate) func(...args);
    };

    const callNow = immediate && !timeout;

    if (timeout) clearTimeout(timeout);
    timeout = setTimeout(later, wait);

    if (callNow) func(...args);
  };
}

// Cache for API responses to reduce redundant calls
class ApiCache {
  private cache: Map<string, { data: any; timestamp: number; ttl: number }> = new Map();
  private defaultTtl: number = 5000; // 5 seconds default TTL

  set(key: string, data: any, ttl?: number): void {
    this.cache.set(key, {
      data,
      timestamp: Date.now(),
      ttl: ttl || this.defaultTtl
    });
  }

  get(key: string): any | null {
    const cached = this.cache.get(key);
    if (!cached) return null;

    if (Date.now() - cached.timestamp > cached.ttl) {
      this.cache.delete(key);
      return null;
    }

    return cached.data;
  }

  clear(): void {
    this.cache.clear();
  }

  delete(key: string): void {
    this.cache.delete(key);
  }
}

const cache = new ApiCache();

// Get configuration from server-injected config or fallback to environment/defaults
// This function is called at runtime to ensure __SERVER_CONFIG__ is available
function getApiConfig() {
  if (typeof window !== 'undefined' && window.__SERVER_CONFIG__) {
    return {
      apiUrl: window.__SERVER_CONFIG__.apiUrl,
      apiBase: window.__SERVER_CONFIG__.apiBasePath,
    };
  }

  // Fallback for development mode - use current window location
  if (typeof window !== 'undefined') {
    const protocol = window.location.protocol;
    const host = window.location.hostname;
    const port = window.location.port;
    return {
      apiUrl: `${protocol}//${host}:${port}`,
      apiBase: '/api',
    };
  }

  // Server-side fallback
  return {
    apiUrl: process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8081',
    apiBase: process.env.NEXT_PUBLIC_API_BASE_PATH || '/api',
  };
}

export interface GenerationRequest {
  model?: string;
  prompt: string;
  negative_prompt?: string;
  width?: number;
  height?: number;
  steps?: number;
  cfg_scale?: number;
  seed?: string;
  sampling_method?: string;
  scheduler?: string;
  batch_count?: number;
  clip_skip?: number;
  strength?: number;
  control_strength?: number;
}

export interface JobInfo {
  id?: string;
  request_id?: string;
  status: 'pending' | 'processing' | 'completed' | 'failed' | 'cancelled' | 'queued';
  progress?: number;
  result?: {
    images: string[];
  };
  outputs?: Array<{
    filename: string;
    url: string;
    path: string;
  }>;
  error?: string;
  created_at?: string;
  updated_at?: string;
  message?: string;
  queue_position?: number;
}

export interface ModelInfo {
  id?: string;
  name: string;
  path?: string;
  type: string;
  size?: number;
  file_size?: number;
  file_size_mb?: number;
  sha256?: string | null;
  sha256_short?: string | null;
  loaded?: boolean;
}

export interface QueueStatus {
  active_generations: number;
  jobs: JobInfo[];
  running: boolean;
  size: number;
}

export interface HealthStatus {
  status: 'ok' | 'error' | 'degraded';
  message: string;
  timestamp: string;
  uptime?: number;
  version?: string;
}

class ApiClient {
  private baseUrl: string = '';
  private isInitialized: boolean = false;

  // Initialize base URL
  private initBaseUrl(): string {
    if (!this.isInitialized) {
      const config = getApiConfig();
      this.baseUrl = `${config.apiUrl}${config.apiBase}`;
      this.isInitialized = true;
    }
    return this.baseUrl;
  }

  // Get base URL dynamically at runtime to ensure server config is loaded
  private getBaseUrl(): string {
    return this.initBaseUrl();
  }

  private async request<T>(
    endpoint: string,
    options: RequestInit = {}
  ): Promise<T> {
    const url = `${this.getBaseUrl()}${endpoint}`;

    // Check request throttling for certain endpoints
    const needsThrottling = endpoint.includes('/queue/status') || endpoint.includes('/health');
    if (needsThrottling) {
      const waitTime = throttler.getWaitTime(endpoint);
      if (waitTime > 0) {
        // Wait before making the request
        await new Promise(resolve => setTimeout(resolve, waitTime));
      }
      if (!throttler.canMakeRequest(endpoint)) {
        throw new Error('Too many requests. Please wait before making another request.');
      }
    }

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

    // Add auth token or Unix user header based on auth method
    const token = typeof window !== 'undefined' ? localStorage.getItem('auth_token') : null;
    const unixUser = typeof window !== 'undefined' ? localStorage.getItem('unix_user') : null;

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
      ...options.headers as Record<string, string>,
    };

    if (authMethod === 'unix' && unixUser) {
      // For Unix auth, send the username in X-Unix-User header
      headers['X-Unix-User'] = unixUser;
    } else if (token) {
      // For JWT auth, send the token in Authorization header
      headers['Authorization'] = `Bearer ${token}`;
    }

    const response = await fetch(url, {
      ...options,
      headers,
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({
        error: { message: response.statusText },
      }));

      // Handle nested error structure: { error: { message: "..." } }
      const errorMessage =
        errorData.error?.message ||
        errorData.message ||
        errorData.error ||
        'API request failed';

      throw new Error(errorMessage);
    }

    return response.json();
  }

  // Enhanced health check with caching and better error handling
  async checkHealth(): Promise<HealthStatus> {
    const cacheKey = 'health_check';
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    const endpoints = ['/queue/status', '/health', '/status', '/'];

    for (const endpoint of endpoints) {
      try {
        const response = await fetch(`${this.getBaseUrl()}${endpoint}`, {
          method: 'GET',
          headers: {
            'Content-Type': 'application/json',
          },
          // Add timeout to prevent hanging requests
          signal: AbortSignal.timeout(3000), // Reduced timeout
        });

        if (response.ok) {
          const data = await response.json();

          // For queue status, consider it healthy if it returns valid structure
          if (endpoint === '/queue/status' && data.queue) {
            const result = {
              status: 'ok' as const,
              message: 'API is running and queue is accessible',
              timestamp: new Date().toISOString(),
            };
            cache.set(cacheKey, result, 10000); // Cache for 10 seconds
            return result;
          }

          // For other health endpoints
          const healthStatus: HealthStatus = {
            status: 'ok',
            message: 'API is running',
            timestamp: new Date().toISOString(),
            uptime: data.uptime,
            version: data.version || data.build || data.git_version,
          };

          // Handle different response formats
          if (data.status) {
            if (data.status === 'healthy' || data.status === 'running' || data.status === 'ok') {
              healthStatus.status = 'ok';
              healthStatus.message = data.message || 'API is running';
            } else if (data.status === 'degraded') {
              healthStatus.status = 'degraded';
              healthStatus.message = data.message || 'API is running in degraded mode';
            } else {
              healthStatus.status = 'error';
              healthStatus.message = data.message || 'API status is unknown';
            }
          }

          cache.set(cacheKey, healthStatus, 10000); // Cache for 10 seconds
          return healthStatus;
        }
      } catch (error) {
        // Continue to next endpoint if this one fails
        console.warn(`Health check failed for endpoint ${endpoint}:`, error);
        continue;
      }
    }

    // If all endpoints fail
    throw new Error('All health check endpoints are unavailable');
  }

  // Alternative simple connectivity check with caching
  async checkConnectivity(): Promise<boolean> {
    const cacheKey = 'connectivity_check';
    const cachedResult = cache.get(cacheKey);
    if (cachedResult !== null) {
      return cachedResult;
    }

    try {
      const response = await fetch(`${this.getBaseUrl()}`, {
        method: 'HEAD',
        signal: AbortSignal.timeout(2000), // Reduced timeout
      });
      const result = response.ok || response.status < 500;
      cache.set(cacheKey, result, 5000); // Cache for 5 seconds
      return result;
    } catch (error) {
      cache.set(cacheKey, false, 5000); // Cache failure for 5 seconds
      return false;
    }
  }

  // Generation endpoints
  async generateImage(params: GenerationRequest): Promise<JobInfo> {
    return this.request<JobInfo>('/generate/text2img', {
      method: 'POST',
      body: JSON.stringify(params),
    });
  }

  async text2img(params: GenerationRequest): Promise<JobInfo> {
    return this.request<JobInfo>('/generate/text2img', {
      method: 'POST',
      body: JSON.stringify(params),
    });
  }

  async img2img(params: GenerationRequest & { image: string }): Promise<JobInfo> {
    // Convert frontend field name to backend field name
    const backendParams = {
      ...params,
      init_image: params.image,
      image: undefined, // Remove the frontend field
    };
    return this.request<JobInfo>('/generate/img2img', {
      method: 'POST',
      body: JSON.stringify(backendParams),
    });
  }

  async inpainting(params: GenerationRequest & { source_image: string; mask_image: string }): Promise<JobInfo> {
    return this.request<JobInfo>('/generate/inpainting', {
      method: 'POST',
      body: JSON.stringify(params),
    });
  }

  // Job management with caching for status checks
  async getJobStatus(jobId: string): Promise<JobInfo> {
    const cacheKey = `job_status_${jobId}`;
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    const result = await this.request<JobInfo>(`/queue/job/${jobId}`);

    // Cache job status for a short time
    if (result.status === 'processing' || result.status === 'queued') {
      cache.set(cacheKey, result, 2000); // Cache for 2 seconds for active jobs
    } else {
      cache.set(cacheKey, result, 10000); // Cache for 10 seconds for completed jobs
    }

    return result;
  }

  // Get authenticated image URL with cache-busting
  getImageUrl(jobId: string, filename: string): string {
    const baseUrl = this.getBaseUrl();

    // Add cache-busting timestamp
    const timestamp = Date.now();
    const url = `${baseUrl}/queue/job/${jobId}/output/${filename}?t=${timestamp}`;

    return url;
  }

  // Download image with authentication
  async downloadImage(jobId: string, filename: string): Promise<Blob> {
    const url = this.getImageUrl(jobId, filename);

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

    // Add auth token or Unix user header based on auth method
    const token = typeof window !== 'undefined' ? localStorage.getItem('auth_token') : null;
    const unixUser = typeof window !== 'undefined' ? localStorage.getItem('unix_user') : null;

    const headers: Record<string, string> = {};

    if (authMethod === 'unix' && unixUser) {
      // For Unix auth, send the username in X-Unix-User header
      headers['X-Unix-User'] = unixUser;
    } else if (token) {
      // For JWT auth, send the token in Authorization header
      headers['Authorization'] = `Bearer ${token}`;
    }

    const response = await fetch(url, {
      headers,
    });

    if (!response.ok) {
      const errorData = await response.json().catch(() => ({
        error: { message: response.statusText },
      }));

      // Handle nested error structure: { error: { message: "..." } }
      const errorMessage =
        errorData.error?.message ||
        errorData.message ||
        errorData.error ||
        'Failed to download image';

      throw new Error(errorMessage);
    }

    return response.blob();
  }

  async cancelJob(jobId: string): Promise<void> {
    // Clear job status cache when cancelling
    cache.delete(`job_status_${jobId}`);

    return this.request<void>('/queue/cancel', {
      method: 'POST',
      body: JSON.stringify({ job_id: jobId }),
    });
  }

  // Get queue status with caching and throttling
  async getQueueStatus(): Promise<QueueStatus> {
    const cacheKey = 'queue_status';
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    const response = await this.request<{ queue: QueueStatus }>('/queue/status');

    // Cache queue status based on current activity
    const hasActiveJobs = response.queue.jobs.some(job =>
      job.status === 'processing' || job.status === 'queued'
    );

    // Cache for shorter time if there are active jobs
    const cacheTime = hasActiveJobs ? 1000 : 5000; // 1 second for active, 5 seconds for idle
    cache.set(cacheKey, response.queue, cacheTime);

    return response.queue;
  }

  async clearQueue(): Promise<void> {
    // Clear all related caches
    cache.delete('queue_status');

    return this.request<void>('/queue/clear', {
      method: 'POST',
    });
  }

  // Model management
  async getModels(type?: string, loaded?: boolean, page: number = 1, limit: number = 50, search?: string): Promise<{ models: ModelInfo[]; pagination: any; statistics: any }> {
    const cacheKey = `models_${type || 'all'}_${loaded ? 'loaded' : 'all'}_${page}_${limit}_${search || 'all'}`;
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    let endpoint = '/models';
    const params = [];
    if (type && type !== 'loaded') params.push(`type=${type}`);
    if (type === 'loaded' || loaded) params.push('loaded=true');
    params.push(`page=${page}`);
    params.push(`limit=${limit}`);
    if (search) params.push(`search=${encodeURIComponent(search)}`);

    // Add include_metadata for additional information
    params.push('include_metadata=true');

    if (params.length > 0) endpoint += '?' + params.join('&');

    const response = await this.request<{
      models: ModelInfo[];
      pagination: {
        page: number;
        limit: number;
        total_count: number;
        total_pages: number;
        has_next: boolean;
        has_prev: boolean
      };
      statistics: any;
    }>(endpoint);

    const models = response.models.map(model => ({
      ...model,
      id: model.sha256_short || model.name,
      size: model.file_size || model.size,
      path: model.path || model.name,
    }));

    const result = {
      models,
      pagination: response.pagination,
      statistics: response.statistics || {}
    };

    // Cache models for 30 seconds as they don't change frequently
    cache.set(cacheKey, result, 30000);

    return result;
  }

  // Get all models (for backward compatibility)
  async getAllModels(type?: string, loaded?: boolean): Promise<ModelInfo[]> {
    const allModels: ModelInfo[] = [];
    let page = 1;
    const limit = 100;

    while (true) {
      const response = await this.getModels(type, loaded, page, limit);
      allModels.push(...response.models);

      if (!response.pagination.has_next) {
        break;
      }

      page++;
    }

    return allModels;
  }

  async getModelInfo(modelId: string): Promise<ModelInfo> {
    const cacheKey = `model_info_${modelId}`;
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    const result = await this.request<ModelInfo>(`/models/${modelId}`);
    cache.set(cacheKey, result, 30000); // Cache for 30 seconds
    return result;
  }

  async loadModel(modelId: string): Promise<void> {
    // Clear model cache when loading
    cache.delete(`model_info_${modelId}`);

    return this.request<void>(`/models/${modelId}/load`, {
      method: 'POST',
    });
  }

  async unloadModel(modelId: string): Promise<void> {
    // Clear model cache when unloading
    cache.delete(`model_info_${modelId}`);

    return this.request<void>(`/models/${modelId}/unload`, {
      method: 'POST',
    });
  }

  async scanModels(): Promise<void> {
    // Clear all model caches when scanning
    cache.clear();

    return this.request<void>('/models/refresh', {
      method: 'POST',
    });
  }

  async convertModel(modelName: string, quantizationType: string, outputPath?: string): Promise<{ request_id: string; message: string }> {
    return this.request<{ request_id: string; message: string }>('/models/convert', {
      method: 'POST',
      body: JSON.stringify({
        model_name: modelName,
        quantization_type: quantizationType,
        output_path: outputPath,
      }),
    });
  }

  // System endpoints
  async getHealth(): Promise<{ status: string }> {
    return this.request<{ status: string }>('/health');
  }

  async getStatus(): Promise<any> {
    return this.request<any>('/status');
  }

  async getSystemInfo(): Promise<any> {
    return this.request<any>('/system');
  }

  async restartServer(): Promise<{ message: string }> {
    return this.request<{ message: string }>('/system/restart', {
      method: 'POST',
      body: JSON.stringify({}),
    });
  }

  // Configuration endpoints with caching
  async getSamplers(): Promise<Array<{ name: string; description: string; recommended_steps: number }>> {
    const cacheKey = 'samplers';
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    const response = await this.request<{ samplers: Array<{ name: string; description: string; recommended_steps: number }> }>('/samplers');
    cache.set(cacheKey, response.samplers, 60000); // Cache for 1 minute
    return response.samplers;
  }

  async getSchedulers(): Promise<Array<{ name: string; description: string }>> {
    const cacheKey = 'schedulers';
    const cachedResult = cache.get(cacheKey);
    if (cachedResult) {
      return cachedResult;
    }

    const response = await this.request<{ schedulers: Array<{ name: string; description: string }> }>('/schedulers');
    cache.set(cacheKey, response.schedulers, 60000); // Cache for 1 minute
    return response.schedulers;
  }

  // Cache management methods
  clearCache(): void {
    cache.clear();
  }

  clearCacheByPrefix(prefix: string): void {
    const keysToDelete: string[] = [];
    (cache as any).cache.forEach((_: any, key: string) => {
      if (key.startsWith(prefix)) {
        keysToDelete.push(key);
      }
    });
    keysToDelete.forEach(key => cache.delete(key));
  }
}

// Generic API request function for authentication
export async function apiRequest(
  endpoint: string,
  options: RequestInit = {}
): Promise<Response> {
  const { apiUrl, apiBase } = getApiConfig();
  const url = `${apiUrl}${apiBase}${endpoint}`;

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

  // Add auth token or Unix user header based on auth method
  const token = typeof window !== 'undefined' ? localStorage.getItem('auth_token') : null;
  const unixUser = typeof window !== 'undefined' ? localStorage.getItem('unix_user') : null;

  const headers: Record<string, string> = {
    'Content-Type': 'application/json',
    ...options.headers as Record<string, string>,
  };

  if (authMethod === 'unix' && unixUser) {
    // For Unix auth, send the username in X-Unix-User header
    headers['X-Unix-User'] = unixUser;
  } else if (token) {
    // For JWT auth, send the token in Authorization header
    headers['Authorization'] = `Bearer ${token}`;
  }

  return fetch(url, {
    ...options,
    headers,
  });
}

// Authentication API endpoints
export const authApi = {
  async login(username: string, password?: string) {
    // Get authentication method from server config
    const authMethod = typeof window !== 'undefined' && window.__SERVER_CONFIG__
      ? window.__SERVER_CONFIG__.authMethod
      : 'jwt';

    // For both Unix and JWT auth, send username and password
    // The server will handle whether password is required based on PAM availability
    const response = await apiRequest('/auth/login', {
      method: 'POST',
      body: JSON.stringify({ username, password }),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Login failed' }));
      throw new Error(error.message || 'Login failed');
    }

    return response.json();
  },

  async validateToken(token: string) {
    const response = await apiRequest('/auth/validate', {
      headers: { 'Authorization': `Bearer ${token}` },
    });

    if (!response.ok) {
      throw new Error('Token validation failed');
    }

    return response.json();
  },

  async refreshToken() {
    const response = await apiRequest('/auth/refresh', {
      method: 'POST',
    });

    if (!response.ok) {
      throw new Error('Token refresh failed');
    }

    return response.json();
  },

  async logout() {
    await apiRequest('/auth/logout', {
      method: 'POST',
    });
  },

  async getCurrentUser() {
    const response = await apiRequest('/auth/me');

    if (!response.ok) {
      throw new Error('Failed to get current user');
    }

    return response.json();
  },

  async changePassword(oldPassword: string, newPassword: string) {
    const response = await apiRequest('/auth/change-password', {
      method: 'POST',
      body: JSON.stringify({ old_password: oldPassword, new_password: newPassword }),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Password change failed' }));
      throw new Error(error.message || 'Password change failed');
    }

    return response.json();
  },

  // API Key management
  async getApiKeys() {
    const response = await apiRequest('/auth/api-keys');

    if (!response.ok) {
      throw new Error('Failed to get API keys');
    }

    return response.json();
  },

  async createApiKey(name: string, scopes?: string[]) {
    const response = await apiRequest('/auth/api-keys', {
      method: 'POST',
      body: JSON.stringify({ name, scopes }),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Failed to create API key' }));
      throw new Error(error.message || 'Failed to create API key');
    }

    return response.json();
  },

  async deleteApiKey(keyId: string) {
    const response = await apiRequest(`/auth/api-keys/${keyId}`, {
      method: 'DELETE',
    });

    if (!response.ok) {
      throw new Error('Failed to delete API key');
    }

    return response.json();
  },

  // User management (admin only)
  async getUsers() {
    const response = await apiRequest('/auth/users');

    if (!response.ok) {
      throw new Error('Failed to get users');
    }

    return response.json();
  },

  async createUser(userData: { username: string; email?: string; password: string; role?: string }) {
    const response = await apiRequest('/auth/users', {
      method: 'POST',
      body: JSON.stringify(userData),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Failed to create user' }));
      throw new Error(error.message || 'Failed to create user');
    }

    return response.json();
  },

  async updateUser(userId: string, userData: { email?: string; role?: string; active?: boolean }) {
    const response = await apiRequest(`/auth/users/${userId}`, {
      method: 'PUT',
      body: JSON.stringify(userData),
    });

    if (!response.ok) {
      const error = await response.json().catch(() => ({ message: 'Failed to update user' }));
      throw new Error(error.message || 'Failed to update user');
    }

    return response.json();
  },

  async deleteUser(userId: string) {
    const response = await apiRequest(`/auth/users/${userId}`, {
      method: 'DELETE',
    });

    if (!response.ok) {
      throw new Error('Failed to delete user');
    }

    return response.json();
  }
};

export const apiClient = new ApiClient();