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

import 'package:web_socket_channel/web_socket_channel.dart';

import 'models/websocket_event.dart';
import 'models/progress_event.dart';
import 'models/execution_event.dart';
import 'types/callback_types.dart';
import 'utils/websocket_event_handler.dart';

class WebSocketManager {
  final String host;
  final String clientId;

  WebSocketChannel? _wsChannel;

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

  // Event callbacks
  final Map<WebSocketEventType, List<WebSocketEventCallback>>
      _typedEventCallbacks = {
    for (var type in WebSocketEventType.values) type: [],
  };
  final List<ProgressEventCallback> _progressEventCallbacks = [];
  final Map<String, List<PromptEventCallback>> _eventCallbacks = {
    'onPromptStart': [],
    'onPromptFinished': [],
  };

  WebSocketManager({required this.host, required this.clientId});

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

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

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

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

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

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

  /// Connects to the WebSocket for progress updates
  Future<void> connect() 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);
        }
      }

      // Handle progress updates
      if (event.eventType == WebSocketEventType.progress) {
        for (final callback in _progressEventCallbacks) {
          callback(ProgressEvent.fromJson(event.data));
        }
      }
    }, onError: (error) {
      print('WebSocket error: $error');
    }, onDone: () {
      print('WebSocket connection closed');
    });
  }

  /// Attempts to create a ProgressEvent from a WebSocketEvent
  void _tryCreateProgressEvent(WebSocketEvent event) {
    WebSocketEventHandler.tryCreateProgressEvent(
      event,
      _progressEventController,
      _progressEventCallbacks,
    );
  }

  /// Attempts to create an ExecutionEvent from a WebSocketEvent
  void _tryCreateExecutionEvent(WebSocketEvent event) {
    WebSocketEventHandler.tryCreateExecutionEvent(
      event,
      _executionEventController,
    );
  }

  /// Trigger the onPromptStart callbacks
  void triggerOnPromptStart(String promptId) {
    for (final callback in _eventCallbacks['onPromptStart']!) {
      callback(promptId);
    }
  }

  /// Trigger the onPromptFinished callbacks
  void triggerOnPromptFinished(String promptId) {
    for (final callback in _eventCallbacks['onPromptFinished']!) {
      callback(promptId);
    }
  }

  /// Trigger the onProgressChanged callbacks
  void triggerOnProgressChanged(Map<String, dynamic> progressData) {
    for (final callback in _progressEventCallbacks) {
      callback(ProgressEvent.fromJson(progressData));
    }
  }

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