abawo-bt-app/lib/pages/bootloader_recovery_update_page.dart

import 'dart:async';

import 'package:abawo_bt_app/controller/bluetooth.dart';
import 'package:abawo_bt_app/model/firmware_file_selection.dart';
import 'package:abawo_bt_app/model/shifter_types.dart';
import 'package:abawo_bt_app/service/firmware_update_service.dart';
import 'package:abawo_bt_app/widgets/firmware_update_fullscreen.dart';
import 'package:flutter/material.dart';
import 'package:flutter_riverpod/flutter_riverpod.dart';
import 'package:go_router/go_router.dart';

class BootloaderRecoveryUpdateArgs {
  const BootloaderRecoveryUpdateArgs({
    required this.bootloaderDeviceId,
    required this.firmware,
  });

  final String bootloaderDeviceId;
  final BootloaderDfuPreparedFirmware firmware;
}

class BootloaderRecoveryUpdatePage extends ConsumerStatefulWidget {
  const BootloaderRecoveryUpdatePage({
    required this.args,
    super.key,
  });

  final BootloaderRecoveryUpdateArgs args;

  @override
  ConsumerState<BootloaderRecoveryUpdatePage> createState() =>
      _BootloaderRecoveryUpdatePageState();
}

class _BootloaderRecoveryUpdatePageState
    extends ConsumerState<BootloaderRecoveryUpdatePage> {
  FirmwareUpdateService? _firmwareUpdateService;
  StreamSubscription<DfuUpdateProgress>? _firmwareProgressSubscription;
  DfuUpdateProgress _dfuProgress = const DfuUpdateProgress(
    state: DfuUpdateState.idle,
    totalBytes: 0,
    sentBytes: 0,
    expectedOffset: 0,
    sessionId: 0,
    flags: DfuUpdateFlags(),
  );
  String? _firmwareUserMessage = 'Preparing US-DFU recovery update...';
  bool _hasStarted = false;

  @override
  void initState() {
    super.initState();
    WidgetsBinding.instance.addPostFrameCallback((_) {
      if (mounted) {
        unawaited(_startUpdate());
      }
    });
  }

  @override
  void dispose() {
    final service = _firmwareUpdateService;
    final bluetooth = ref.read(bluetoothProvider).valueOrNull;
    _firmwareUpdateService = null;
    unawaited(_firmwareProgressSubscription?.cancel());
    unawaited(() async {
      await service?.dispose();
      await _disconnectBootloaderIfStillConnected(bluetooth: bluetooth);
    }());
    super.dispose();
  }

  Future<void> _disconnectBootloaderIfStillConnected({
    BluetoothController? bluetooth,
  }) async {
    bluetooth ??= ref.read(bluetoothProvider).valueOrNull;
    if (bluetooth == null) {
      return;
    }

    final currentState = bluetooth.currentConnectionState;
    if (currentState.$2 != widget.args.bootloaderDeviceId ||
        (currentState.$1 != ConnectionStatus.connected &&
            currentState.$1 != ConnectionStatus.connecting)) {
      return;
    }

    await bluetooth.disconnect();
  }

  Future<void> _dismissToDevices() async {
    await _disconnectBootloaderIfStillConnected();
    if (!mounted) {
      return;
    }
    context.go('/devices');
  }

  Future<FirmwareUpdateService?> _ensureFirmwareUpdateService() async {
    if (_firmwareUpdateService != null) {
      return _firmwareUpdateService;
    }

    final bluetooth = ref.read(bluetoothProvider).valueOrNull;
    if (bluetooth == null) {
      return null;
    }

    final service = FirmwareUpdateService(
      verifyAfterFinish: false,
      transport: ShifterFirmwareUpdateTransport(
        shifterService: null,
        bluetoothController: bluetooth,
        buttonDeviceId: widget.args.bootloaderDeviceId,
      ),
    );
    _firmwareProgressSubscription = service.progressStream.listen((progress) {
      if (!mounted) {
        return;
      }
      setState(() {
        _dfuProgress = progress;
        if (progress.state == DfuUpdateState.failed &&
            progress.errorMessage != null) {
          _firmwareUserMessage = progress.errorMessage;
        } else if (progress.state == DfuUpdateState.completed) {
          _firmwareUserMessage =
              'Firmware update completed. The bootloader accepted FINISH and reset; reconnect to the device when it starts advertising again.';
        } else if (progress.state == DfuUpdateState.aborted) {
          _firmwareUserMessage = 'Firmware update canceled.';
        } else if (progress.errorMessage != null) {
          _firmwareUserMessage = progress.errorMessage;
        }
      });
    });
    _firmwareUpdateService = service;
    return service;
  }

  Future<void> _startUpdate() async {
    if (_hasStarted) {
      return;
    }
    _hasStarted = true;

    final updater = await _ensureFirmwareUpdateService();
    if (!mounted) {
      return;
    }
    if (updater == null) {
      setState(() {
        _dfuProgress = DfuUpdateProgress(
          state: DfuUpdateState.failed,
          totalBytes: widget.args.firmware.fileBytes.length,
          sentBytes: 0,
          expectedOffset: 0,
          sessionId: widget.args.firmware.metadata.sessionId,
          flags: DfuUpdateFlags.fromRaw(widget.args.firmware.metadata.flags),
          errorMessage:
              'Firmware updater is not ready. Reconnect to US-DFU and retry.',
        );
        _firmwareUserMessage = _dfuProgress.errorMessage;
      });
      return;
    }

    final firmware = widget.args.firmware;
    final result = await updater.startUpdate(
      imageBytes: firmware.fileBytes,
      sessionId: firmware.metadata.sessionId,
      appStart: firmware.metadata.appStart,
      imageVersion: firmware.metadata.imageVersion,
      flags: DfuUpdateFlags.fromRaw(firmware.metadata.flags),
    );

    if (!mounted || result.isOk()) {
      return;
    }
    await _disconnectBootloaderIfStillConnected();
    if (!mounted) {
      return;
    }
    final errorMessage = result.unwrapErr().toString();
    setState(() {
      _firmwareUserMessage = errorMessage;
    });

    if (errorMessage.startsWith(universalShifterBootMetadataWarningMessage)) {
      ScaffoldMessenger.of(context).showSnackBar(
        const SnackBar(
          content: Text(
            'WARNING: The device failed writing to its internal storage. This might be a temporary problem. If it continues happening, the internal storage may be broken. If so, please contact abawo support',
          ),
        ),
      );
    }
  }

  String _dfuPhaseText(DfuUpdateState state) {
    return switch (state) {
      DfuUpdateState.idle => 'Preparing recovery update',
      DfuUpdateState.starting => 'Preparing update',
      DfuUpdateState.enteringBootloader => 'Checking bootloader mode',
      DfuUpdateState.connectingBootloader => 'Connecting to bootloader',
      DfuUpdateState.waitingForStatus => 'Waiting for bootloader status',
      DfuUpdateState.erasing => 'Starting destructive bootloader update',
      DfuUpdateState.transferring => 'Transferring firmware image',
      DfuUpdateState.finishing => 'Finalizing bootloader update',
      DfuUpdateState.rebooting => 'Waiting for bootloader reset',
      DfuUpdateState.verifying => 'Verifying updated app',
      DfuUpdateState.completed => 'Update completed',
      DfuUpdateState.aborted => 'Update canceled',
      DfuUpdateState.failed => 'Update failed',
    };
  }

  String _formatBytes(int bytes) {
    if (bytes < 1024) {
      return '$bytes B';
    }
    if (bytes < 1024 * 1024) {
      return '${(bytes / 1024).toStringAsFixed(1)} KB';
    }
    return '${(bytes / (1024 * 1024)).toStringAsFixed(2)} MB';
  }

  @override
  Widget build(BuildContext context) {
    return FirmwareUpdateFullscreen(
      progress: _dfuProgress,
      selectedFirmware: widget.args.firmware,
      phaseText: _dfuPhaseText(_dfuProgress.state),
      statusText: _firmwareUserMessage,
      formattedProgressBytes:
          '${_formatBytes(_dfuProgress.sentBytes)} / ${_formatBytes(_dfuProgress.totalBytes)}',
      expectedOffsetHex:
          '0x${_dfuProgress.expectedOffset.toRadixString(16).padLeft(8, '0').toUpperCase()}',
      doneLabel: 'Back to devices',
      failedLabel: 'Back to devices',
      onDismiss: () => unawaited(_dismissToDevices()),
    );
  }
}