trainhub-flutter/lib/presentation/settings/ai_model_settings_controller.dart

import 'dart:async';
import 'dart:io';

import 'package:archive/archive_io.dart';
import 'package:dio/dio.dart';
import 'package:path/path.dart' as p;
import 'package:path_provider/path_provider.dart';
import 'package:riverpod_annotation/riverpod_annotation.dart';
import 'package:trainhub_flutter/core/constants/ai_constants.dart';
import 'package:trainhub_flutter/data/services/ai_process_manager.dart';
import 'package:trainhub_flutter/injection.dart' as di;
import 'package:trainhub_flutter/presentation/settings/ai_model_settings_state.dart';

part 'ai_model_settings_controller.g.dart';

Future<String> _llamaArchiveUrl() async {
  final build = AiConstants.llamaBuild;
  if (Platform.isMacOS) {
    final result = await Process.run('uname', ['-m']);
    final arch = (result.stdout as String).trim();
    final suffix = arch == 'arm64' ? 'macos-arm64' : 'macos-x64';
    return 'https://github.com/ggml-org/llama.cpp/releases/download/$build/llama-$build-bin-$suffix.tar.gz';
  } else if (Platform.isWindows) {
    return 'https://github.com/ggml-org/llama.cpp/releases/download/$build/llama-$build-bin-win-vulkan-x64.zip';
  } else if (Platform.isLinux) {
    return 'https://github.com/ggml-org/llama.cpp/releases/download/$build/llama-$build-bin-ubuntu-vulkan-x64.tar.gz';
  }
  throw UnsupportedError('Unsupported platform: ${Platform.operatingSystem}');
}

// keepAlive: a download must survive navigating away from the settings /
// welcome screen — otherwise leaving the page silently drops its progress.
@Riverpod(keepAlive: true)
class AiModelSettingsController extends _$AiModelSettingsController {
  final _dio = Dio();

  @override
  AiModelSettingsState build() {
    // Check installed files right away so every screen (settings, welcome,
    // chat) sees the real state without a manual "Validate" click.
    Future.microtask(validateModels);
    return const AiModelSettingsState();
  }

  /// A file that exists but is suspiciously small is a leftover from an
  /// interrupted download — treat it as missing.
  static bool _isComplete(File file, int minBytes) =>
      file.existsSync() && file.lengthSync() >= minBytes;

  Future<void> validateModels() async {
    state = state.copyWith(
      currentTask: 'Checking installed files…',
      errorMessage: null,
    );
    try {
      final dir = await getApplicationDocumentsDirectory();
      final base = dir.path;
      final validated =
          _isComplete(
            File(p.join(base, AiConstants.serverBinaryName)),
            AiConstants.serverBinaryMinBytes,
          ) &&
          _isComplete(
            File(p.join(base, AiConstants.nomicModelFile)),
            AiConstants.nomicModelMinBytes,
          ) &&
          _isComplete(
            File(p.join(base, AiConstants.qwenModelFile)),
            AiConstants.qwenModelMinBytes,
          );
      state = state.copyWith(
        areModelsValidated: validated,
        currentTask: validated ? 'All files present.' : 'Files missing.',
      );
    } catch (e) {
      state = state.copyWith(
        areModelsValidated: false,
        currentTask: 'Validation failed.',
        errorMessage: e.toString(),
      );
    }
  }

  Future<void> downloadAll() async {
    if (state.isDownloading) return;
    try {
      await di.getIt<AiProcessManager>().stopServers();
    } catch (_) {}
    state = state.copyWith(
      isDownloading: true,
      progress: 0.0,
      areModelsValidated: false,
      errorMessage: null,
    );
    try {
      final dir = await getApplicationDocumentsDirectory();
      final archiveUrl = await _llamaArchiveUrl();
      final archiveExt = archiveUrl.endsWith('.zip') ? '.zip' : '.tar.gz';
      final archivePath = p.join(dir.path, 'llama_binary$archiveExt');
      await _downloadFile(
        url: archiveUrl,
        savePath: archivePath,
        taskLabel: 'Downloading llama.cpp binary…',
        overallStart: 0.0,
        overallEnd: 0.2,
      );
      state = state.copyWith(
        currentTask: 'Extracting llama.cpp binary…',
        progress: 0.2,
      );
      await _extractBinary(archivePath, dir.path);
      final archiveFile = File(archivePath);
      if (archiveFile.existsSync()) archiveFile.deleteSync();
      await _downloadFile(
        url: AiConstants.nomicModelUrl,
        savePath: p.join(dir.path, AiConstants.nomicModelFile),
        taskLabel: 'Downloading Nomic embedding model…',
        overallStart: 0.2,
        overallEnd: 0.35,
        minBytes: AiConstants.nomicModelMinBytes,
      );
      await _downloadFile(
        url: AiConstants.qwenModelUrl,
        savePath: p.join(dir.path, AiConstants.qwenModelFile),
        taskLabel: 'Downloading Qwen3 4B model…',
        overallStart: 0.35,
        overallEnd: 1.0,
        minBytes: AiConstants.qwenModelMinBytes,
      );
      state = state.copyWith(
        isDownloading: false,
        progress: 1.0,
        currentTask: 'Download complete.',
      );
      await validateModels();
      // Bring the AI online immediately — the user shouldn't have to visit
      // the chat page or click anything for the servers to start loading.
      if (state.areModelsValidated) {
        unawaited(di.getIt<AiProcessManager>().startServers());
      }
    } on DioException catch (e) {
      state = state.copyWith(
        isDownloading: false,
        currentTask: 'Download failed.',
        errorMessage: 'Network error: ${e.message}',
      );
    } catch (e) {
      state = state.copyWith(
        isDownloading: false,
        currentTask: 'Download failed.',
        errorMessage: e.toString(),
      );
    }
  }

  Future<void> _downloadFile({
    required String url,
    required String savePath,
    required String taskLabel,
    required double overallStart,
    required double overallEnd,
    int minBytes = 0,
  }) async {
    // Already downloaded (e.g. a retry after a failure further down the
    // list) — don't pull gigabytes again.
    if (_isComplete(File(savePath), minBytes) && minBytes > 0) {
      state = state.copyWith(progress: overallEnd);
      return;
    }

    state = state.copyWith(currentTask: taskLabel, progress: overallStart);

    // Download to a temp file and rename at the end, so an interrupted
    // download can never masquerade as a complete model file.
    final partPath = '$savePath.part';
    await _dio.download(
      url,
      partPath,
      onReceiveProgress: (received, total) {
        if (total <= 0) return;
        final fileProgress = received / total;
        final overall =
            overallStart + fileProgress * (overallEnd - overallStart);
        state = state.copyWith(progress: overall);
      },
      options: Options(
        followRedirects: true,
        maxRedirects: 5,
        receiveTimeout: const Duration(hours: 2),
      ),
    );

    final target = File(savePath);
    if (target.existsSync()) target.deleteSync();
    File(partPath).renameSync(savePath);
  }

  Future<void> _extractBinary(String archivePath, String destDir) async {
    final extractDir = p.join(destDir, '_llama_extract_tmp');
    final extractDirObj = Directory(extractDir);
    if (extractDirObj.existsSync()) extractDirObj.deleteSync(recursive: true);
    extractDirObj.createSync(recursive: true);
    try {
      await extractFileToDisk(archivePath, extractDir);
      bool foundServer = false;
      final binaryName = AiConstants.serverBinaryName;
      for (final entity in extractDirObj.listSync(recursive: true)) {
        if (entity is File) {
          final ext = p.extension(entity.path).toLowerCase();
          final name = p.basename(entity.path);
          if (name == binaryName ||
              ext == '.dll' ||
              ext == '.so' ||
              ext == '.dylib') {
            final destFile = p.join(destDir, name);
            int retryCount = 0;
            bool success = false;
            while (!success && retryCount < 5) {
              try {
                if (File(destFile).existsSync()) {
                  File(destFile).deleteSync();
                }
                entity.copySync(destFile);
                success = true;
              } on FileSystemException catch (_) {
                if (retryCount >= 4) {
                  throw Exception(
                    'Failed to overwrite $name. Ensure no other applications are using it.',
                  );
                }
                await Future.delayed(const Duration(milliseconds: 500));
                retryCount++;
              }
            }
            if (name == binaryName) {
              foundServer = true;
              if (Platform.isMacOS || Platform.isLinux) {
                await Process.run('chmod', ['+x', destFile]);
              }
            }
          }
        }
      }
      if (!foundServer) {
        throw FileSystemException(
          'llama-server binary not found in archive.',
          archivePath,
        );
      }
    } finally {
      if (extractDirObj.existsSync()) {
        extractDirObj.deleteSync(recursive: true);
      }
    }
  }
}