using System;
using System.Reflection;
using System.Threading.Tasks;
using AIRuntime.Runtime.Providers;
using NUnit.Framework;

namespace AIRuntime.Tests
{
    public sealed class ProviderTokenUsageTests
    {
        [Test]
        public void TransportCountsHttpResponseWithoutCountingDiagnosticCopies()
        {
            var before = ProviderTokenUsage.Shared.Snapshot();
            var publish = typeof(ProviderTransportDebug).GetMethod("Publish", BindingFlags.NonPublic | BindingFlags.Static);
            using (ProviderTransportDebug.BeginAttempt("usage-test", Guid.NewGuid().ToString("N"), 1))
                foreach (string direction in new[] { "REQUEST", "HTTP RESPONSE", "TRANSPORT END", "ADAPTER INPUT" })
                    publish.Invoke(null, new object[] { "usage-test", direction, "",
                        "{\"usage\":{\"prompt_tokens\":20,\"completion_tokens\":5}}", false, null });
            var after = ProviderTokenUsage.Shared.Snapshot();
            Assert.AreEqual(20, after.Sent - before.Sent);
            Assert.AreEqual(5, after.Received - before.Received);
            Assert.AreEqual(1, after.ReportedResponses - before.ReportedResponses);
        }

        [Test]
        public void CountsResponsesAcrossProvidersAndAttempts()
        {
            var usage = new ProviderTokenUsage();
            usage.RecordResponse("first", "{\"usage\":{\"prompt_tokens\":123,\"completion_tokens\":45}}");
            usage.RecordResponse("fallback", "{\"prompt_eval_count\":20,\"eval_count\":7}");
            usage.RecordResponse("third", "{\"usage\":{\"input_tokens\":10,\"output_tokens\":2}}");
            var totals = usage.Snapshot();
            Assert.AreEqual(153, totals.Sent);
            Assert.AreEqual(54, totals.Received);
            Assert.AreEqual(3, totals.ReportedResponses);
        }

        [Test]
        public void StreamingUsageIsCumulativeAndResponseIsCountedOnce()
        {
            var usage = new ProviderTokenUsage();
            string response = "data: {\"choices\":[]}\n\ndata: {\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":4}}\n\n" +
                "data: {\"usage\":{\"prompt_tokens\":100,\"completion_tokens\":8}}\n\ndata: [DONE]\n";
            usage.RecordResponse("request", response);
            usage.RecordResponse("request", response);
            Assert.AreEqual(100, usage.Snapshot().Sent);
            Assert.AreEqual(8, usage.Snapshot().Received);
            Assert.AreEqual(1, usage.Snapshot().ReportedResponses);
        }

        [Test]
        public void MissingMalformedAndPartialUsageAreNotReportedAsZeroUsage()
        {
            var usage = new ProviderTokenUsage();
            usage.RecordResponse("empty", "");
            usage.RecordResponse("error", "{\"error\":\"Bad request\"}");
            usage.RecordResponse("broken", "data: {broken\n");
            usage.RecordResponse("partial", "{\"usage\":{\"prompt_tokens\":100}}");
            usage.RecordResponse("negative", "{\"usage\":{\"prompt_tokens\":-1,\"completion_tokens\":1}}");
            usage.RecordResponse("zero", "{\"usage\":{\"prompt_tokens\":0,\"completion_tokens\":0}}");
            Assert.AreEqual(5, usage.Snapshot().MissingResponses);
            Assert.AreEqual(1, usage.Snapshot().ReportedResponses);
            Assert.AreEqual(0, usage.Snapshot().Sent);
        }

        [Test]
        public void ResetDoesNotRecountAnOldResponse()
        {
            var usage = new ProviderTokenUsage();
            const string response = "{\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":3}}";
            usage.RecordResponse("old", response);
            usage.Reset();
            usage.RecordResponse("old", response);
            Assert.AreEqual(0, usage.Snapshot().ReportedResponses);
            usage.RecordResponse("new", response);
            Assert.AreEqual(10, usage.Snapshot().Sent);
        }

        [Test]
        public void ConcurrentResponsesAndLargeCountsRemainAccurate()
        {
            var usage = new ProviderTokenUsage();
            Parallel.For(0, 100, i => usage.RecordResponse(i.ToString(),
                "{\"usage\":{\"prompt_tokens\":30000000,\"completion_tokens\":2}}"));
            Assert.AreEqual(3000000000L, usage.Snapshot().Sent);
            Assert.AreEqual(200, usage.Snapshot().Received);
            Assert.AreEqual(100, usage.Snapshot().ReportedResponses);
        }
    }
}
