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

import 'package:dio/dio.dart';
import 'package:riverpod_annotation/riverpod_annotation.dart';
import 'package:trainhub_flutter/domain/repositories/chat_repository.dart';
import 'package:trainhub_flutter/domain/repositories/note_repository.dart';
import 'package:trainhub_flutter/injection.dart';
import 'package:trainhub_flutter/presentation/chat/chat_state.dart';

part 'chat_controller.g.dart';

const _chatApiUrl = 'http://localhost:8080/v1/chat/completions';

/// Base system prompt that is always included.
const _baseSystemPrompt =
    'You are a helpful AI fitness assistant for personal trainers. '
    'Help users design training plans, analyse exercise technique, '
    'and answer questions about sports science and nutrition.';

@riverpod
class ChatController extends _$ChatController {
  late ChatRepository _repo;
  late NoteRepository _noteRepo;

  // Shared Dio client — generous timeout for 7B models running on CPU.
  final _dio = Dio(
    BaseOptions(
      connectTimeout: const Duration(seconds: 30),
      receiveTimeout: const Duration(minutes: 5),
    ),
  );

  @override
  Future<ChatState> build() async {
    _repo = getIt<ChatRepository>();
    _noteRepo = getIt<NoteRepository>();
    final sessions = await _repo.getAllSessions();
    return ChatState(sessions: sessions);
  }

  // -------------------------------------------------------------------------
  // Session management (unchanged)
  // -------------------------------------------------------------------------

  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,
      ),
    );
  }

  // -------------------------------------------------------------------------
  // Send message  (RAG + Step D)
  // -------------------------------------------------------------------------

  Future<void> sendMessage(String content) async {
    final current = state.valueOrNull;
    if (current == null) return;

    // ── 1. Resolve / create a session ─────────────────────────────────────
    String sessionId;
    if (current.activeSession == null) {
      final session = await _repo.createSession();
      sessionId = session.id;
      final sessions = await _repo.getAllSessions();
      state = AsyncValue.data(
        current.copyWith(sessions: sessions, activeSession: session),
      );
    } else {
      sessionId = current.activeSession!.id;
    }

    // ── 2. Persist user message & show typing indicator ───────────────────
    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,
      ),
    );

    // ── 3. RAG: retrieve relevant chunks from the knowledge base ──────────
    // Gracefully degrades — if Nomic server is unavailable or no chunks
    // exist, the chat still works with the base system prompt alone.
    List<String> contextChunks = [];
    try {
      contextChunks = await _noteRepo.searchSimilar(content, topK: 3);
    } catch (_) {
      // Nomic server not running or no chunks stored — continue without RAG.
    }

    // ── 4. Build enriched system prompt (Step D) ──────────────────────────
    final systemPrompt = _buildSystemPrompt(contextChunks);

    // Build the full conversation history so the model maintains context.
    final history = messagesAfterUser
        .map(
          (m) => <String, String>{
            'role': m.isUser ? 'user' : 'assistant',
            'content': m.content,
          },
        )
        .toList();

    // ── 5. POST to Qwen (http://localhost:8080/v1/chat/completions) ────────
    String aiResponse;
    try {
      final response = await _dio.post<Map<String, dynamic>>(
        _chatApiUrl,
        data: {
          'messages': [
            {'role': 'system', 'content': systemPrompt},
            ...history,
          ],
          'temperature': 0.7,
        },
      );
      aiResponse =
          response.data!['choices'][0]['message']['content'] as String;
    } on DioException catch (e) {
      aiResponse =
          'Could not reach the AI server (${e.message}). '
          'Make sure AI models are downloaded and the inference servers have '
          'had time to start.';
    } catch (e) {
      aiResponse = 'An unexpected error occurred: $e';
    }

    // ── 6. Persist response & update session title on first exchange ───────
    await _repo.addMessage(
      sessionId: sessionId,
      role: 'assistant',
      content: aiResponse,
    );

    final messagesAfterAi = await _repo.getMessages(sessionId);
    if (messagesAfterAi.length <= 2) {
      final title =
          content.length > 30 ? '${content.substring(0, 30)}…' : content;
      await _repo.updateSessionTitle(sessionId, title);
    }

    final sessions = await _repo.getAllSessions();
    state = AsyncValue.data(
      state.valueOrNull!.copyWith(
        messages: messagesAfterAi,
        isTyping: false,
        sessions: sessions,
      ),
    );
  }

  // -------------------------------------------------------------------------
  // Helpers
  // -------------------------------------------------------------------------

  /// Builds the system prompt, injecting RAG context when available.
  static String _buildSystemPrompt(List<String> chunks) {
    if (chunks.isEmpty) return _baseSystemPrompt;

    final contextBlock = chunks
        .asMap()
        .entries
        .map((e) => '[${e.key + 1}] ${e.value}')
        .join('\n\n');

    return '$_baseSystemPrompt\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.';
  }
}