using System.Text.Json;
using BuildFlow.Application.RagQA.DTOs;
using BuildFlow.Application.RagQA.Queries;
using BuildFlow.Host.Endpoints.Rag.RAG;
using FastEndpoints;
using FluentAssertions;
using MediatR;
using NSubstitute;
using Xunit;

namespace BuildFlow.Host.Tests;

/// <summary>
/// RAG 流式 Endpoint 测试，固定既有 Server-Sent Events 协议。
/// </summary>
public class RagStreamEndpointTests
{
    [Fact]
    public async Task EndpointEncodesSourcesContentAndDoneUsingExistingSseProtocol()
    {
        var documentId = Guid.NewGuid();
        var sources = new List<SourceChunk>
        {
            new(documentId, "架构说明", "来源片段", 0.9),
        };
        var mediator = Substitute.For<IMediator>();
        mediator.CreateStream(Arg.Any<RAGStreamQuery>(), Arg.Any<CancellationToken>())
            .Returns(CreateEvents(sources));
        await using var body = new MemoryStream();
        var endpoint = Factory.Create<RAGStreamEndpoint>(
            context => context.Response.Body = body,
            mediator);
        var request = new RAGQueryRequest(Guid.NewGuid(), "问题");

        await endpoint.HandleAsync(request, CancellationToken.None);

        _ = mediator.Received(1).CreateStream(
            Arg.Is<RAGStreamQuery>(query =>
                query.KnowledgeBaseId == request.KnowledgeBaseId
                && query.Question == request.Question),
            Arg.Any<CancellationToken>());
        body.Position = 0;
        var response = await new StreamReader(body).ReadToEndAsync();
        var sourcesJson = JsonSerializer.Serialize(new { sources });
        var contentJson = JsonSerializer.Serialize(new { content = "回答" });
        response.Should().Be(
            $"event: sources\ndata: {sourcesJson}\n\n" +
            $"data: {contentJson}\n\n" +
            "data: [DONE]\n\n");
        endpoint.HttpContext.Response.ContentType.Should().Be("text/event-stream");
        endpoint.HttpContext.Response.Headers.CacheControl.ToString().Should().Be("no-cache");
        endpoint.HttpContext.Response.Headers.Connection.ToString().Should().Be("keep-alive");
    }

    private static async IAsyncEnumerable<RAGStreamEvent> CreateEvents(
        IReadOnlyList<SourceChunk> sources)
    {
        await Task.Yield();
        yield return RAGStreamEvent.FromSources(sources);
        yield return RAGStreamEvent.Chunk("回答");
        yield return RAGStreamEvent.Done();
    }
}