peyawallet/lib/state/wallet_controller.dart

import 'dart:io';

import 'package:flutter/scheduler.dart';
import 'package:flutter_riverpod/flutter_riverpod.dart';
import 'package:logger/logger.dart';
import 'package:path/path.dart' as p;

import '../data/wallet_repository.dart';
import '../domain/models.dart';
import '../domain/send.dart';
import '../domain/staking.dart';
import '../domain/transactions.dart';
import '../services/wallet_cache_recovery.dart';

class WalletState {
  static const Object _sentinel = Object();

  const WalletState({
    required this.walletInfo,
    required this.balanceAtomic,
    required this.unlockedAtomic,
    required this.transactions,
    required this.subaddresses,
    required this.syncStatus,
    required this.isLoading,
    required this.isBlockingOperation,
    required this.operationMessage,
    required this.isSyncing,
    required this.error,
  });

  final WalletInfo? walletInfo;
  final int balanceAtomic;
  final int unlockedAtomic;
  final List<WalletTransaction> transactions;
  final List<SubaddressInfo> subaddresses;
  final SyncStatus syncStatus;
  final bool isLoading;
  final bool isBlockingOperation;
  final String? operationMessage;
  final bool isSyncing;
  final String? error;

  WalletState copyWith({
    WalletInfo? walletInfo,
    int? balanceAtomic,
    int? unlockedAtomic,
    List<WalletTransaction>? transactions,
    List<SubaddressInfo>? subaddresses,
    SyncStatus? syncStatus,
    bool? isLoading,
    bool? isBlockingOperation,
    Object? operationMessage = _sentinel,
    bool? isSyncing,
    String? error,
  }) {
    return WalletState(
      walletInfo: walletInfo ?? this.walletInfo,
      balanceAtomic: balanceAtomic ?? this.balanceAtomic,
      unlockedAtomic: unlockedAtomic ?? this.unlockedAtomic,
      transactions: transactions ?? this.transactions,
      subaddresses: subaddresses ?? this.subaddresses,
      syncStatus: syncStatus ?? this.syncStatus,
      isLoading: isLoading ?? this.isLoading,
      isBlockingOperation: isBlockingOperation ?? this.isBlockingOperation,
      operationMessage: identical(operationMessage, _sentinel)
          ? this.operationMessage
          : operationMessage as String?,
      isSyncing: isSyncing ?? this.isSyncing,
      error: error,
    );
  }

  factory WalletState.initial() {
    return WalletState(
      walletInfo: null,
      balanceAtomic: 0,
      unlockedAtomic: 0,
      transactions: const [],
      subaddresses: const [],
      syncStatus: SyncStatus.initial(),
      isLoading: false,
      isBlockingOperation: false,
      operationMessage: null,
      isSyncing: false,
      error: null,
    );
  }
}

class WalletController extends StateNotifier<WalletState> {
  WalletController({
    required WalletRepository repository,
    required Logger logger,
    required AppConfig Function() readConfig,
  })  : _repository = repository,
        _logger = logger,
        _readConfig = readConfig,
        _cacheRecovery = WalletCacheRecovery(logger: logger),
        super(WalletState.initial());

  final WalletRepository _repository;
  final Logger _logger;
  final AppConfig Function() _readConfig;
  final WalletCacheRecovery _cacheRecovery;

  Future<bool> tryAutoOpen(AppConfig config) async {
    final wallet = config.lastWallet;
    if (wallet == null) {
      return false;
    }
    final resolvedPath = _resolveOpenPath(wallet.path);
    if (_isDirectoryPath(resolvedPath)) {
      state = state.copyWith(error: 'Wallet file not found in ${wallet.path}');
      return false;
    }
    if (!File(resolvedPath).existsSync() &&
        !File('$resolvedPath.keys').existsSync()) {
      state = state.copyWith(error: 'Wallet not found at ${wallet.path}');
      return false;
    }
    if (!await verifyWalletPassword(path: resolvedPath, password: '')) {
      return true;
    }
    await _cacheRecovery.recoverIfNeeded(resolvedPath);
    await openWallet(path: resolvedPath, password: '');
    return false;
  }

  Future<void> createWallet({
    required String name,
    required String path,
    required String password,
  }) async {
    await _runWithLoading(() async {
      _updateOperationStatus('Preparing wallet directory...');
      await Directory(path).create(recursive: true);
      final resolvedPath = _resolveWalletFilePath(path, name: name);
      _updateOperationStatus('Creating wallet...');
      await _cacheRecovery.markOpenAttempt(resolvedPath);
      await _repository.createWallet(
          name: name, path: path, password: password);
      try {
        _updateOperationStatus('Loading wallet data...');
        await _postOpen(name: name, path: resolvedPath);
      } finally {
        await _cacheRecovery.clearOpenAttempt();
      }
    }, message: 'Creating wallet...');
  }

  Future<String> prepareWalletSeedDraft({
    required String name,
    required String password,
  }) async {
    return _runWithLoading(() async {
      _updateOperationStatus('Preparing temporary wallet...');
      final tempDir = await Directory.systemTemp.createTemp('peyawallet-draft-');
      try {
        _updateOperationStatus('Generating seed phrase...');
        await _repository.createWallet(
          name: name,
          path: tempDir.path,
          password: password,
        );
        final seed = await _repository.getSeed();
        await _repository.closeWallet();
        return seed;
      } finally {
        try {
          await _repository.closeWallet();
        } catch (_) {}
        if (tempDir.existsSync()) {
          await tempDir.delete(recursive: true);
        }
      }
    }, message: 'Preparing seed backup...');
  }

  Future<void> restoreWalletFromSeed({
    required String name,
    required String seed,
    required String path,
    required String password,
  }) async {
    await _runWithLoading(() async {
      _updateOperationStatus('Preparing wallet directory...');
      await Directory(path).create(recursive: true);
      final resolvedPath = _resolveWalletFilePath(path, name: name);
      _updateOperationStatus('Restoring wallet from seed...');
      await _cacheRecovery.markOpenAttempt(resolvedPath);
      await _repository.restoreWalletFromSeed(
        name: name,
        seed: seed,
        path: path,
        password: password,
      );
      try {
        _updateOperationStatus('Loading wallet data...');
        await _postOpen(name: name, path: resolvedPath);
      } finally {
        await _cacheRecovery.clearOpenAttempt();
      }
    }, message: 'Restoring wallet...');
  }

  Future<void> openWallet({
    required String path,
    required String password,
  }) async {
    await _runWithLoading(() async {
      final resolvedPath = _resolveOpenPath(path);
      _updateOperationStatus('Checking wallet files...');
      await _cacheRecovery.recoverIfNeeded(resolvedPath);
      _updateOperationStatus('Opening wallet...');
      await _cacheRecovery.markOpenAttempt(resolvedPath);
      try {
        await _repository.openWallet(path: resolvedPath, password: password);
        _updateOperationStatus('Loading wallet data...');
        await _postOpen(name: p.basename(resolvedPath), path: resolvedPath);
      } finally {
        await _cacheRecovery.clearOpenAttempt();
      }
    }, message: 'Opening wallet...');
  }

  Future<bool> verifyWalletPassword({
    required String path,
    required String password,
  }) async {
    final resolvedPath = _resolveOpenPath(path);
    return _repository.verifyWalletPassword(
      path: resolvedPath,
      password: password,
    );
  }

  Future<bool> walletRequiresPassword(String path) async {
    return !await verifyWalletPassword(path: path, password: '');
  }

  Future<void> connectNode(NodeConfig nodeConfig) async {
    await _runWithLoading(() async {
      await _repository.connectNode(nodeConfig);
    }, blocking: false, message: 'Connecting to node...');
  }

  Future<void> switchWallet() async {
    await _repository.closeWallet();
    await _cacheRecovery.clearOpenAttempt();
    state = WalletState.initial();
  }

  Future<String> getSeed() async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    return _repository.getSeed();
  }

  Future<String> createSubaddress({String? label}) async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    return _runWithLoading(
      () => _repository.createSubaddress(label: label),
      blocking: false,
      message: 'Creating subaddress...',
    );
  }

  Future<void> setSubaddressLabel({
    required int accountIndex,
    required int addressIndex,
    required String label,
  }) async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    await _runWithLoading(() {
      return _repository.setSubaddressLabel(
        accountIndex: accountIndex,
        addressIndex: addressIndex,
        label: label,
      );
    }, blocking: false, message: 'Updating subaddress...');
    await refreshSubaddresses();
  }

  Future<void> refreshSubaddresses() async {
    if (state.walletInfo == null) {
      return;
    }
    final subaddresses = await _repository.getSubaddresses();
    state = state.copyWith(subaddresses: subaddresses);
  }

  Future<SendPreview> prepareSend(SendRequest request) async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    return _runWithLoading(
      () => _repository.prepareSend(request),
      blocking: false,
      message: 'Preparing transaction...',
    );
  }

  Future<SendResult> commitPreparedSend() async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    return _runWithLoading(() async {
      final result = await _repository.commitPreparedSend();
      await _refreshSnapshot();
      return result;
    }, blocking: false, message: 'Submitting transaction...');
  }

  Future<void> discardPreparedSend() async {
    await _repository.discardPreparedSend();
  }

  Future<StakePreview> prepareStake({
    required int amountAtomic,
    int accountIndex = 0,
  }) async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    return _runWithLoading(() {
      return _repository.prepareStake(
        amountAtomic: amountAtomic,
        accountIndex: accountIndex,
      );
    }, blocking: false, message: 'Preparing stake...');
  }

  Future<Map<String, StakeYieldInfo>> getActiveStakeYields({
    required List<WalletTransaction> activeStakes,
    required int currentHeight,
  }) async {
    if (state.walletInfo == null || activeStakes.isEmpty) {
      return const {};
    }
    return _repository.getActiveStakeYields(
      activeStakes: activeStakes,
      currentHeight: currentHeight,
    );
  }

  Future<StakeResult> commitPreparedStake() async {
    if (state.walletInfo == null) {
      throw StateError('No wallet loaded');
    }
    return _runWithLoading(() async {
      final result = await _repository.commitPreparedStake();
      await _refreshSnapshot();
      return result;
    }, blocking: false, message: 'Submitting stake...');
  }

  Future<void> discardPreparedStake() async {
    await _repository.discardPreparedStake();
  }

  Future<void> syncNow({bool announceStages = false}) async {
    if (state.isSyncing || state.walletInfo == null) {
      return;
    }
    state = state.copyWith(isSyncing: true);
    try {
      await _ensureConnected();
      await _refreshSnapshot(announceStages: announceStages);
    } catch (error, stack) {
      _logger.e('Sync failed', error: error, stackTrace: stack);
      state = state.copyWith(error: error.toString());
    } finally {
      if (announceStages) {
        _updateOperationStatus(null);
      }
      state = state.copyWith(isSyncing: false);
    }
  }

  Future<void> rescanBlockchain({int? fromHeight}) async {
    if (state.walletInfo == null) {
      return;
    }
    state = state.copyWith(isSyncing: true, error: null);
    try {
      await _ensureConnected();
      await _repository.rescanBlockchain(fromHeight: fromHeight);
      await _refreshSnapshot();
    } catch (error, stack) {
      _logger.e('Rescan failed', error: error, stackTrace: stack);
      state = state.copyWith(error: error.toString());
    } finally {
      state = state.copyWith(isSyncing: false);
    }
  }

  Future<void> _ensureConnected() async {
    if (await _repository.isConnected()) {
      return;
    }
    final config = _readConfig();
    await _repository.connectNode(config.nodeConfig);
    if (!await _repository.isConnected()) {
      throw StateError('Failed to connect to node');
    }
  }

  Future<void> _postOpen({required String name, required String path}) async {
    final address = await _repository.getAddress();
    final subaddresses = await _repository.getSubaddresses();
    state = state.copyWith(
      walletInfo: WalletInfo(name: name, path: path, address: address),
      balanceAtomic: 0,
      syncStatus: SyncStatus.initial(),
      transactions: const [],
      subaddresses: subaddresses,
      operationMessage: null,
      error: null,
    );
  }

  Future<WalletSnapshot> _refreshSnapshot({bool announceStages = false}) async {
    if (announceStages) {
      _updateOperationStatus('Refreshing balances...');
    }
    final snapshot = await _repository.refresh();
    if (announceStages) {
      _updateOperationStatus('Loading transactions...');
    }
    final transactions = await _repository.getTransactions();
    final syncStatus = SyncStatus(
      synced: snapshot.synced,
      progress: snapshot.progress,
      nodeHeight: snapshot.nodeHeight,
      walletHeight: snapshot.walletHeight,
      lastSync: DateTime.now(),
    );
    final updatedWalletInfo =
        state.walletInfo?.copyWith(address: snapshot.address);
    state = state.copyWith(
      walletInfo: updatedWalletInfo ?? state.walletInfo,
      balanceAtomic: snapshot.balanceAtomic,
      unlockedAtomic: snapshot.unlockedAtomic,
      transactions: transactions,
      subaddresses: snapshot.subaddresses.isEmpty
          ? state.subaddresses
          : snapshot.subaddresses,
      syncStatus: syncStatus,
      error: null,
    );
    return snapshot;
  }

  Future<T> _runWithLoading<T>(
    Future<T> Function() action, {
    bool blocking = true,
    String? message,
  }) async {
    state = state.copyWith(
      isLoading: true,
      isBlockingOperation: blocking,
      operationMessage: message,
      error: null,
    );
    try {
      await SchedulerBinding.instance.endOfFrame;
      return await action();
    } catch (error, stack) {
      _logger.e('Wallet action failed', error: error, stackTrace: stack);
      state = state.copyWith(error: error.toString());
      rethrow;
    } finally {
      state = state.copyWith(
        isLoading: false,
        isBlockingOperation: false,
        operationMessage: null,
      );
    }
  }

  void _updateOperationStatus(String? message) {
    state = state.copyWith(
      operationMessage: message,
      isBlockingOperation: false,
    );
  }

  String _resolveWalletFilePath(String basePath, {required String name}) {
    if (basePath.isEmpty) {
      return basePath;
    }
    if (_isDirectoryPath(basePath)) {
      return p.join(basePath, name);
    }
    return basePath;
  }

  String _resolveOpenPath(String basePath) {
    if (!_isDirectoryPath(basePath)) {
      return basePath;
    }
    final dirName = p.basename(basePath);
    final candidate = p.join(basePath, dirName);
    if (File(candidate).existsSync() || File('$candidate.keys').existsSync()) {
      return candidate;
    }
    final dir = Directory(basePath);
    if (!dir.existsSync()) {
      return basePath;
    }
    final entries = dir
        .listSync()
        .whereType<File>()
        .map((entry) => entry.path)
        .where((entryPath) => !entryPath.endsWith('.keys'))
        .toList()
      ..sort();
    for (final entryPath in entries) {
      if (File('$entryPath.keys').existsSync()) {
        return entryPath;
      }
    }
    return basePath;
  }

  bool _isDirectoryPath(String path) {
    if (path.endsWith(p.separator)) {
      return true;
    }
    return Directory(path).existsSync();
  }
}