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

    public class ClientTextProvider {
        public static readonly ClientTextProvider Provider = new();

        private readonly Dictionary<String, TextEntry> _streamingTexts = new();

        private readonly HashSet<IClientTextListener> _listeners = new();

        public void Clear() { _streamingTexts.Clear(); }

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

        public String
        HandleNewStreamingText(String llmId, String newText, Int32 knownByteCount, bool completed) {
            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);

            _listeners.Where(x => x.TextId() == llmId)
                    .ToList()
                    .ForEach(x => x.OnTextUpdate(updatedText, completed));

            return updatedText;
        }

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

            if (ClientPregeneratedText.Provider.TryGetText(streamId, out var text)) {
                return new TextEntry(text, true);
            }

            return _streamingTexts.GetValueOrDefault(streamId, null);
        }

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