trainhub-flutter/lib/presentation/chat/chat_controller.dart

import 'package:dio/dio.dart';
import 'package:riverpod_annotation/riverpod_annotation.dart';
import 'package:trainhub_flutter/core/constants/ai_constants.dart';
import 'package:trainhub_flutter/domain/repositories/chat_repository.dart';
import 'package:trainhub_flutter/domain/repositories/exercise_repository.dart';
import 'package:trainhub_flutter/domain/repositories/note_repository.dart';
import 'package:trainhub_flutter/domain/repositories/training_plan_repository.dart';
import 'package:trainhub_flutter/data/services/ai_process_manager.dart';
import 'package:trainhub_flutter/data/services/ai_settings_service.dart';
import 'package:trainhub_flutter/data/services/llm_client.dart';
import 'package:trainhub_flutter/injection.dart';
import 'package:trainhub_flutter/presentation/chat/chat_state.dart';
import 'package:uuid/uuid.dart';

part 'chat_controller.g.dart';

@riverpod
AiProcessManager aiProcessManager(AiProcessManagerRef ref) {
  final manager = getIt<AiProcessManager>();
  manager.addListener(() => ref.notifyListeners());
  return manager;
}

@riverpod
AiSettingsService aiSettingsService(AiSettingsServiceRef ref) {
  final service = getIt<AiSettingsService>();
  service.addListener(() => ref.notifyListeners());
  return service;
}

@riverpod
class ChatController extends _$ChatController {
  late ChatRepository _repo;
  late NoteRepository _noteRepo;
  late LlmClient _llm;
  CancelToken? _cancelToken;

  @override
  Future<ChatState> build() async {
    _repo = getIt<ChatRepository>();
    _noteRepo = getIt<NoteRepository>();
    _llm = getIt<LlmClient>();
    // Abort any in-flight generation when the user leaves the chat page.
    ref.onDispose(() => _cancelToken?.cancel());
    final aiManager = ref.read(aiProcessManagerProvider);
    if (aiManager.status == AiServerStatus.offline) {
      aiManager.startServers();
    }
    final sessions = await _repo.getAllSessions();
    return ChatState(sessions: sessions);
  }

  Future<void> createSession() async {
    final session = await _repo.createSession();
    final sessions = await _repo.getAllSessions();
    state = AsyncValue.data(
      ChatState(sessions: sessions, activeSession: session),
    );
  }

  Future<void> loadSession(String id) async {
    final session = await _repo.getSession(id);
    if (session == null) return;
    final messages = await _repo.getMessages(id);
    final current = state.valueOrNull ?? const ChatState();
    state = AsyncValue.data(
      current.copyWith(activeSession: session, messages: messages),
    );
  }

  Future<void> deleteSession(String id) async {
    await _repo.deleteSession(id);
    final sessions = await _repo.getAllSessions();
    final current = state.valueOrNull ?? const ChatState();
    state = AsyncValue.data(
      current.copyWith(
        sessions: sessions,
        activeSession: current.activeSession?.id == id
            ? null
            : current.activeSession,
        messages: current.activeSession?.id == id ? [] : current.messages,
      ),
    );
  }

  Future<void> sendMessage(String content) async {
    final current = state.valueOrNull;
    if (current == null) return;
    final sessionId = await _resolveSession(current, content);
    await _persistUserMessage(sessionId, content);
    final contextChunks = await _searchKnowledgeBase(content);
    final trainingContext = await _buildTrainingContext();
    final systemPrompt = _buildSystemPrompt(contextChunks, trainingContext);
    final history = _buildHistory();
    final fullAiResponse = await _streamResponse(systemPrompt, history);
    await _persistAssistantResponse(sessionId, content, fullAiResponse);
  }

  /// Summarizes the user's exercise library and training plans so the model
  /// can reference and plan around real data. Only attached when a cloud
  /// provider is active — the local 4B model's context is too small for it.
  Future<String> _buildTrainingContext() async {
    if (getIt<LlmClient>().activeProvider == AiProvider.local) return '';
    try {
      final exercises = await getIt<ExerciseRepository>().getAll();
      final plans = await getIt<TrainingPlanRepository>().getAll();

      final buffer = StringBuffer();
      if (exercises.isNotEmpty) {
        buffer.writeln("### The trainer's exercise library:");
        for (final e in exercises.take(150)) {
          buffer.write('- ${e.name}');
          final tags = e.tags;
          if (tags != null && tags.isNotEmpty) buffer.write(' [$tags]');
          buffer.writeln();
        }
      }
      if (plans.isNotEmpty) {
        buffer.writeln("\n### The trainer's training plans:");
        for (final plan in plans.take(30)) {
          buffer.writeln('- ${plan.name}');
        }
      }
      return buffer.toString();
    } catch (_) {
      return '';
    }
  }

  Future<String> _resolveSession(ChatState current, String content) async {
    if (current.activeSession != null) return current.activeSession!.id;
    final session = await _repo.createSession();
    final sessions = await _repo.getAllSessions();
    state = AsyncValue.data(
      current.copyWith(sessions: sessions, activeSession: session),
    );
    return session.id;
  }

  Future<void> _persistUserMessage(String sessionId, String content) async {
    await _repo.addMessage(
      sessionId: sessionId,
      role: 'user',
      content: content,
    );
    final messagesAfterUser = await _repo.getMessages(sessionId);
    state = AsyncValue.data(
      state.valueOrNull!.copyWith(
        messages: messagesAfterUser,
        isTyping: true,
        thinkingSteps: [],
        streamingContent: '',
      ),
    );
  }

  Future<List<String>> _searchKnowledgeBase(String query) async {
    final searchStep = _createStep('Searching knowledge base...');
    List<String> contextChunks = [];
    try {
      contextChunks = await _noteRepo.searchSimilar(query, topK: 3);
      if (contextChunks.isNotEmpty) {
        _updateStep(
          searchStep.id,
          status: ThinkingStepStatus.completed,
          title: 'Found ${contextChunks.length} documents',
          details: 'Context added for assistant.',
        );
      } else {
        _updateStep(
          searchStep.id,
          status: ThinkingStepStatus.completed,
          title: 'No matching documents in knowledge base',
          details: 'Responding based on general knowledge.',
        );
      }
    } catch (e) {
      _updateStep(
        searchStep.id,
        status: ThinkingStepStatus.error,
        title: 'Knowledge base search error',
        details: e.toString(),
      );
    }
    return contextChunks;
  }

  /// Most recent messages only — an unbounded history would eventually
  /// overflow the model context and slow every request down.
  List<Map<String, String>> _buildHistory() {
    final messages = state.valueOrNull?.messages ?? [];
    final recent = messages.length > AiConstants.chatHistoryLimit
        ? messages.sublist(messages.length - AiConstants.chatHistoryLimit)
        : messages;
    return recent
        .map(
          (m) => <String, String>{
            'role': m.isUser ? 'user' : 'assistant',
            'content': m.content,
          },
        )
        .toList();
  }

  /// Stops an in-flight generation. The partial response streamed so far is
  /// kept and persisted like a normal reply.
  void stopGeneration() => _cancelToken?.cancel();

  Future<String> _streamResponse(
    String systemPrompt,
    List<Map<String, String>> history,
  ) async {
    final generateStep = _createStep('Generating response...');
    String fullAiResponse = '';
    _cancelToken = CancelToken();
    try {
      final stream = _llm.streamChat([
        {'role': 'system', 'content': systemPrompt},
        ...history,
      ], cancelToken: _cancelToken);
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.running,
        title: 'Writing...',
      );
      await for (final delta in stream) {
        fullAiResponse += delta;
        final updatedState = state.valueOrNull;
        if (updatedState != null) {
          state = AsyncValue.data(
            updatedState.copyWith(streamingContent: fullAiResponse),
          );
        }
      }
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.completed,
        title: 'Response generated',
      );
    } on DioException catch (e) {
      if (CancelToken.isCancel(e)) {
        _updateStep(
          generateStep.id,
          status: ThinkingStepStatus.completed,
          title: 'Stopped by user',
        );
      } else {
        _updateStep(
          generateStep.id,
          status: ThinkingStepStatus.error,
          title: 'Generation failed',
          details: e.message ?? e.toString(),
        );
      }
    } catch (e) {
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.error,
        title: 'Generation failed',
        details: e.toString(),
      );
    } finally {
      _cancelToken = null;
    }
    return fullAiResponse;
  }

  Future<void> _persistAssistantResponse(
    String sessionId,
    String userContent,
    String aiResponse,
  ) async {
    // A failed generation yields an empty response — leave the error visible
    // in the thinking steps instead of saving an empty assistant message.
    if (aiResponse.trim().isEmpty) {
      final current = state.valueOrNull;
      if (current != null) {
        state = AsyncValue.data(
          current.copyWith(isTyping: false, streamingContent: null),
        );
      }
      return;
    }
    await _repo.addMessage(
      sessionId: sessionId,
      role: 'assistant',
      content: aiResponse,
    );
    final messagesAfterAi = await _repo.getMessages(sessionId);
    if (messagesAfterAi.length <= 2) {
      final title = userContent.length > 30
          ? '${userContent.substring(0, 30)}…'
          : userContent;
      await _repo.updateSessionTitle(sessionId, title);
    }
    final sessions = await _repo.getAllSessions();
    state = AsyncValue.data(
      state.valueOrNull!.copyWith(
        messages: messagesAfterAi,
        isTyping: false,
        streamingContent: null,
        thinkingSteps: [],
        sessions: sessions,
      ),
    );
  }

  ThinkingStep _createStep(String title) {
    final step = ThinkingStep(
      id: const Uuid().v4(),
      title: title,
      status: ThinkingStepStatus.pending,
    );
    final current = state.valueOrNull;
    if (current != null) {
      state = AsyncValue.data(
        current.copyWith(thinkingSteps: [...current.thinkingSteps, step]),
      );
    }
    return step;
  }

  void _updateStep(
    String id, {
    ThinkingStepStatus? status,
    String? title,
    String? details,
  }) {
    final current = state.valueOrNull;
    if (current == null) return;
    final updatedSteps = current.thinkingSteps.map((s) {
      if (s.id != id) return s;
      return s.copyWith(
        status: status ?? s.status,
        title: title ?? s.title,
        details: details ?? s.details,
      );
    }).toList();
    state = AsyncValue.data(current.copyWith(thinkingSteps: updatedSteps));
  }

  static String _buildSystemPrompt(
    List<String> chunks,
    String trainingContext,
  ) {
    final buffer = StringBuffer(AiConstants.baseSystemPrompt);
    if (trainingContext.isNotEmpty) {
      buffer.write('\n\n$trainingContext');
      buffer.write(
        '\nWhen designing or discussing training plans, prefer exercises '
        'from the library above and reference existing plans by name.',
      );
    }
    if (chunks.isNotEmpty) {
      final contextBlock = chunks
          .asMap()
          .entries
          .map((e) => '[${e.key + 1}] ${e.value}')
          .join('\n\n');
      buffer.write(
        '\n\n### Relevant notes from the trainer\'s knowledge base:\n'
        '$contextBlock\n\n'
        'Use the above context to inform your response when relevant. '
        'If the context is not directly applicable, rely on your general '
        'fitness knowledge.',
      );
    }
    return buffer.toString();
  }
}