TermRemoteCtl/apps/mobile_app/lib/features/terminal/terminal_socket_session.dart

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

import 'package:flutter_riverpod/flutter_riverpod.dart';
import 'package:web_socket_channel/web_socket_channel.dart';

import '../../core/network/agent_api_client.dart';
import '../../core/network/agent_socket_client.dart';
import '../sessions/session.dart';
import 'terminal_output_payload.dart';
import 'terminal_restore_payload.dart';

typedef TerminalSocketTransportFactory =
    TerminalSocketTransport Function(Uri uri);

final terminalSocketSessionFactoryProvider =
    Provider<TerminalSocketSessionFactory>((ref) {
      return TerminalSocketSessionFactory(
        transportFactory: WebSocketTerminalSocketTransport.connect,
      );
    });

class TerminalSocketSessionFactory {
  TerminalSocketSessionFactory({
    TerminalSocketTransportFactory? transportFactory,
  }) : _transportFactory =
           transportFactory ?? WebSocketTerminalSocketTransport.connect;

  final TerminalSocketTransportFactory _transportFactory;

  TerminalSocketSession create({
    required Uri baseUri,
    required Session session,
  }) {
    return TerminalSocketSession(
      sessionId: session.sessionId,
      socketClient: AgentSocketClient(baseUri),
      transportFactory: _transportFactory,
    );
  }
}

class TerminalSocketSession {
  TerminalSocketSession({
    required this.sessionId,
    required this.socketClient,
    TerminalSocketTransportFactory? transportFactory,
  }) : _transportFactory =
           transportFactory ?? WebSocketTerminalSocketTransport.connect;

  final String sessionId;
  final AgentSocketClient socketClient;
  final TerminalSocketTransportFactory _transportFactory;

  TerminalSocketTransport? _transport;
  StreamSubscription<dynamic>? _subscription;
  bool _isAttached = false;

  Future<void> connect({
    required void Function(TerminalOutputPayload output) onOutput,
    required void Function(TerminalRestorePayload restore) onRestore,
    void Function(String inputId)? onInputAck,
    void Function()? onDisconnected,
  }) async {
    if (_transport != null || _subscription != null) {
      await dispose();
    }

    final transport = _transportFactory(
      socketClient.buildTerminalSocketUri(sessionId),
    );
    _transport = transport;

    final attachedCompleter = Completer<void>();
    _subscription = transport.stream.listen(
      (message) {
        if (message is! String) {
          return;
        }

        if (!_isAttached && _handleAttachedAck(message)) {
          _isAttached = true;
          if (!attachedCompleter.isCompleted) {
            attachedCompleter.complete();
          }
          return;
        }

        if (_handleRestoreFrame(message, onRestore)) {
          return;
        }

        if (_handleInputAckFrame(message, onInputAck)) {
          return;
        }

        final output = _decodeOutputFrame(message);
        if (output != null) {
          onOutput(output);
        }
      },
      onError: (error, stackTrace) {
        final wasAttached = _isAttached;
        _handleTransportClosed(transport);
        if (!attachedCompleter.isCompleted) {
          attachedCompleter.completeError(error, stackTrace);
          return;
        }

        if (wasAttached) {
          onDisconnected?.call();
        }
      },
      onDone: () {
        final wasAttached = _isAttached;
        _handleTransportClosed(transport);
        if (!attachedCompleter.isCompleted) {
          attachedCompleter.completeError(
            StateError('Terminal socket closed before attach acknowledgement.'),
          );
          return;
        }

        if (wasAttached) {
          onDisconnected?.call();
        }
      },
    );

    transport.send(jsonEncode(socketClient.buildAttachMessage(sessionId)));

    try {
      await attachedCompleter.future;
    } catch (_) {
      await dispose();
      rethrow;
    }
  }

  TerminalSocketDispatchResult sendInput(String input, {String? inputId}) {
    final transport = _transport;
    if (input.isEmpty) {
      return TerminalSocketDispatchResult.emptyInput;
    }

    if (transport == null) {
      return TerminalSocketDispatchResult.noTransport;
    }

    try {
      transport.send(
        jsonEncode(socketClient.buildInputMessage(input, inputId: inputId)),
      );
      return TerminalSocketDispatchResult.sent;
    } catch (_) {
      _handleTransportClosed(transport);
      return TerminalSocketDispatchResult.noTransport;
    }
  }

  void sendResize(int columns, int rows) {
    final transport = _transport;
    if (transport == null || columns <= 0 || rows <= 0) {
      return;
    }

    try {
      transport.send(
        jsonEncode(socketClient.buildResizeMessage(columns, rows)),
      );
    } catch (_) {
      _handleTransportClosed(transport);
    }
  }

  Future<void> dispose() async {
    final subscription = _subscription;
    final transport = _transport;
    _handleTransportClosed(transport);
    await subscription?.cancel();
    try {
      await transport?.close();
    } catch (_) {}
  }

  bool _handleAttachedAck(String frame) {
    try {
      final decoded = jsonDecode(frame);
      if (decoded is Map &&
          decoded['type'] == 'attached' &&
          _matchesSessionId(decoded)) {
        return true;
      }
    } catch (_) {}

    return false;
  }

  bool _handleRestoreFrame(
    String frame,
    void Function(TerminalRestorePayload restore) onRestore,
  ) {
    try {
      final decoded = jsonDecode(frame);
      if (decoded is Map &&
          decoded['type'] == 'restore' &&
          _matchesSessionId(decoded)) {
        onRestore(
          TerminalRestorePayload.fromJson(Map<String, dynamic>.from(decoded)),
        );
        return true;
      }
    } catch (_) {}

    return false;
  }

  bool _handleInputAckFrame(
    String frame,
    void Function(String inputId)? onInputAck,
  ) {
    if (onInputAck == null) {
      return false;
    }

    try {
      final decoded = jsonDecode(frame);
      if (decoded is Map &&
          decoded['type'] == 'inputAck' &&
          _matchesSessionId(decoded)) {
        final inputId = decoded['inputId'];
        if (inputId is String && inputId.isNotEmpty) {
          onInputAck(inputId);
          return true;
        }
      }
    } catch (_) {}

    return false;
  }

  TerminalOutputPayload? _decodeOutputFrame(String frame) {
    try {
      final decoded = jsonDecode(frame);
      if (decoded is Map &&
          decoded['type'] == 'output' &&
          _matchesSessionId(decoded)) {
        return TerminalOutputPayload.fromJson(
          Map<String, dynamic>.from(decoded),
        );
      }

      if (decoded is Map && decoded['type'] != null) {
        return null;
      }
    } catch (_) {}

    return TerminalOutputPayload(
      sessionId: sessionId,
      sequence: 0,
      chunk: frame,
    );
  }

  bool _matchesSessionId(Map decoded) {
    final messageSessionId = decoded['sessionId'];
    return messageSessionId is String && messageSessionId == sessionId;
  }

  void _handleTransportClosed(TerminalSocketTransport? transport) {
    if (transport == null) {
      _subscription = null;
      _transport = null;
      _isAttached = false;
      return;
    }

    if (identical(_transport, transport)) {
      _subscription = null;
      _transport = null;
      _isAttached = false;
    }
  }
}

enum TerminalSocketDispatchResult { sent, noTransport, emptyInput }

abstract class TerminalSocketTransport {
  Stream<dynamic> get stream;
  void send(String message);
  Future<void> close();
}

class WebSocketTerminalSocketTransport implements TerminalSocketTransport {
  WebSocketTerminalSocketTransport(this._channel);

  final WebSocketChannel _channel;

  static WebSocketTerminalSocketTransport connect(Uri uri) {
    return WebSocketTerminalSocketTransport(WebSocketChannel.connect(uri));
  }

  @override
  Stream<dynamic> get stream => _channel.stream;

  @override
  void send(String message) {
    _channel.sink.add(message);
  }

  @override
  Future<void> close() {
    return _channel.sink.close();
  }
}