FabledApp/lib/providers/chat_provider.dart

import 'dart:async';

import 'package:flutter_riverpod/flutter_riverpod.dart';

import '../data/models/conversation.dart';
import '../data/models/message.dart';
import '../data/repositories/chat_repository.dart';
import 'api_client_provider.dart';

// Separate NotifierProvider.family so UI re-builds immediately when streaming starts/stops.
final isStreamingProvider =
    NotifierProvider.family<_IsStreamingNotifier, bool, int>(
  (convId) => _IsStreamingNotifier(convId),
);

class _IsStreamingNotifier extends Notifier<bool> {
  // ignore: avoid_unused_constructor_parameters
  _IsStreamingNotifier(int convId);

  @override
  bool build() => false;
}

// Tracks the current tool status text during generation (empty = no status).
final streamingStatusProvider =
    NotifierProvider.family<_StreamingStatusNotifier, String, int>(
  (convId) => _StreamingStatusNotifier(convId),
);

class _StreamingStatusNotifier extends Notifier<String> {
  // ignore: avoid_unused_constructor_parameters
  _StreamingStatusNotifier(int convId);

  @override
  String build() => '';
}

final conversationsProvider =
    AsyncNotifierProvider<ConversationsNotifier, List<Conversation>>(
        ConversationsNotifier.new);

class ConversationsNotifier extends AsyncNotifier<List<Conversation>> {
  @override
  Future<List<Conversation>> build() async {
    return ref.watch(chatRepositoryProvider).getConversations();
  }

  Future<Conversation> create(String title) async {
    final conv =
        await ref.read(chatRepositoryProvider).createConversation(title);
    state = AsyncData([conv, ...state.value ?? []]);
    return conv;
  }

  Future<void> delete(int id) async {
    await ref.read(chatRepositoryProvider).deleteConversation(id);
    state = AsyncData([
      for (final c in state.value ?? [])
        if (c.id != id) c,
    ]);
  }

  /// Re-fetch conversations without clearing the current list (no flicker).
  Future<void> refresh() async {
    final fresh = await ref.read(chatRepositoryProvider).getConversations();
    state = AsyncData(fresh);
  }

  // Called after a message is sent to patch the server-generated title
  // in-place without triggering a full reload or loading state.
  void patchConversation(Conversation updated) {
    final list = state.value;
    if (list == null) return;
    state = AsyncData([
      for (final c in list)
        if (c.id == updated.id) updated else c,
    ]);
  }
}

final messagesProvider =
    AsyncNotifierProvider.family<MessagesNotifier, List<Message>, int>(
  (convId) => MessagesNotifier(convId),
);

class MessagesNotifier extends AsyncNotifier<List<Message>> {
  final int _convId;
  MessagesNotifier(this._convId);

  @override
  Future<List<Message>> build() async {
    final (_, messages) =
        await ref.watch(chatRepositoryProvider).getMessages(_convId);
    return messages;
  }

  /// Re-fetch messages without clearing the current list (no flicker).
  ///
  /// Also unfreezes the UI if streaming state got stuck true — this happens
  /// when an SSE connection dies silently (mobile network handoff, app
  /// backgrounded mid-stream, reverse proxy dropping idle sockets) and the
  /// send loop never observes a close. If the server-side message is already
  /// done, we clear `isStreamingProvider` so the input unlocks.
  Future<void> refresh() async {
    final (_, messages) =
        await ref.read(chatRepositoryProvider).getMessages(_convId);
    state = AsyncData(messages);
    Message? lastAssistant;
    for (var i = messages.length - 1; i >= 0; i--) {
      if (messages[i].role == MessageRole.assistant) {
        lastAssistant = messages[i];
        break;
      }
    }
    if (lastAssistant != null && lastAssistant.status != 'generating') {
      ref.read(isStreamingProvider(_convId).notifier).state = false;
      ref.read(streamingStatusProvider(_convId).notifier).state = '';
    }
  }

  /// Attach to an already-running generation for this conversation.
  ///
  /// Used when the chat screen lands on a conversation that was started by
  /// something other than a direct user message — e.g. the /news discuss
  /// button, which creates a conversation on the backend and auto-kicks a
  /// generation before navigating. Without this the stream runs to
  /// completion invisibly and the screen only shows the final persisted
  /// message after a manual refresh.
  ///
  /// Safe to call unconditionally on screen init: no-ops when there is no
  /// generating assistant message. Mirrors the web chat store's
  /// reconnectIfGenerating() helper.
  Future<void> attachToGeneration() async {
    final convId = _convId;
    final repo = ref.read(chatRepositoryProvider);

    // Make sure we're looking at fresh server state before deciding whether
    // to attach. The provider's build() fetches once; if the conversation
    // was seeded via a POST that happened between build() and this call the
    // generating placeholder won't be in our in-memory list yet.
    try {
      final (_, fresh) = await repo.getMessages(convId);
      state = AsyncData(fresh);
    } catch (_) {
      // If we can't load messages we can't attach either — bail cleanly.
      return;
    }

    if (ref.read(isStreamingProvider(convId))) return;
    final msgs = state.value ?? const <Message>[];
    final hasGeneratingAssistant = msgs.any(
      (m) => m.role == MessageRole.assistant && m.status == 'generating',
    );
    if (!hasGeneratingAssistant) return;

    ref.read(isStreamingProvider(convId).notifier).state = true;

    const stallTimeout = Duration(seconds: 45);
    bool streamedContent = false;
    final iter = StreamIterator(repo.streamGeneration(convId));
    try {
      while (await iter.moveNext().timeout(stallTimeout)) {
        final event = iter.current;
        if (event is ChatTextChunk) {
          streamedContent = true;
          ref.read(streamingStatusProvider(convId).notifier).state = '';
          final cur = state.requireValue;
          if (cur.isEmpty) continue;
          // Route text chunks into the generating assistant message. The
          // last message is usually the placeholder, but tool-call fan-in
          // means we can't rely on that universally.
          final idx = _findGeneratingAssistantIndex(cur);
          if (idx < 0) continue;
          final updated = cur[idx].copyWith(content: cur[idx].content + event.text);
          state = AsyncData([
            ...cur.sublist(0, idx),
            updated,
            ...cur.sublist(idx + 1),
          ]);
        } else if (event is ChatStatusUpdate) {
          ref.read(streamingStatusProvider(convId).notifier).state = event.status;
        } else if (event is ChatToolCall) {
          final cur = state.requireValue;
          if (cur.isEmpty) continue;
          final idx = _findGeneratingAssistantIndex(cur);
          if (idx < 0) continue;
          final nextCalls = [...?cur[idx].toolCalls, event.toolCall];
          final updated = cur[idx].copyWith(toolCalls: nextCalls);
          state = AsyncData([
            ...cur.sublist(0, idx),
            updated,
            ...cur.sublist(idx + 1),
          ]);
        }
      }
    } on TimeoutException {
      // Stall — fall through to polling.
    } catch (_) {
      // Stream failed — fall through to polling.
    } finally {
      await iter.cancel();
    }

    try {
      for (var attempt = 0; attempt < 20; attempt++) {
        if (attempt > 0) await Future.delayed(const Duration(seconds: 2));
        final (_, fresh) = await repo.getMessages(convId);
        final done = fresh.any(
          (m) => m.role == MessageRole.assistant && m.status != 'generating',
        );
        final polledHasContent = fresh.any(
          (m) => m.role == MessageRole.assistant && m.content.isNotEmpty,
        );
        if (!streamedContent || done || polledHasContent) {
          state = AsyncData(fresh);
        }
        if (done) break;
      }
    } catch (_) {
      // Give up silently — user can pull-to-refresh.
    } finally {
      ref.read(isStreamingProvider(convId).notifier).state = false;
      ref.read(streamingStatusProvider(convId).notifier).state = '';
    }
  }

  int _findGeneratingAssistantIndex(List<Message> msgs) {
    for (var i = msgs.length - 1; i >= 0; i--) {
      final m = msgs[i];
      if (m.role == MessageRole.assistant && m.status == 'generating') return i;
    }
    return -1;
  }

  Future<void> sendMessage(String content) async {
    final convId = _convId;
    final repo = ref.read(chatRepositoryProvider);
    final previousMessages = state.value ?? [];

    // Optimistic UI: show the user message + assistant placeholder immediately.
    final userMsg = Message(
      conversationId: convId,
      role: MessageRole.user,
      content: content,
    );
    final placeholder = Message(
      conversationId: convId,
      role: MessageRole.assistant,
      content: '',
      status: 'generating',
    );
    state = AsyncData([...previousMessages, userMsg, placeholder]);
    ref.read(isStreamingProvider(convId).notifier).state = true;

    try {
      // ── Step 1: POST the message. If this fails the server never got it. ──
      await repo.sendMessage(convId, content);
    } catch (e) {
      state = AsyncData(previousMessages);
      ref.read(isStreamingProvider(convId).notifier).state = false;
      rethrow;
    }

    // ── Step 2: Stream the response (best effort — silent on failure). ──
    //
    // We use a StreamIterator with a per-event timeout as a stall watchdog.
    // Mobile networks occasionally drop SSE sockets silently: the TCP
    // connection is half-closed, Dio never sees the close, and `await for`
    // hangs forever with `isStreaming=true`, freezing the input. If no
    // event arrives within the watchdog window we bail out and let the
    // polling pass below reconcile state from the server.
    const stallTimeout = Duration(seconds: 45);
    bool streamedContent = false;
    final iter = StreamIterator(repo.streamGeneration(convId));
    try {
      while (await iter.moveNext().timeout(stallTimeout)) {
        final event = iter.current;
        if (event is ChatTextChunk) {
          streamedContent = true;
          ref.read(streamingStatusProvider(convId).notifier).state = '';
          final msgs = state.requireValue;
          if (msgs.isEmpty) continue;
          final updated =
              msgs.last.copyWith(content: msgs.last.content + event.text);
          state = AsyncData([...msgs.sublist(0, msgs.length - 1), updated]);
        } else if (event is ChatStatusUpdate) {
          ref.read(streamingStatusProvider(convId).notifier).state =
              event.status;
        } else if (event is ChatToolCall) {
          // Append the tool call to the in-flight assistant message so the
          // chip appears live. The reload pass at the end of this function
          // will overwrite with the persisted version, which carries the
          // same shape — no de-dup needed.
          final msgs = state.requireValue;
          if (msgs.isEmpty) continue;
          final last = msgs.last;
          if (last.role != MessageRole.assistant) continue;
          final nextCalls = [...?last.toolCalls, event.toolCall];
          final updated = last.copyWith(toolCalls: nextCalls);
          state = AsyncData([...msgs.sublist(0, msgs.length - 1), updated]);
        }
      }
    } on TimeoutException {
      // Stall watchdog — no SSE event for stallTimeout. Fall through to
      // polling so the UI eventually unfreezes even if the socket is dead.
    } catch (_) {
      // SSE failed — fall through to the polling reload below.
    } finally {
      await iter.cancel();
    }

    // ── Step 3: Poll the API until we have a completed assistant response.
    //
    // We always retry up to 20 times regardless of whether SSE streamed
    // content. SSE ending does NOT mean the server has finished writing to
    // the DB — there can be a brief gap between the stream closing and the
    // row being marked complete.
    //
    // To avoid clobbering SSE-streamed text with a stale empty DB row, we
    // only replace state when the polled data is complete OR has content.
    try {
      for (var attempt = 0; attempt < 20; attempt++) {
        if (attempt > 0) await Future.delayed(const Duration(seconds: 2));
        final (conv, fresh) = await repo.getMessages(convId);
        ref.read(conversationsProvider.notifier).patchConversation(conv);

        final done = fresh.any(
          (m) => m.role == MessageRole.assistant && m.status != 'generating',
        );
        final polledHasContent = fresh.any(
          (m) => m.role == MessageRole.assistant && m.content.isNotEmpty,
        );

        // Update UI state when:
        //   • SSE never delivered anything (polling is the only data source), or
        //   • the polled response is complete or has content to show.
        if (!streamedContent || done || polledHasContent) {
          state = AsyncData(fresh);
        }

        if (done) break;
      }
    } catch (_) {
      // Polling failed entirely — clear the generating placeholder so the UI
      // doesn't spin forever.
      final msgs = state.value;
      if (msgs != null && msgs.isNotEmpty && msgs.last.status == 'generating') {
        state = AsyncData([
          ...msgs.sublist(0, msgs.length - 1),
          msgs.last.copyWith(status: 'complete'),
        ]);
      }
    } finally {
      ref.read(isStreamingProvider(convId).notifier).state = false;
      ref.read(streamingStatusProvider(convId).notifier).state = '';
    }
  }
}