The-Agency/lib/ui/providers/chat_provider.dart

import "package:flutter/foundation.dart";
import "dart:convert";

import "../../src/chat/tool_loop_service.dart";
import "../../src/api/openrouter_client.dart";
import "../../src/hooks/hook_loader.dart";
import "../../src/hooks/hook_runner.dart";
import "../../src/hooks/hook_types.dart";
import "../../src/permissions/permission_types.dart";
import "../../src/session/conversation_history.dart";
import "../../src/session/session_store.dart";
import "../../src/session/session_types.dart";
import "../../src/services/cost_tracker.dart" as cost_tracker;
import "settings_provider.dart";

enum QueuePriority {
  now(0),
  next(1),
  later(2);

  final int order;
  const QueuePriority(this.order);
}

class QueuedMessage {
  final String text;
  final QueuePriority priority;

  const QueuedMessage({required this.text, required this.priority});
}


class ChatProvider extends ChangeNotifier {
  ChatProvider(this._settingsProvider) {
    _initHooks();
  }

  final SettingsProvider _settingsProvider;
  ToolLoopService _toolLoopService = ToolLoopService();
  HookRunner? _hookRunner;
  ConversationHistory? _conversationHistory;
  OpenRouterClient? _client;
  bool _stopRequested = false;
  PendingPermission? _pendingPermission;

  PendingPermission? get pendingPermission => _pendingPermission;

  Future<void> _initHooks() async {
    try {
      final hooks = await HookLoader.loadHooks();
      _hookRunner = HookRunner(hooks: hooks);
      _toolLoopService = ToolLoopService(hookRunner: _hookRunner);
    } catch (e) {
      // hooks are optional, carry on without them
      print("Hook init failed: $e");
    }
  }

  List<Map<String, dynamic>> _apiMessages = <Map<String, dynamic>>[];
  bool isLoading = false;
  final List<QueuedMessage> _messageQueue = [];

  List<Message> get messages => _conversationHistory?.getMessages() ?? const [];
  int get messageCount => messages.length;
  String? get workingDirectory => _conversationHistory?.session?.workingDirectory;

  /// Context window size from the last API response — derived from persisted
  /// message data, same as Claude Code (walks backwards to find the last
  /// assistant message that has contextTokens set).
  int get contextTokens {
    final msgs = messages;
    for (var i = msgs.length - 1; i >= 0; i--) {
      final ct = msgs[i].contextTokens;
      if (ct != null && ct > 0) return ct;
    }
    return 0;
  }
  bool get hasConversation => _conversationHistory != null;
  bool get isStopping => _stopRequested;
  int get queuedMessageCount => _messageQueue.length;

  // only user-visible messages (priority != now)
  List<String> get queuedMessages =>
    List.unmodifiable(_messageQueue.map((m) => m.text));

  void removeQueuedMessage(int index) {
    if (index < 0 || index >= _messageQueue.length) return;
    _messageQueue.removeAt(index);
    notifyListeners();
  }

  QueuedMessage? _dequeue() {
    if (_messageQueue.isEmpty) return null;

    int bestIdx = 0;
    for (int i = 1; i < _messageQueue.length; i++) {
      if (_messageQueue[i].priority.order < _messageQueue[bestIdx].priority.order) {
        bestIdx = i;
      }
    }

    final cmd = _messageQueue[bestIdx];
    _messageQueue.removeAt(bestIdx);
    return cmd;
  }

  void setConversation(ConversationHistory history) {
    _conversationHistory = history;
    _apiMessages = _buildApiMessages(history.getMessages());
    notifyListeners();
  }

  void clearConversation() {
    _conversationHistory = null;
    _apiMessages = <Map<String, dynamic>>[];
    _messageQueue.clear();
    isLoading = false;
    notifyListeners();
  }

  Future<void> sendMessage(String text, {QueuePriority priority = QueuePriority.next}) async {
    if (text.isEmpty || _conversationHistory == null) return;

    if (isLoading) {
      _messageQueue.add(QueuedMessage(text: text, priority: priority));
      notifyListeners();
      return;
    }

    final apiKey = _settingsProvider.settings.openRouterApiKey;
    if (apiKey == null || apiKey.isEmpty) {
      throw Exception(
        "OpenRouter API key not set. Please configure it in settings.",
      );
    }

    final savedModel = _settingsProvider.settings.model;
    final model = _settingsProvider.normalizeModelId(savedModel);
    if (savedModel != model) {
      print("Normalizing legacy model ID from $savedModel to $model");
      await _settingsProvider.updateModel(model);
    }

    try {
      _stopRequested = false;
      bool hasStreamingAssistantMessage = false;
      _client = await OpenRouterClientFactory.create(apiKey: apiKey);
      final session = _conversationHistory!.session;
      final workingDirectory = session?.workingDirectory;
      if (session != null) {
        session.model = model;
        if (session.name == "New Chat") {
          session.name = _buildSessionName(text);
        }
      }

      // fire UserPromptSubmit hook
      await _hookRunner?.runHooksForKind(
        HookKind.userPromptSubmit,
        input: {"message": text},
      );

      // add user message to conversation
      _conversationHistory!.addMessage("user", text);

      _apiMessages.add(<String, dynamic>{"role": "user", "content": text});
      isLoading = true;
      notifyListeners();

      final advisorModel = _settingsProvider.settings.advisorModel;

      final toolLoopResult = await _toolLoopService.runTurn(
        client: _client!,
        model: model,
        apiKey: apiKey,
        getSettings: () => _settingsProvider.settings,
        apiMessages: _apiMessages.take(_apiMessages.length - 1).toList(),
        userText: text,
        workingDirectory: workingDirectory,
        advisorModel: advisorModel,
        onToolCall: (toolName, input) {
          _conversationHistory!.addMessage(
            "tool",
            _formatToolCall(toolName, input),
          );
          notifyListeners();
        },
        onToolResult: (toolName, result) {
          _conversationHistory!.addMessage(
            "tool",
            _formatToolResult(toolName, result),
          );
          notifyListeners();
        },
        onAssistantTextDelta: (delta) {
          if (!hasStreamingAssistantMessage) {
            _conversationHistory!.addMessage("assistant", "");
            hasStreamingAssistantMessage = true;
          }
          _conversationHistory!.appendToLastMessage(delta);
          notifyListeners();
        },
        onAssistantMessageComplete: () {
          hasStreamingAssistantMessage = false;
          notifyListeners();
        },
        onPermissionRequired: (toolName, input) async {
          final pending = PendingPermission(toolName: toolName, input: input);
          _pendingPermission = pending;
          notifyListeners();
          final decision = await pending.future;
          _pendingPermission = null;
          notifyListeners();
          return decision;
        },
      );
      _apiMessages = toolLoopResult.apiMessages;

      final ct = toolLoopResult.response.contextTokens;

      // add assistant message to visible conversation
      if (!toolLoopResult.finalResponseWasStreamed) {
        _conversationHistory!.addMessage(
          "assistant",
          toolLoopResult.responseText,
          tokens: toolLoopResult.response.outputTokens,
          contextTokens: ct,
        );
      } else {
        // streamed message was built incrementally — patch contextTokens onto it
        _conversationHistory!.setLastMessageContextTokens(ct);
      }

      // track cost (set to 0 for now — OpenRouter pricing varies by model)
      final inputTokens = toolLoopResult.response.inputTokens ?? 0;
      final outputTokens = toolLoopResult.response.outputTokens ?? 0;

      cost_tracker.addToTotalSessionCost(
        cost: 0.0,
        inputTokens: inputTokens,
        outputTokens: outputTokens,
        cacheReadTokens: 0,
        cacheCreationTokens: 0,
        webSearchRequests: toolLoopResult.webSearchRequests,
        webFetchRequests: toolLoopResult.webFetchRequests,
        model: toolLoopResult.response.model,
      );

      // save session
      if (session != null) {
        await SessionStore.instance.saveSession(session);
      }

      notifyListeners();
    } catch (error, stackTrace) {
      print("Failed to send message: $error");
      print(stackTrace);

      if (error is RequestCancelledException) {
        _conversationHistory!.addMessage("assistant", "Generation stopped.");
        final session = _conversationHistory!.session;
  
        if (session != null) {
          await SessionStore.instance.saveSession(session);
        }
        return;
      }

      if (error is ToolLoopException) {
        _apiMessages = List<Map<String, dynamic>>.from(error.apiMessages);
      }

      _conversationHistory!.addMessage(
        "assistant",
        _buildTurnFailureMessage(error),
      );

      final session = _conversationHistory!.session;

      if (session != null) {
        await SessionStore.instance.saveSession(session);
      }
      rethrow;
    } finally {
      _client?.close();
      _client = null;
      _stopRequested = false;
      isLoading = false;
      notifyListeners();
    }

    final next = _dequeue();
    if (next != null) {
      notifyListeners();
      await sendMessage(next.text, priority: next.priority);
    }
  }

  void resolvePermission(PermissionDecision decision) async {
    final pending = _pendingPermission;
    if (pending == null) return;

    if (decision == PermissionDecision.allowAlways) {
      // persist to settings so this tool is auto-allowed from now on
      await _settingsProvider.addAlwaysAllowRule(pending.toolName);
    }

    pending.resolve(decision);
    _pendingPermission = null;
    notifyListeners();
  }

  void stopGenerating() {
    if (!isLoading) {
      return;
    }

    _pendingPermission?.resolve(PermissionDecision.reject);
    _pendingPermission = null;
    _messageQueue.clear();
    _stopRequested = true;
    print("Stopping active turn");
    _client?.cancelActiveRequest();
    notifyListeners();

    _hookRunner?.runHooksForKind(HookKind.stop);
  }

  @override
  void dispose() {
    _client?.close();
    super.dispose();
  }

  String _buildSessionName(String text) {
    final sanitized = text.replaceAll(RegExp(r"\s+"), " ").trim();
    if (sanitized.isEmpty) {
      return "New Chat";
    }

    const maxLength = 48;
    if (sanitized.length <= maxLength) {
      return sanitized;
    }

    return "${sanitized.substring(0, maxLength - 1).trimRight()}…";
  }

  List<Map<String, dynamic>> _buildApiMessages(List<Message> messages) {
    return messages
        .where(
          (message) => message.role == "user" || message.role == "assistant",
        )
        .map(
          (message) => <String, dynamic>{
            "role": message.role,
            "content": message.content,
          },
        )
        .toList(growable: true);
  }

  String _formatToolCall(String toolName, Map<String, dynamic> input) {
    const encoder = JsonEncoder.withIndent("  ");
    final visibleInput = Map<String, dynamic>.fromEntries(
      input.entries.where((entry) => !entry.key.startsWith("_")),
    );
    return "$toolName call\n${encoder.convert(visibleInput)}";
  }

  String _formatToolResult(String toolName, String result) {
    return "$toolName result\n$result";
  }

  String _buildTurnFailureMessage(Object error) {
    return "This turn failed before the assistant could finish: $error";
  }
}