comfyui_api_sdk/lib/src/comfyui_api.dart

import 'dart:async';

import 'package:http/http.dart' as http;

import 'package:logger/logger.dart';
import 'models/websocket_event.dart';
import 'models/progress_event.dart';
import 'models/execution_event.dart';
import 'types/callback_types.dart';
import 'models/history_response.dart';
import 'models/checkpoint.dart';
import 'models/vae.dart';
import 'models/lora.dart';
import 'websocket_manager.dart';
import 'models/submit_prompt_response.dart';
import '_http_endpoints.dart'; // Import the new helper class

/// A Dart SDK for interacting with the ComfyUI API
class ComfyUiApi {
  final String host;
  final String clientId;
  final http.Client _httpClient; // Keep for potential direct use or disposal
  final WebSocketManager _webSocketManager;
  final Logger logger;
  late final ComfyApiHttpEndpoints _httpEndpoints; // Add the helper instance

  /// Creates a new ComfyUI API client
  ComfyUiApi({
    required this.host,
    required this.clientId,
    Logger? logger,
    http.Client? httpClient,
  })  : _httpClient = httpClient ?? http.Client(),
        _webSocketManager = WebSocketManager(host: host, clientId: clientId),
        logger = logger ?? Logger() {
    _httpEndpoints = ComfyApiHttpEndpoints(
      host: host,
      httpClient: _httpClient,
      logger: this.logger, // Use the instance logger
      webSocketManager: _webSocketManager,
    );
  }

  /// Expose WebSocketManager streams and methods
  Stream<WebSocketEvent> get events => _webSocketManager.events;
  Stream<Map<String, dynamic>> get progressUpdates =>
      _webSocketManager.progressUpdates;
  Stream<ProgressEvent> get progressEvents => _webSocketManager.progressEvents;
  Stream<ExecutionEvent> get executionEvents =>
      _webSocketManager.executionEvents;
  Stream<int> get executingNodeStream => _webSocketManager.executingNodeStream;
  Stream<void> get executionInterruptedStream =>
      _webSocketManager.executionInterruptedStream;

  /// Access to the WebSocketManager instance
  WebSocketManager get webSocketManager => _webSocketManager;

  void onExecutingNodeChanged(void Function(int nodeId) callback) {
    _webSocketManager.executingNodeStream.listen(callback);
  }

  void onEventType(WebSocketEventType type, WebSocketEventCallback callback) {
    _webSocketManager.onEventType(type, callback);
  }

  void onProgressChanged(ProgressEventCallback callback) {
    _webSocketManager.onProgressChanged(callback);
  }

  void onPromptStart(PromptEventCallback callback) {
    _webSocketManager.onPromptStart(callback);
  }

  void onPromptFinished(PromptEventCallback callback) {
    _webSocketManager.onPromptFinished(callback);
  }

  /// Register a callback for when execution is interrupted
  void onExecutionInterrupted(void Function() callback) {
    _webSocketManager.executionInterruptedStream.listen((_) => callback());
  }

  Future<void> connectWebSocket() => _webSocketManager.connect();

  void dispose() {
    _webSocketManager.dispose();
    _httpClient.close();
  }

  // --- HTTP Endpoint Methods (Delegated) ---

  /// Gets the current queue status
  Future<Map<String, dynamic>> getQueue() => _httpEndpoints.getQueue();

  /// Gets the history of the queue
  Future<HistoryResponse> getHistory({int maxItems = 64}) =>
      _httpEndpoints.getHistory(maxItems: maxItems);

  /// Gets image data by filename, optionally specifying subfolder and type
  Future<List<int>> getImage(String filename,
          {String? subfolder, String? type}) =>
      _httpEndpoints.getImage(filename, subfolder: subfolder, type: type);

  /// Gets a list of all available models
  Future<Map<String, dynamic>> getModels() => _httpEndpoints.getModels();

  /// Gets a list of checkpoints
  Future<List<Checkpoint>> getCheckpoints() => _httpEndpoints.getCheckpoints();

  /// Gets details for a specific checkpoint
  Future<Map<String, dynamic>> getCheckpointDetails(String pathAndFileName) =>
      _httpEndpoints.getCheckpointDetails(pathAndFileName);

  /// Gets a list of LoRAs
  Future<List<Lora>> getLoras() => _httpEndpoints.getLoras();

  /// Gets details for a specific LoRA
  Future<Map<String, dynamic>> getLoraDetails(String pathAndFileName) =>
      _httpEndpoints.getLoraDetails(pathAndFileName);

  /// Gets a list of VAEs
  Future<List<Vae>> getVaes() => _httpEndpoints.getVaes();

  /// Gets details for a specific VAE
  Future<Map<String, dynamic>> getVaeDetails(String pathAndFileName) =>
      _httpEndpoints.getVaeDetails(pathAndFileName);

  /// Gets a list of upscale models
  Future<List<dynamic>> getUpscaleModels() => _httpEndpoints.getUpscaleModels();

  /// Gets details for a specific upscale model
  Future<Map<String, dynamic>> getUpscaleModelDetails(String pathAndFileName) =>
      _httpEndpoints.getUpscaleModelDetails(pathAndFileName);

  /// Gets a list of embeddings
  Future<List<dynamic>> getEmbeddings() => _httpEndpoints.getEmbeddings();

  /// Gets details for a specific embedding
  Future<Map<String, dynamic>> getEmbeddingDetails(String pathAndFileName) =>
      _httpEndpoints.getEmbeddingDetails(pathAndFileName);

  /// Gets information about all available objects (nodes)
  Future<Map<String, dynamic>> getObjectInfo() =>
      _httpEndpoints.getObjectInfo();

  /// Gets a list of possible samplers from the object info
  Future<List<String>> getKSamplers() => _httpEndpoints.getKSamplers();

  /// Gets a list of possible schedulers from the object info
  Future<List<String>> getSchedulers() => _httpEndpoints.getSchedulers();

  /// Submits a prompt (workflow) to generate an image
  Future<SubmitPromptResponse> submitPrompt(Map<String, dynamic> prompt) =>
      _httpEndpoints.submitPrompt(prompt);

  /// Interrupts the current execution and optionally clears the queue
  Future<bool> interrupt({bool clearQueue = false}) =>
      _httpEndpoints.interrupt(clearQueue: clearQueue);
}