buzz/mobile/lib/features/pairing/pairing_socket.dart

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

import 'package:web_socket_channel/web_socket_channel.dart';

import '../../shared/relay/nostr_models.dart';
import 'package:buzz/shared/crypto/event_signing.dart';

const _desktopPairingAuthChallengeGrace = Duration(seconds: 3);
const _pairingAuthOkTimeout = Duration(seconds: 8);

/// Ephemeral WebSocket connection for NIP-AB pairing.
///
/// Uses ephemeral keys for NIP-42 auth (not the stored user keys).
/// Single-use — disposed after the pairing session completes.
class PairingAuthException implements Exception {
  final String message;

  const PairingAuthException(this.message);

  @override
  String toString() => 'PairingAuthException: $message';
}

class PairingSocket {
  final String _wsUrl;
  final String _ephemeralPrivkey;
  final void Function(List<dynamic> message) _onMessage;
  final void Function(Object? error) _onDisconnected;
  final Duration _authChallengeTimeout;
  final Duration _authResponseTimeout;
  final WebSocketChannel Function(Uri uri) _channelFactory;

  WebSocketChannel? _channel;
  StreamSubscription<dynamic>? _subscription;
  Completer<void>? _authCompleter;
  Timer? _authChallengeTimer;
  Timer? _authResponseTimer;
  String? _pendingAuthEventId;
  bool _connected = false;
  final Completer<void> _closed = Completer<void>();

  PairingSocket({
    required String wsUrl,
    required String ephemeralPrivkey,
    required void Function(List<dynamic> message) onMessage,
    required void Function(Object? error) onDisconnected,
    Duration authChallengeTimeout = _desktopPairingAuthChallengeGrace,
    Duration authResponseTimeout = _pairingAuthOkTimeout,
    WebSocketChannel Function(Uri uri) channelFactory =
        WebSocketChannel.connect,
  }) : _wsUrl = wsUrl,
       _ephemeralPrivkey = ephemeralPrivkey,
       _onMessage = onMessage,
       _onDisconnected = onDisconnected,
       _authChallengeTimeout = authChallengeTimeout,
       _authResponseTimeout = authResponseTimeout,
       _channelFactory = channelFactory;

  bool get isConnected => _connected;

  /// Connect and answer a NIP-42 challenge when the relay requires one.
  Future<void> connect() async {
    if (_closed.isCompleted) {
      throw const PairingAuthException('Pairing connection cancelled');
    }
    final channel = _channelFactory(Uri.parse(_wsUrl));
    _channel = channel;
    try {
      await Future.any([
        channel.ready,
        _closed.future,
      ]).timeout(const Duration(seconds: 10));
    } catch (_) {
      dispose();
      rethrow;
    }
    if (_closed.isCompleted) {
      throw const PairingAuthException('Pairing connection cancelled');
    }

    _authCompleter = Completer<void>();

    _subscription = _channel!.stream.listen(
      _handleRawMessage,
      onError: _failAuth,
      onDone: () => _failAuth(null),
    );

    // Dedicated pairing relays may be open and send no NIP-42 challenge.
    _authChallengeTimer = Timer(_authChallengeTimeout, () {
      if (_pendingAuthEventId == null &&
          _authCompleter != null &&
          !_authCompleter!.isCompleted) {
        _authCompleter!.complete();
      }
    });

    try {
      await _authCompleter!.future;
      if (_closed.isCompleted) {
        throw const PairingAuthException('Pairing connection cancelled');
      }
      _connected = true;
    } catch (error) {
      final deliberatelyClosed = _closed.isCompleted;
      await disconnect();
      if (!deliberatelyClosed) _onDisconnected(error);
      rethrow;
    } finally {
      _authChallengeTimer?.cancel();
      _authResponseTimer?.cancel();
    }
  }

  /// Send a raw JSON array.
  void send(List<dynamic> payload) {
    _channel?.sink.add(jsonEncode(payload));
  }

  /// Send a subscribe request.
  void subscribe(String subId, int kind, String pubkeyHex) {
    send([
      'REQ',
      subId,
      {
        'kinds': [kind],
        '#p': [pubkeyHex],
      },
    ]);
  }

  /// Publish a Nostr event (already JSON-encoded map).
  void publishEvent(Map<String, dynamic> event) {
    send(['EVENT', event]);
  }

  Future<void> disconnect() async {
    _retireConnection();
    _connected = false;
    _subscription?.cancel();
    _subscription = null;
    _authChallengeTimer?.cancel();
    _authResponseTimer?.cancel();
    final channel = _channel;
    _channel = null;
    if (channel != null) {
      await channel.sink.close();
    }
  }

  void dispose() {
    _retireConnection();
    _connected = false;
    _subscription?.cancel();
    _channel?.sink.close();
    _channel = null;
    _authChallengeTimer?.cancel();
    _authResponseTimer?.cancel();
  }

  void _retireConnection() {
    if (!_closed.isCompleted) _closed.complete();
    final auth = _authCompleter;
    if (auth != null && !auth.isCompleted) auth.complete();
  }

  void _failAuth(Object? error) {
    final authError = error ?? Exception('Connection closed');
    if (_authCompleter != null && !_authCompleter!.isCompleted) {
      _authCompleter!.completeError(authError);
      return;
    }
    if (_connected) {
      unawaited(disconnect());
      _onDisconnected(authError);
    }
  }

  void _handleRawMessage(dynamic raw) {
    if (raw is! String) return;

    final List<dynamic> data;
    try {
      data = jsonDecode(raw) as List<dynamic>;
    } catch (_) {
      return;
    }

    if (data.isEmpty) return;
    final type = data[0] as String;

    switch (type) {
      case 'AUTH':
        _handleAuthChallenge(data);
      case 'OK':
        _handleOk(data);
      default:
        // Pass EVENT, EOSE, NOTICE upstream.
        _onMessage(data);
    }
  }

  void _handleAuthChallenge(List<dynamic> data) {
    if (data.length < 2) return;
    final challenge = data[1] as String;

    _authChallengeTimer?.cancel();
    _authResponseTimer?.cancel();
    _authResponseTimer = Timer(_authResponseTimeout, () {
      _failAuth(
        const PairingAuthException('Relay did not confirm authentication'),
      );
    });

    try {
      // Build NIP-42 auth event (kind:22242) with ephemeral keys.
      final tags = <List<String>>[
        ['relay', _wsUrl],
        ['challenge', challenge],
      ];

      final event = signNostrEvent(
        kind: EventKind.auth,
        content: '',
        tags: tags,
        secretKey: _ephemeralPrivkey,
        createdAt: DateTime.now().millisecondsSinceEpoch ~/ 1000,
      );

      _pendingAuthEventId = event.id;
      send(['AUTH', event.toMap()]);
    } catch (e) {
      _failAuth(e);
    }
  }

  void _handleOk(List<dynamic> data) {
    if (data.length < 3) return;
    final eventId = data[1] as String;
    final accepted = data[2] as bool;

    if (_pendingAuthEventId != null && eventId == _pendingAuthEventId) {
      _pendingAuthEventId = null;
      _authResponseTimer?.cancel();
      if (accepted) {
        if (_authCompleter != null && !_authCompleter!.isCompleted) {
          _authCompleter!.complete();
        }
      } else {
        final message = data.length > 3
            ? data[3] as String
            : 'Auth rejected by relay';
        _failAuth(PairingAuthException(message));
      }
      return;
    }

    // Pass non-auth OK upstream.
    _onMessage(data);
  }
}