abawo-bt-app/lib/service/firmware_update_service.dart

import 'dart:async';

import 'package:abawo_bt_app/controller/bluetooth.dart';
import 'package:abawo_bt_app/model/shifter_types.dart';
import 'package:abawo_bt_app/service/dfu_protocol.dart';
import 'package:abawo_bt_app/service/shifter_service.dart';
import 'package:anyhow/anyhow.dart';

const int _initialAckSequence = 0xFF;

class FirmwareUpdateService {
  FirmwareUpdateService({
    required FirmwareUpdateTransport transport,
    this.defaultWindowSize = 8,
    this.maxNoProgressRetries = 5,
    this.defaultAckTimeout = const Duration(milliseconds: 800),
    this.defaultPostFinishResetTimeout = const Duration(seconds: 8),
    this.defaultReconnectTimeout = const Duration(seconds: 12),
    this.defaultVerificationTimeout = const Duration(seconds: 5),
  }) : _transport = transport;

  final FirmwareUpdateTransport _transport;
  final int defaultWindowSize;
  final int maxNoProgressRetries;
  final Duration defaultAckTimeout;
  final Duration defaultPostFinishResetTimeout;
  final Duration defaultReconnectTimeout;
  final Duration defaultVerificationTimeout;

  final StreamController<DfuUpdateProgress> _progressController =
      StreamController<DfuUpdateProgress>.broadcast();

  DfuUpdateProgress _currentProgress = const DfuUpdateProgress(
    state: DfuUpdateState.idle,
    totalBytes: 0,
    sentBytes: 0,
    lastAckedSequence: _initialAckSequence,
    sessionId: 0,
    flags: DfuUpdateFlags(),
  );

  StreamSubscription<List<int>>? _ackSubscription;
  Completer<void>? _ackSignal;
  Completer<void>? _cancelSignal;
  int _ackEventCount = 0;
  String? _ackStreamError;
  bool _isRunning = false;
  bool _cancelRequested = false;
  int _latestAckSequence = _initialAckSequence;
  int _ackedFrames = 0;
  int _totalFrames = 0;
  int _totalBytes = 0;

  Stream<DfuUpdateProgress> get progressStream => _progressController.stream;

  DfuUpdateProgress get currentProgress => _currentProgress;

  bool get isUpdating => _isRunning;

  Future<Result<void>> startUpdate({
    required List<int> imageBytes,
    required int sessionId,
    DfuUpdateFlags flags = const DfuUpdateFlags(),
    int requestedMtu = universalShifterDfuPreferredMtu,
    int? windowSize,
    Duration? ackTimeout,
    int? noProgressRetries,
    Duration? postFinishResetTimeout,
    Duration? reconnectTimeout,
    Duration? verificationTimeout,
  }) async {
    if (_isRunning) {
      return bail(
          'Firmware update is already running. Cancel or wait for completion before starting a new upload.');
    }
    if (imageBytes.isEmpty) {
      return bail(
          'Firmware image is empty. Select a valid .bin file and retry.');
    }

    final effectiveWindowSize = windowSize ?? defaultWindowSize;
    final effectiveAckTimeout = ackTimeout ?? defaultAckTimeout;
    final effectiveNoProgressRetries =
        noProgressRetries ?? maxNoProgressRetries;
    final effectivePostFinishResetTimeout =
        postFinishResetTimeout ?? defaultPostFinishResetTimeout;
    final effectiveReconnectTimeout =
        reconnectTimeout ?? defaultReconnectTimeout;
    final effectiveVerificationTimeout =
        verificationTimeout ?? defaultVerificationTimeout;

    if (effectiveWindowSize <= 0) {
      return bail(
          'DFU window size must be at least 1 frame. Got $effectiveWindowSize.');
    }
    if (effectiveNoProgressRetries < 0) {
      return bail(
          'No-progress retry limit cannot be negative. Got $effectiveNoProgressRetries.');
    }

    _isRunning = true;
    _cancelRequested = false;
    _cancelSignal = Completer<void>();
    _ackSignal = null;
    _ackEventCount = 0;
    _ackStreamError = null;
    _latestAckSequence = _initialAckSequence;
    _ackedFrames = 0;
    _totalFrames =
        (imageBytes.length + universalShifterDfuFramePayloadSizeBytes - 1) ~/
            universalShifterDfuFramePayloadSizeBytes;
    _totalBytes = imageBytes.length;

    final normalizedSessionId = sessionId & 0xFF;
    final crc32 = DfuProtocol.crc32(imageBytes);
    final frames = DfuProtocol.buildDataFrames(imageBytes);
    var shouldAbortForCleanup = false;

    _emitProgress(
      state: DfuUpdateState.starting,
      totalBytes: imageBytes.length,
      sentBytes: 0,
      lastAckedSequence: _initialAckSequence,
      sessionId: normalizedSessionId,
      flags: flags,
    );

    try {
      final preflightResult = await _transport.runPreflight(
        requestedMtu: requestedMtu,
      );
      if (preflightResult.isErr()) {
        throw _DfuFailure(
          'DFU preflight check failed due to transport error: ${preflightResult.unwrapErr()}',
        );
      }
      final preflight = preflightResult.unwrap();
      if (!preflight.canStart) {
        throw _DfuFailure(
          preflight.message ??
              'DFU preflight failed. Ensure button connection and MTU are ready, then retry.',
        );
      }

      await _ackSubscription?.cancel();
      _ackSubscription = _transport.subscribeToAck().listen(
        _handleAckPayload,
        onError: (Object error) {
          _ackStreamError =
              'ACK indication stream failed: $error. Reconnect and retry the update.';
          _signalAckWaiters();
        },
      );

      _emitProgress(state: DfuUpdateState.waitingForAck);
      final startEventCount = _ackEventCount;
      final startWriteResult = await _transport.writeControl(
        DfuProtocol.encodeStartPayload(
          DfuStartPayload(
            totalLength: imageBytes.length,
            imageCrc32: crc32,
            sessionId: normalizedSessionId,
            flags: flags.rawValue,
          ),
        ),
      );
      if (startWriteResult.isErr()) {
        throw _DfuFailure(
          'Failed to send DFU START command: ${startWriteResult.unwrapErr()}',
        );
      }
      shouldAbortForCleanup = true;

      final initialAck = await _waitForInitialAck(
        afterEventCount: startEventCount,
        timeout: effectiveAckTimeout,
      );
      if (initialAck != _initialAckSequence) {
        throw _DfuFailure(
          'Device did not acknowledge START correctly (expected ACK 0xFF, got 0x${initialAck.toRadixString(16).padLeft(2, '0').toUpperCase()}). Send ABORT, reconnect if needed, and retry.',
        );
      }

      _emitProgress(state: DfuUpdateState.transferring);

      var nextFrameIndex = 0;
      var retriesWithoutProgress = 0;

      while (_ackedFrames < _totalFrames) {
        _throwIfCancelled();
        _throwIfAckStreamErrored();

        final ackedBeforeWindow = _ackedFrames;
        final endExclusive =
            (nextFrameIndex + effectiveWindowSize).clamp(0, frames.length);

        for (var frameIndex = nextFrameIndex;
            frameIndex < endExclusive;
            frameIndex++) {
          _throwIfCancelled();
          final writeResult =
              await _transport.writeDataFrame(frames[frameIndex].bytes);
          if (writeResult.isErr()) {
            throw _DfuFailure(
              'Failed sending DFU data frame #$frameIndex (seq 0x${frames[frameIndex].sequence.toRadixString(16).padLeft(2, '0').toUpperCase()}): ${writeResult.unwrapErr()}',
            );
          }
        }

        nextFrameIndex = endExclusive;

        if (_ackedFrames > ackedBeforeWindow) {
          retriesWithoutProgress = 0;
          nextFrameIndex = _ackedFrames;
          continue;
        }

        final gotProgress = await _waitForAckProgress(
          ackedFramesBeforeWait: ackedBeforeWindow,
          timeout: effectiveAckTimeout,
        );

        if (gotProgress) {
          retriesWithoutProgress = 0;
          nextFrameIndex = _ackedFrames;
          continue;
        }

        retriesWithoutProgress += 1;
        if (retriesWithoutProgress > effectiveNoProgressRetries) {
          throw _DfuFailure(
            'Upload stalled: no ACK progress after $retriesWithoutProgress retries (last ACK 0x${_latestAckSequence.toRadixString(16).padLeft(2, '0').toUpperCase()}). Check BLE signal quality and retry.',
          );
        }

        nextFrameIndex = _ackedFrames;
      }

      _emitProgress(
          state: DfuUpdateState.finishing, sentBytes: imageBytes.length);
      final finishResult =
          await _transport.writeControl(DfuProtocol.encodeFinishPayload());
      if (finishResult.isErr()) {
        throw _DfuFailure(
          'Failed to send DFU FINISH command: ${finishResult.unwrapErr()}',
        );
      }

      await _ackSubscription?.cancel();
      _ackSubscription = null;

      final resetDisconnectResult =
          await _transport.waitForExpectedResetDisconnect(
        timeout: effectivePostFinishResetTimeout,
      );
      if (resetDisconnectResult.isErr()) {
        throw _DfuFailure(
          'Device did not perform the expected post-FINISH reset disconnect: ${resetDisconnectResult.unwrapErr()}',
        );
      }

      final reconnectResult = await _transport.reconnectForVerification(
        timeout: effectiveReconnectTimeout,
      );
      if (reconnectResult.isErr()) {
        throw _DfuFailure(
          'Device did not reconnect after DFU reset: ${reconnectResult.unwrapErr()}',
        );
      }

      final verificationResult = await _transport.verifyDeviceReachable(
        timeout: effectiveVerificationTimeout,
      );
      if (verificationResult.isErr()) {
        throw _DfuFailure(
          'Device reconnected but post-update verification failed: ${verificationResult.unwrapErr()} '
          'Firmware version cannot be compared yet because the device does not expose a version characteristic.',
        );
      }

      shouldAbortForCleanup = false;
      _emitProgress(
          state: DfuUpdateState.completed, sentBytes: imageBytes.length);
      return Ok(null);
    } on _DfuCancelled {
      if (shouldAbortForCleanup) {
        await _sendAbortForCleanup();
      }
      _emitProgress(state: DfuUpdateState.aborted);
      return bail('Firmware update canceled by user.');
    } on _DfuFailure catch (failure) {
      if (shouldAbortForCleanup) {
        await _sendAbortForCleanup();
      }
      _emitProgress(
          state: DfuUpdateState.failed, errorMessage: failure.message);
      return bail(failure.message);
    } catch (error) {
      if (shouldAbortForCleanup) {
        await _sendAbortForCleanup();
      }
      final message =
          'Firmware update failed unexpectedly: $error. Reconnect to the button and retry.';
      _emitProgress(state: DfuUpdateState.failed, errorMessage: message);
      return bail(message);
    } finally {
      await _ackSubscription?.cancel();
      _ackSubscription = null;
      _isRunning = false;
      _cancelRequested = false;
      _cancelSignal = null;
      _ackSignal = null;
      _ackEventCount = 0;
      _ackStreamError = null;
      _latestAckSequence = _currentProgress.lastAckedSequence;
      _ackedFrames = 0;
      _totalFrames = 0;
      _totalBytes = 0;
    }
  }

  Future<void> cancelUpdate() async {
    if (!_isRunning || _cancelRequested) {
      return;
    }
    _cancelRequested = true;
    _cancelSignal?.complete();
    _signalAckWaiters();
  }

  Future<void> dispose() async {
    await cancelUpdate();
    await _ackSubscription?.cancel();
    _ackSubscription = null;
    await _progressController.close();
  }

  void _handleAckPayload(List<int> payload) {
    try {
      final sequence = DfuProtocol.parseAckPayload(payload);
      final previousAck = _latestAckSequence;
      _latestAckSequence = sequence;

      if (_totalFrames > 0 &&
          _currentProgress.state == DfuUpdateState.transferring) {
        final delta = DfuProtocol.sequenceDistance(previousAck, sequence);
        if (delta > 0) {
          _ackedFrames = (_ackedFrames + delta).clamp(0, _totalFrames);
        }

        _emitProgress(
          lastAckedSequence: sequence,
          sentBytes:
              _ackedBytesFromFrames(_ackedFrames, _totalFrames, _totalBytes),
        );
      } else {
        _emitProgress(lastAckedSequence: sequence);
      }
    } on FormatException catch (error) {
      _ackStreamError =
          'Received malformed ACK indication: $error. Reconnect and retry.';
    } finally {
      _ackEventCount += 1;
      _signalAckWaiters();
    }
  }

  void _emitProgress({
    DfuUpdateState? state,
    int? totalBytes,
    int? sentBytes,
    int? lastAckedSequence,
    int? sessionId,
    DfuUpdateFlags? flags,
    String? errorMessage,
  }) {
    final next = DfuUpdateProgress(
      state: state ?? _currentProgress.state,
      totalBytes: totalBytes ?? _currentProgress.totalBytes,
      sentBytes: sentBytes ?? _currentProgress.sentBytes,
      lastAckedSequence:
          lastAckedSequence ?? _currentProgress.lastAckedSequence,
      sessionId: sessionId ?? _currentProgress.sessionId,
      flags: flags ?? _currentProgress.flags,
      errorMessage: errorMessage,
    );
    _currentProgress = next;
    _progressController.add(next);
  }

  Future<int> _waitForInitialAck({
    required int afterEventCount,
    required Duration timeout,
  }) async {
    final deadline = DateTime.now().add(timeout);
    var observedEvents = afterEventCount;

    while (true) {
      _throwIfCancelled();
      _throwIfAckStreamErrored();
      final remaining = deadline.difference(DateTime.now());
      if (remaining <= Duration.zero) {
        throw _DfuFailure(
          'Timed out waiting for initial DFU ACK after START. Ensure indications are enabled and retry.',
        );
      }

      final gotEvent = await _waitForNextAckEvent(
        afterEventCount: observedEvents,
        timeout: remaining,
      );
      if (!gotEvent) {
        continue;
      }
      observedEvents = _ackEventCount;
      return _latestAckSequence;
    }
  }

  Future<bool> _waitForAckProgress({
    required int ackedFramesBeforeWait,
    required Duration timeout,
  }) async {
    final deadline = DateTime.now().add(timeout);
    var observedEvents = _ackEventCount;

    while (true) {
      _throwIfCancelled();
      _throwIfAckStreamErrored();

      if (_ackedFrames > ackedFramesBeforeWait) {
        return true;
      }

      final remaining = deadline.difference(DateTime.now());
      if (remaining <= Duration.zero) {
        return false;
      }

      final gotEvent = await _waitForNextAckEvent(
        afterEventCount: observedEvents,
        timeout: remaining,
      );

      if (!gotEvent) {
        continue;
      }

      observedEvents = _ackEventCount;
    }
  }

  Future<bool> _waitForNextAckEvent({
    required int afterEventCount,
    required Duration timeout,
  }) async {
    if (_ackEventCount > afterEventCount) {
      return true;
    }

    _ackSignal ??= Completer<void>();
    final signal = _ackSignal!;

    try {
      await Future.any<void>([
        signal.future,
        _cancelSignal?.future ?? Future<void>.value(),
      ]).timeout(timeout);
    } on TimeoutException {
      return false;
    }

    if (identical(_ackSignal, signal)) {
      _ackSignal = null;
    }

    _throwIfCancelled();
    _throwIfAckStreamErrored();
    return _ackEventCount > afterEventCount;
  }

  void _throwIfCancelled() {
    if (_cancelRequested) {
      throw const _DfuCancelled();
    }
  }

  void _throwIfAckStreamErrored() {
    final error = _ackStreamError;
    if (error != null) {
      throw _DfuFailure(error);
    }
  }

  Future<void> _sendAbortForCleanup() async {
    final result =
        await _transport.writeControl(DfuProtocol.encodeAbortPayload());
    if (result.isErr()) {
      final cleanupMessage =
          'Could not send DFU ABORT during cleanup: ${result.unwrapErr()}';
      if (_currentProgress.state == DfuUpdateState.failed &&
          _currentProgress.errorMessage != null) {
        _emitProgress(
          errorMessage: '${_currentProgress.errorMessage} $cleanupMessage',
        );
      }
    }
  }

  void _signalAckWaiters() {
    final signal = _ackSignal;
    if (signal != null && !signal.isCompleted) {
      signal.complete();
    }
  }

  int _ackedBytesFromFrames(int ackedFrames, int totalFrames, int totalBytes) {
    if (totalFrames == 0 || ackedFrames <= 0) {
      return 0;
    }
    if (ackedFrames >= totalFrames) {
      return totalBytes;
    }
    return ackedFrames * universalShifterDfuFramePayloadSizeBytes;
  }
}

abstract interface class FirmwareUpdateTransport {
  Future<Result<DfuPreflightResult>> runPreflight({required int requestedMtu});

  Stream<List<int>> subscribeToAck();

  Future<Result<void>> writeControl(List<int> payload);

  Future<Result<void>> writeDataFrame(List<int> frame);

  Future<Result<void>> waitForExpectedResetDisconnect({
    required Duration timeout,
  });

  Future<Result<void>> reconnectForVerification({
    required Duration timeout,
  });

  /// Verifies that the device is reachable after reconnect.
  ///
  /// Current limitation: strict firmware version comparison is not possible
  /// yet because no firmware version characteristic is exposed by the device.
  Future<Result<void>> verifyDeviceReachable({
    required Duration timeout,
  });
}

class ShifterFirmwareUpdateTransport implements FirmwareUpdateTransport {
  ShifterFirmwareUpdateTransport({
    required this.shifterService,
    required this.bluetoothController,
    required this.buttonDeviceId,
  });

  final ShifterService shifterService;
  final BluetoothController bluetoothController;
  final String buttonDeviceId;

  @override
  Future<Result<DfuPreflightResult>> runPreflight({
    required int requestedMtu,
  }) {
    return shifterService.runDfuPreflight(requestedMtu: requestedMtu);
  }

  @override
  Stream<List<int>> subscribeToAck() {
    return bluetoothController.subscribeToCharacteristic(
      buttonDeviceId,
      universalShifterControlServiceUuid,
      universalShifterDfuAckCharacteristicUuid,
    );
  }

  @override
  Future<Result<void>> writeControl(List<int> payload) {
    return bluetoothController.writeCharacteristic(
      buttonDeviceId,
      universalShifterControlServiceUuid,
      universalShifterDfuControlCharacteristicUuid,
      payload,
    );
  }

  @override
  Future<Result<void>> writeDataFrame(List<int> frame) {
    return bluetoothController.writeCharacteristic(
      buttonDeviceId,
      universalShifterControlServiceUuid,
      universalShifterDfuDataCharacteristicUuid,
      frame,
      withResponse: false,
    );
  }

  @override
  Future<Result<void>> waitForExpectedResetDisconnect({
    required Duration timeout,
  }) async {
    final currentState = bluetoothController.currentConnectionState;
    if (currentState.$1 == ConnectionStatus.disconnected) {
      return Ok(null);
    }

    try {
      await bluetoothController.connectionStateStream
          .firstWhere((state) => state.$1 == ConnectionStatus.disconnected)
          .timeout(timeout);
      return Ok(null);
    } on TimeoutException {
      return bail(
        'Timed out after ${timeout.inMilliseconds}ms waiting for the expected reset disconnect.',
      );
    } catch (error) {
      return bail('Failed while waiting for expected reset disconnect: $error');
    }
  }

  @override
  Future<Result<void>> reconnectForVerification({
    required Duration timeout,
  }) async {
    final connectResult =
        await bluetoothController.connectById(buttonDeviceId, timeout: timeout);
    if (connectResult.isErr()) {
      return bail(connectResult.unwrapErr());
    }

    final currentState = bluetoothController.currentConnectionState;
    if (currentState.$1 == ConnectionStatus.connected &&
        currentState.$2 == buttonDeviceId) {
      return Ok(null);
    }

    try {
      await bluetoothController.connectionStateStream
          .firstWhere(
            (state) =>
                state.$1 == ConnectionStatus.connected &&
                state.$2 == buttonDeviceId,
          )
          .timeout(timeout);
      return Ok(null);
    } on TimeoutException {
      return bail(
        'Timed out after ${timeout.inMilliseconds}ms waiting for reconnect.',
      );
    } catch (error) {
      return bail('Reconnect wait failed: $error');
    }
  }

  @override
  Future<Result<void>> verifyDeviceReachable({
    required Duration timeout,
  }) async {
    try {
      final statusResult = await shifterService.readStatus().timeout(timeout);
      if (statusResult.isErr()) {
        return bail(statusResult.unwrapErr());
      }
      return Ok(null);
    } on TimeoutException {
      return bail(
        'Timed out after ${timeout.inMilliseconds}ms while reading status for post-update verification.',
      );
    } catch (error) {
      return bail('Post-update verification failed: $error');
    }
  }
}

class _DfuFailure implements Exception {
  const _DfuFailure(this.message);

  final String message;

  @override
  String toString() => message;
}

class _DfuCancelled implements Exception {
  const _DfuCancelled();
}