import 'dart:async';
import 'dart:convert';

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

import 'models/websocket_event.dart';
import 'models/progress_event.dart';
import 'models/execution_event.dart';

/// Callback function type for prompt events
typedef PromptEventCallback = void Function(String promptId);

/// Callback function type for typed WebSocket events
typedef WebSocketEventCallback = void Function(WebSocketEvent event);

/// Callback function type for progress events
typedef ProgressEventCallback = void Function(ProgressEvent event);

/// A Dart SDK for interacting with the ComfyUI API
class ComfyUiApi {
  final String host;
  final String clientId;
  final http.Client _httpClient;
  WebSocketChannel? _wsChannel;

  // Controllers for different event streams
  final StreamController<WebSocketEvent> _eventController =
      StreamController.broadcast();
  final StreamController<Map<String, dynamic>> _progressController =
      StreamController.broadcast();

  // Add new controllers for specific event types
  final StreamController<ProgressEvent> _progressEventController =
      StreamController.broadcast();
  final StreamController<ExecutionEvent> _executionEventController =
      StreamController.broadcast();

  /// Stream of typed progress events
  Stream<ProgressEvent> get progressEvents => _progressEventController.stream;

  /// Stream of typed execution events
  Stream<ExecutionEvent> get executionEvents =>
      _executionEventController.stream;

  // Event callbacks
  final Map<String, List<PromptEventCallback>> _eventCallbacks = {
    'onPromptStart': [],
    'onPromptFinished': [],
  };

  final Map<WebSocketEventType, List<WebSocketEventCallback>>
      _typedEventCallbacks = {
    for (var type in WebSocketEventType.values) type: [],
  };

  // Add a separate map for progress event callbacks
  final List<ProgressEventCallback> _progressEventCallbacks = [];

  /// Stream of typed WebSocket events
  Stream<WebSocketEvent> get events => _eventController.stream;

  /// Stream of progress updates from ComfyUI (legacy format)
  Stream<Map<String, dynamic>> get progressUpdates =>
      _progressController.stream;

  /// Creates a new ComfyUI API client
  ///
  /// [host] The host of the ComfyUI server (e.g. 'http://localhost:7860')
  /// [clientId] Optional client ID, will be automatically generated if not provided
  ComfyUiApi({
    required this.host,
    String? clientId,
    http.Client? httpClient,
  })  : clientId = clientId ?? const Uuid().v4(),
        _httpClient = httpClient ?? http.Client();

  /// Register a callback for when a prompt starts processing
  void onPromptStart(PromptEventCallback callback) {
    _eventCallbacks['onPromptStart']!.add(callback);
  }

  /// Register a callback for when a prompt finishes processing
  void onPromptFinished(PromptEventCallback callback) {
    _eventCallbacks['onPromptFinished']!.add(callback);
  }

  /// Register a callback for specific WebSocket event types
  void onEventType(WebSocketEventType type, WebSocketEventCallback callback) {
    _typedEventCallbacks[type]!.add(callback);
  }

  /// Register a callback for progress updates
  void onProgressChanged(ProgressEventCallback callback) {
    _progressEventCallbacks.add(callback);
  }

  /// Connects to the WebSocket for progress updates
  Future<void> connectWebSocket() async {
    final wsUrl =
        'ws://${host.replaceFirst(RegExp(r'^https?://'), '')}/ws?clientId=$clientId';
    _wsChannel = WebSocketChannel.connect(Uri.parse(wsUrl));

    _wsChannel!.stream.listen((message) {
      final jsonData = jsonDecode(message);

      // Create a typed event
      final event = WebSocketEvent.fromJson(jsonData);

      // Add to the typed event stream
      _eventController.add(event);

      // Also add to the legacy progress stream
      _progressController.add(jsonData);

      // Convert to more specific event types if possible
      if (event.eventType == WebSocketEventType.progress) {
        _tryCreateProgressEvent(event);
      } else if ([
        WebSocketEventType.executionStart,
        WebSocketEventType.executionSuccess,
        WebSocketEventType.executionCached,
        WebSocketEventType.executed,
        WebSocketEventType.executing,
      ].contains(event.eventType)) {
        _tryCreateExecutionEvent(event);
      }

      // Trigger event type specific callbacks
      for (final callback in _typedEventCallbacks[event.eventType]!) {
        callback(event);
      }

      // Handle execution_success event (prompt finished)
      if (event.eventType == WebSocketEventType.executionSuccess &&
          event.promptId != null) {
        final promptId = event.promptId!;
        for (final callback in _eventCallbacks['onPromptFinished']!) {
          callback(promptId);
        }
      }
    }, onError: (error) {
      print('WebSocket error: $error');
    }, onDone: () {
      print('WebSocket connection closed');
    });
  }

  /// Attempts to create a ProgressEvent from a WebSocketEvent
  void _tryCreateProgressEvent(WebSocketEvent event) {
    if (event.data.containsKey('value') &&
        event.data.containsKey('max') &&
        event.data.containsKey('prompt_id')) {
      try {
        final progressEvent = ProgressEvent(
          value: event.data['value'] as int,
          max: event.data['max'] as int,
          promptId: event.data['prompt_id'] as String,
          node: event.data['node']?.toString(),
        );
        _progressEventController.add(progressEvent);

        // Trigger all registered progress callbacks
        for (final callback in _progressEventCallbacks) {
          callback(progressEvent);
        }
      } catch (e) {
        print('Error creating ProgressEvent: $e');
      }
    }
  }

  /// Attempts to create an ExecutionEvent from a WebSocketEvent
  void _tryCreateExecutionEvent(WebSocketEvent event) {
    if (event.data.containsKey('prompt_id')) {
      try {
        final executionEvent = ExecutionEvent(
          promptId: event.data['prompt_id'] as String,
          timestamp: event.data['timestamp'] as int? ??
              DateTime.now().millisecondsSinceEpoch,
          node: event.data['node']?.toString(),
          extra: event.data['extra'] as Map<String, dynamic>?,
        );
        _executionEventController.add(executionEvent);
      } catch (e) {
        print('Error creating ExecutionEvent: $e');
      }
    }
  }

  /// Closes the WebSocket connection and cleans up resources
  void dispose() {
    _wsChannel?.sink.close();
    _progressController.close();
    _eventController.close();
    _progressEventController.close();
    _executionEventController.close();
    _httpClient.close();
  }

  /// Gets the current queue status
  Future<Map<String, dynamic>> getQueue() async {
    final response = await _httpClient.get(Uri.parse('$host/queue'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets the history of the queue
  Future<Map<String, dynamic>> getHistory({int maxItems = 64}) async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/history?max_items=$maxItems'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets image data by filename
  Future<List<int>> getImage(String filename) async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/view?filename=$filename'));
    _validateResponse(response);
    return response.bodyBytes;
  }

  /// Gets a list of all available models
  Future<Map<String, dynamic>> getModels() async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/experiment/models'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of checkpoints
  Future<List<dynamic>> getCheckpoints() async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/experiment/models/checkpoints'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific checkpoint
  Future<Map<String, dynamic>> getCheckpointDetails(
      String pathAndFileName) async {
    final response = await _httpClient.get(Uri.parse(
        '$host/api/view_metadata/checkpoints?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of LoRAs
  Future<List<dynamic>> getLoras() async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/experiment/models/loras'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific LoRA
  Future<Map<String, dynamic>> getLoraDetails(String pathAndFileName) async {
    final response = await _httpClient.get(
        Uri.parse('$host/api/view_metadata/loras?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of VAEs
  Future<List<dynamic>> getVaes() async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/experiment/models/vae'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific VAE
  Future<Map<String, dynamic>> getVaeDetails(String pathAndFileName) async {
    final response = await _httpClient.get(
        Uri.parse('$host/api/view_metadata/vae?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of upscale models
  Future<List<dynamic>> getUpscaleModels() async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/experiment/models/upscale_models'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific upscale model
  Future<Map<String, dynamic>> getUpscaleModelDetails(
      String pathAndFileName) async {
    final response = await _httpClient.get(Uri.parse(
        '$host/api/view_metadata/upscale_models?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of embeddings
  Future<List<dynamic>> getEmbeddings() async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/experiment/models/embeddings'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific embedding
  Future<Map<String, dynamic>> getEmbeddingDetails(
      String pathAndFileName) async {
    final response = await _httpClient.get(Uri.parse(
        '$host/api/view_metadata/embeddings?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets information about all available objects (nodes)
  Future<Map<String, dynamic>> getObjectInfo() async {
    final response = await _httpClient.get(Uri.parse('$host/api/object_info'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Submits a prompt (workflow) to generate an image
  Future<Map<String, dynamic>> submitPrompt(Map<String, dynamic> prompt) async {
    final response = await _httpClient.post(
      Uri.parse('$host/api/prompt'),
      headers: {'Content-Type': 'application/json'},
      body: jsonEncode(prompt),
    );
    _validateResponse(response);
    final responseData = jsonDecode(response.body);

    // Trigger onPromptStart event if prompt_id exists
    if (responseData.containsKey('prompt_id')) {
      final promptId = responseData['prompt_id'];
      for (final callback in _eventCallbacks['onPromptStart']!) {
        callback(promptId);
      }
    }

    return responseData;
  }

  /// Validates HTTP response and throws an exception if needed
  void _validateResponse(http.Response response) {
    if (response.statusCode < 200 || response.statusCode >= 300) {
      throw ComfyUiApiException(
          statusCode: response.statusCode,
          message: 'API request failed: ${response.body}');
    }
  }
}

/// Exception thrown when the ComfyUI API returns an error
class ComfyUiApiException implements Exception {
  final int statusCode;
  final String message;

  ComfyUiApiException({required this.statusCode, required this.message});

  @override
  String toString() => 'ComfyUiApiException: $statusCode - $message';
}
