using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace eagle {
    public interface IClientTextListener {
        void OnTextUpdate(string text, bool completed);
        string TextId();
    }

    public class TextEntry {
        public string Text { get; }
        public bool Completed { get; }

        public TextEntry(string text, bool completed) {
            Text = text;
            Completed = completed;
        }
    }

    /// <summary>
    /// Thread-safe provider for streaming LLM text content.
    ///
    /// HandleNewStreamingText can be called from any thread (e.g., gRPC thread).
    /// ProcessPendingUpdates must be called from the main thread (once per frame)
    /// to notify listeners of changes.
    /// </summary>
    public class ClientTextProvider {
        public static readonly ClientTextProvider Provider = new();

        // Lock for thread-safe access to text dictionary
        // Using a lock instead of ConcurrentDictionary because HandleNewStreamingText
        // does a read-modify-write that must be atomic
        private readonly object _lock = new();
        private readonly Dictionary<String, TextEntry> _streamingTexts = new();

        // Track which text IDs have pending updates
        private readonly HashSet<String> _pendingUpdates = new();

        // Listeners are only added/removed from main thread
        private readonly HashSet<IClientTextListener> _listeners = new();

        public void Clear() {
            lock (_lock) {
                _streamingTexts.Clear();
                _pendingUpdates.Clear();
            }
        }

        public Dictionary<String, TextEntry> All() {
            lock (_lock) { return new Dictionary<string, TextEntry>(_streamingTexts); }
        }

        /// <summary>
        /// Updates the text dictionary. Thread-safe - can be called from any thread.
        /// Listeners are NOT notified here; call ProcessPendingUpdates from main thread.
        /// </summary>
        public String
        HandleNewStreamingText(String llmId, String newText, Int32 knownByteCount, bool completed) {
            lock (_lock) {
                var currentText = "";
                if (_streamingTexts.TryGetValue(llmId, out var entry)) { currentText = entry.Text; }

                var currentTextBytes = Encoding.UTF8.GetBytes(currentText);
                var truncatedBytes = currentTextBytes.Take(knownByteCount).ToArray();

                var updatedText = Encoding.UTF8.GetString(truncatedBytes) + newText;
                _streamingTexts[llmId] = new TextEntry(updatedText, completed);

                // Mark this text ID as having pending updates
                _pendingUpdates.Add(llmId);

                return updatedText;
            }
        }

        /// <summary>
        /// Process pending updates and notify listeners. Must be called from main thread.
        /// This batches multiple updates to the same text ID into a single notification per frame.
        /// </summary>
        public void ProcessPendingUpdates() {
            List<(String llmId, TextEntry entry)> updates;

            lock (_lock) {
                if (_pendingUpdates.Count == 0) return;

                // Collect pending updates and their current values
                updates = _pendingUpdates.Where(id => _streamingTexts.ContainsKey(id))
                                  .Select(id => (id, _streamingTexts[id]))
                                  .ToList();

                _pendingUpdates.Clear();
            }

            // Notify listeners outside the lock to avoid potential deadlocks
            foreach (var (llmId, textEntry) in updates) {
                foreach (var listener in _listeners.Where(x => x.TextId() == llmId)) {
                    listener.OnTextUpdate(textEntry.Text, textEntry.Completed);
                }
            }
        }

        public TextEntry GetTextEntry(string streamId) {
            if (String.IsNullOrEmpty(streamId)) return new TextEntry("", false);

            lock (_lock) {
                _streamingTexts.TryGetValue(streamId, out var entry);
                return entry;
            }
        }

        public void AddListener(IClientTextListener listener) {
            _listeners.Add(listener);
            if (listener.TextId() is {} textId && !String.IsNullOrEmpty(textId)) {
                var entry = GetTextEntry(textId);
                if (entry != null) { listener.OnTextUpdate(entry.Text, entry.Completed); }
            }
        }

        public void RemoveListener(IClientTextListener listener) { _listeners.Remove(listener); }
    }
}
