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

import 'dart:convert';
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/note_repository.dart';
import 'package:trainhub_flutter/data/services/ai_process_manager.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
class ChatController extends _$ChatController {
  late ChatRepository _repo;
  late NoteRepository _noteRepo;

  final _dio = Dio(
    BaseOptions(
      connectTimeout: AiConstants.serverConnectTimeout,
      receiveTimeout: AiConstants.serverReceiveTimeout,
    ),
  );

  @override
  Future<ChatState> build() async {
    _repo = getIt<ChatRepository>();
    _noteRepo = getIt<NoteRepository>();
    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 systemPrompt = _buildSystemPrompt(contextChunks);
    final history = _buildHistory();
    final fullAiResponse = await _streamResponse(systemPrompt, history);
    await _persistAssistantResponse(sessionId, content, fullAiResponse);
  }

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

  List<Map<String, String>> _buildHistory() {
    final messages = state.valueOrNull?.messages ?? [];
    return messages
        .map((m) => <String, String>{
              'role': m.isUser ? 'user' : 'assistant',
              'content': m.content,
            })
        .toList();
  }

  Future<String> _streamResponse(
    String systemPrompt,
    List<Map<String, String>> history,
  ) async {
    final generateStep = _createStep('Generating response...');
    String fullAiResponse = '';
    try {
      final response = await _dio.post<ResponseBody>(
        AiConstants.chatApiUrl,
        options: Options(responseType: ResponseType.stream),
        data: {
          'messages': [
            {'role': 'system', 'content': systemPrompt},
            ...history,
          ],
          'temperature': AiConstants.chatTemperature,
          'stream': true,
        },
      );
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.running,
        title: 'Writing...',
      );
      final stream = response.data!.stream;
      await for (final chunk in stream) {
        final textChunk = utf8.decode(chunk);
        for (final line in textChunk.split('\n')) {
          if (!line.startsWith('data: ')) continue;
          final dataStr = line.substring(6).trim();
          if (dataStr == '[DONE]') break;
          if (dataStr.isEmpty) continue;
          try {
            final data = jsonDecode(dataStr);
            final delta = data['choices']?[0]?['delta']?['content'] ?? '';
            if (delta.isNotEmpty) {
              fullAiResponse += delta;
              final updatedState = state.valueOrNull;
              if (updatedState != null) {
                state = AsyncValue.data(
                  updatedState.copyWith(streamingContent: fullAiResponse),
                );
              }
            }
          } catch (_) {}
        }
      }
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.completed,
        title: 'Response generated',
      );
    } on DioException catch (e) {
      fullAiResponse += '\n\n[AI model communication error]';
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.error,
        title: 'Generation failed',
        details: '${e.message}',
      );
    } catch (e) {
      fullAiResponse += '\n\n[Unexpected error]';
      _updateStep(
        generateStep.id,
        status: ThinkingStepStatus.error,
        title: 'Generation failed',
        details: e.toString(),
      );
    }
    return fullAiResponse;
  }

  Future<void> _persistAssistantResponse(
    String sessionId,
    String userContent,
    String aiResponse,
  ) async {
    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) {
    if (chunks.isEmpty) return AiConstants.baseSystemPrompt;
    final contextBlock = chunks
        .asMap()
        .entries
        .map((e) => '[${e.key + 1}] ${e.value}')
        .join('\n\n');
    return '${AiConstants.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.';
  }
}