// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
 
using System.Threading.Channels;
using Aspire.Dashboard.Model;
using Aspire.Dashboard.Model.Assistant;
using Aspire.Dashboard.Otlp.Model;
using Aspire.Dashboard.Otlp.Storage;
using Aspire.Dashboard.Tests.Integration.Playwright.Infrastructure;
using Aspire.Dashboard.Tests.Shared;
using Aspire.Tests.Shared.DashboardModel;
using Google.Protobuf.Collections;
using OpenTelemetry.Proto.Logs.V1;
using OpenTelemetry.Proto.Trace.V1;
using Xunit;
using static Aspire.Tests.Shared.Telemetry.TelemetryTestHelpers;
 
namespace Aspire.Dashboard.Tests.Model.AIAssistant;
 
public class AssistantChatDataContextTests
{
    private static readonly DateTime s_testTime = new(1970, 1, 1, 0, 0, 0, DateTimeKind.Utc);
 
    [Fact]
    public async Task GetStructuredLogs_ExceedTokenLimit_ReturnMostRecentItems()
    {
        // Arrange
        var repository = CreateRepository();
 
        var scopeLogs = new ScopeLogs();
        for (var i = 0; i < 20; i++)
        {
            var logRecord = CreateLogRecord(message: $"Log {i}: {new string((char)('a' + i), 10_000)}", time: s_testTime.AddMinutes(i));
            scopeLogs.LogRecords.Add(logRecord);
        }
        var addContext = new AddContext();
        repository.AddLogs(addContext, new RepeatedField<ResourceLogs>()
        {
            new ResourceLogs
            {
                Resource = CreateResource(),
                ScopeLogs = { scopeLogs }
            }
        });
        var dataContext = CreateAssistantChatDataContext(telemetryRepository: repository);
 
        // Act
        var result = await dataContext.GetStructuredLogsAsync(resourceName: null, CancellationToken.None);
 
        // Assert
        for (var i = 6; i < 20; i++)
        {
            Assert.Contains($"Log {i}:", result);
        }
        Assert.Contains("Returned latest 14 log entries. Earlier 6 log entries not returned because of size limits.", result);
    }
 
    [Fact]
    public async Task GetTraces_ExceedTokenLimit_ReturnMostRecentItems()
    {
        // Arrange
        var repository = CreateRepository();
 
        var scopeSpans = new ScopeSpans();
        for (var i = 0; i < 20; i++)
        {
            var span = CreateSpan(traceId: $"{i}", spanId: $"{i}-1", startTime: s_testTime.AddMinutes(i), endTime: s_testTime.AddMinutes(10), attributes: [new KeyValuePair<string, string>("message", $"Log {i}: {new string((char)('a' + i), 10_000)}")]);
            scopeSpans.Spans.Add(span);
        }
        var addContext = new AddContext();
        repository.AddTraces(addContext, new RepeatedField<ResourceSpans>()
        {
            new ResourceSpans
            {
                Resource = CreateResource(),
                ScopeSpans = { scopeSpans }
            }
        });
        var dataContext = CreateAssistantChatDataContext(telemetryRepository: repository);
 
        // Act
        var result = await dataContext.GetTracesAsync(resourceName: null, CancellationToken.None);
 
        // Assert
        for (var i = 7; i < 20; i++)
        {
            Assert.Contains($"Test span. Id: {i}", result);
        }
        Assert.Contains("Returned latest 13 traces. Earlier 7 traces not returned because of size limits.", result);
    }
 
    [Fact]
    public async Task GetConsoleLogs_ExceedTokenLimit_ReturnMostRecentItems()
    {
        // Arrange
        var consoleLogsChannel = Channel.CreateUnbounded<IReadOnlyList<ResourceLogLine>>();
        var testResource = ModelTestHelpers.CreateResource(resourceName: "test-resource", state: KnownResourceState.Running);
        var dashboardClient = new TestDashboardClient(
            isEnabled: true,
            consoleLogsChannelProvider: name =>
            {
                return consoleLogsChannel;
            },
            initialResources: [testResource]);
 
        var dataContext = CreateAssistantChatDataContext(dashboardClient: dashboardClient);
 
        for (var i = 0; i < 20; i++)
        {
            var line = new string((char)('a' + i), 10_000);
            consoleLogsChannel.Writer.TryWrite([new ResourceLogLine(i + 1, line, IsErrorMessage: false)]);
        }
        consoleLogsChannel.Writer.Complete();
 
        // Act
        var result = await dataContext.GetConsoleLogsAsync(resourceName: "test-resource", CancellationToken.None);
 
        // Assert
        for (var i = 5; i < 20; i++)
        {
            var line = AIHelpers.LimitLength(new string((char)('a' + i), 10_000));
            Assert.Contains(line, result);
        }
        Assert.Contains("Returned latest 15 console logs. Earlier 5 console logs not returned because of size limits.", result);
    }
 
    internal static AssistantChatDataContext CreateAssistantChatDataContext(TelemetryRepository? telemetryRepository = null, IDashboardClient? dashboardClient = null)
    {
        var context = new AssistantChatDataContext(
            telemetryRepository ?? CreateRepository(),
            dashboardClient ?? new MockDashboardClient(),
            [],
            new TestStringLocalizer<Dashboard.Resources.AIAssistant>());
 
        return context;
    }
}