using System;
using System.Collections.Generic;
using System.Globalization;
using System.IO;
using Newtonsoft.Json;
using Newtonsoft.Json.Linq;

namespace AIRuntime.Runtime.Providers
{
    /// <summary>Provider-reported totals, independent of the bounded transport log.</summary>
    public sealed class ProviderTokenUsage
    {
        public readonly struct Totals
        {
            public readonly long Sent, Received, ReportedResponses, MissingResponses;
            public Totals(long sent, long received, long reported, long missing)
            { Sent = sent; Received = received; ReportedResponses = reported; MissingResponses = missing; }
        }

        public static ProviderTokenUsage Shared { get; } = new ProviderTokenUsage();
        private readonly object _sync = new object();
        private readonly HashSet<string> _responses = new HashSet<string>(StringComparer.Ordinal);
        private long _sent, _received, _reported, _missing;

        public Totals Snapshot()
        {
            lock (_sync) return new Totals(_sent, _received, _reported, _missing);
        }

        public void Reset()
        {
            lock (_sync)
            {
                _sent = _received = _reported = _missing = 0;
                // Keep response identities so a duplicate terminal event cannot
                // re-add usage from before the reset.
            }
        }

        public void RecordResponse(string requestId, string payload)
        {
            bool reported = TryReadUsage(payload, out long sent, out long received);
            lock (_sync)
            {
                if (!string.IsNullOrEmpty(requestId) && !_responses.Add(requestId)) return;
                if (reported) { _sent += sent; _received += received; _reported++; }
                else _missing++;
            }
        }

        private static bool TryReadUsage(string payload, out long sent, out long received)
        {
            sent = received = 0;
            if (string.IsNullOrWhiteSpace(payload)) return false;
            if (TryReadObject(payload, out sent, out received)) return true;
            bool found = false;
            using (var reader = new StringReader(payload))
            {
                string line;
                while ((line = reader.ReadLine()) != null)
                {
                    line = line.Trim();
                    if (line.StartsWith("data:", StringComparison.Ordinal)) line = line.Substring(5).Trim();
                    if (!TryReadObject(line, out long input, out long output)) continue;
                    // Streaming usage frames are cumulative, not token deltas.
                    sent = Math.Max(sent, input);
                    received = Math.Max(received, output);
                    found = true;
                }
            }
            return found;
        }

        private static bool TryReadObject(string json, out long sent, out long received)
        {
            sent = received = 0;
            if (string.IsNullOrWhiteSpace(json) || json.TrimStart()[0] != '{') return false;
            try
            {
                var root = JObject.Parse(json);
                var usage = root["usage"] as JObject;
                return usage != null
                    ? TryCount(usage["prompt_tokens"] ?? usage["input_tokens"], out sent) &&
                      TryCount(usage["completion_tokens"] ?? usage["output_tokens"], out received)
                    : TryCount(root["prompt_eval_count"], out sent) && TryCount(root["eval_count"], out received);
            }
            catch (JsonException) { return false; }
        }

        private static bool TryCount(JToken token, out long count)
        {
            count = 0;
            return token?.Type == JTokenType.Integer &&
                long.TryParse(token.ToString(), NumberStyles.Integer, CultureInfo.InvariantCulture, out count) && count >= 0;
        }
    }
}
