using System.Text;
using System.Text.Json;
using NATS.Server.Events;
using NATS.Server.Internal;

namespace NATS.Server.Core.Tests.Internal;

/// <summary>
/// Tests for MsgTraceContext: header parsing, event collection, trace propagation,
/// JetStream two-phase send, hop tracking, and JSON serialization.
/// Go reference: msgtrace.go — initMsgTrace, sendEvent, addEgressEvent,
/// addJetStreamEvent, genHeaderMapIfTraceHeadersPresent.
/// </summary>
public class MessageTraceContextTests
{
    private static ReadOnlyMemory<byte> BuildHeaders(params (string key, string value)[] headers)
    {
        var sb = new StringBuilder("NATS/1.0\r\n");
        foreach (var (key, value) in headers)
        {
            sb.Append($"{key}: {value}\r\n");
        }
        sb.Append("\r\n");
        return Encoding.ASCII.GetBytes(sb.ToString());
    }

    // --- Header parsing ---

    [Fact]
    public void ParseTraceHeaders_returns_null_for_no_trace_headers()
    {
        var headers = BuildHeaders(("Content-Type", "text/plain"));
        var result = MsgTraceContext.ParseTraceHeaders(headers.Span);
        result.ShouldBeNull();
    }

    [Fact]
    public void ParseTraceHeaders_returns_map_when_trace_dest_present()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.subject"),
            ("Content-Type", "text/plain"));
        var result = MsgTraceContext.ParseTraceHeaders(headers.Span);
        result.ShouldNotBeNull();
        result.ShouldContainKey(MsgTraceHeaders.TraceDest);
        result[MsgTraceHeaders.TraceDest][0].ShouldBe("trace.subject");
    }

    [Fact]
    public void ParseTraceHeaders_returns_null_when_trace_disabled()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, MsgTraceHeaders.TraceDestDisabled));
        var result = MsgTraceContext.ParseTraceHeaders(headers.Span);
        result.ShouldBeNull();
    }

    [Fact]
    public void ParseTraceHeaders_detects_traceparent_with_sampled_flag()
    {
        // W3C trace context: version-traceid-parentid-flags (01 = sampled)
        var headers = BuildHeaders(
            ("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"));
        var result = MsgTraceContext.ParseTraceHeaders(headers.Span);
        result.ShouldNotBeNull();
        result.ShouldContainKey("traceparent");
    }

    [Fact]
    public void ParseTraceHeaders_ignores_traceparent_without_sampled_flag()
    {
        // flags=00 means not sampled
        var headers = BuildHeaders(
            ("traceparent", "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-00"));
        var result = MsgTraceContext.ParseTraceHeaders(headers.Span);
        result.ShouldBeNull();
    }

    [Fact]
    public void ParseTraceHeaders_returns_null_for_empty_input()
    {
        var result = MsgTraceContext.ParseTraceHeaders(ReadOnlySpan<byte>.Empty);
        result.ShouldBeNull();
    }

    [Fact]
    public void ParseTraceHeaders_returns_null_for_non_nats_header()
    {
        var headers = Encoding.ASCII.GetBytes("HTTP/1.1 200 OK\r\nFoo: bar\r\n\r\n");
        var result = MsgTraceContext.ParseTraceHeaders(headers);
        result.ShouldBeNull();
    }

    // --- Context creation ---

    [Fact]
    public void Create_returns_null_for_empty_headers()
    {
        var ctx = MsgTraceContext.Create(
            ReadOnlyMemory<byte>.Empty,
            clientId: 1,
            clientName: "test",
            accountName: "$G",
            subject: "test.sub",
            msgSize: 10);
        ctx.ShouldBeNull();
    }

    [Fact]
    public void Create_returns_null_for_headers_without_trace()
    {
        var headers = BuildHeaders(("Content-Type", "text/plain"));
        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "test",
            accountName: "$G",
            subject: "test.sub",
            msgSize: 10);
        ctx.ShouldBeNull();
    }

    [Fact]
    public void Create_builds_context_with_ingress_event()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 42,
            clientName: "my-publisher",
            accountName: "$G",
            subject: "orders.new",
            msgSize: 128);

        ctx.ShouldNotBeNull();
        ctx.IsActive.ShouldBeTrue();
        ctx.Destination.ShouldBe("trace.dest");
        ctx.TraceOnly.ShouldBeFalse();
        ctx.AccountName.ShouldBe("$G");

        // Check ingress event
        ctx.Event.Events.Count.ShouldBe(1);
        var ingress = ctx.Event.Events[0].ShouldBeOfType<MsgTraceIngress>();
        ingress.Type.ShouldBe(MsgTraceTypes.Ingress);
        ingress.Cid.ShouldBe(42UL);
        ingress.Name.ShouldBe("my-publisher");
        ingress.Account.ShouldBe("$G");
        ingress.Subject.ShouldBe("orders.new");
        ingress.Error.ShouldBeNull();

        // Check request info
        ctx.Event.Request.MsgSize.ShouldBe(128);
        ctx.Event.Request.Header.ShouldNotBeNull();
        ctx.Event.Request.Header.ShouldContainKey(MsgTraceHeaders.TraceDest);
    }

    [Fact]
    public void Create_with_trace_only_flag()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"),
            (MsgTraceHeaders.TraceOnly, "true"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "test",
            accountName: "$G",
            subject: "test",
            msgSize: 0);

        ctx.ShouldNotBeNull();
        ctx.TraceOnly.ShouldBeTrue();
    }

    [Fact]
    public void Create_with_trace_only_flag_numeric()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"),
            (MsgTraceHeaders.TraceOnly, "1"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "test",
            accountName: "$G",
            subject: "test",
            msgSize: 0);

        ctx.ShouldNotBeNull();
        ctx.TraceOnly.ShouldBeTrue();
    }

    [Fact]
    public void Create_without_trace_only_flag()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"),
            (MsgTraceHeaders.TraceOnly, "false"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "test",
            accountName: "$G",
            subject: "test",
            msgSize: 0);

        ctx.ShouldNotBeNull();
        ctx.TraceOnly.ShouldBeFalse();
    }

    [Fact]
    public void Create_captures_hop_from_non_client_kind()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"),
            (MsgTraceHeaders.TraceHop, "1.2"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "route-1",
            accountName: "$G",
            subject: "test",
            msgSize: 0,
            clientKind: MsgTraceContext.KindRouter);

        ctx.ShouldNotBeNull();
        ctx.Hop.ShouldBe("1.2");
    }

    [Fact]
    public void Create_ignores_hop_from_client_kind()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"),
            (MsgTraceHeaders.TraceHop, "1.2"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "test",
            accountName: "$G",
            subject: "test",
            msgSize: 0,
            clientKind: MsgTraceContext.KindClient);

        ctx.ShouldNotBeNull();
        ctx.Hop.ShouldBe(""); // Client hop is ignored
    }

    // --- Event recording ---

    [Fact]
    public void SetIngressError_sets_error_on_first_event()
    {
        var ctx = CreateSimpleContext();
        ctx.SetIngressError("publish denied");

        var ingress = ctx.Event.Events[0].ShouldBeOfType<MsgTraceIngress>();
        ingress.Error.ShouldBe("publish denied");
    }

    [Fact]
    public void AddSubjectMappingEvent_appends_mapping()
    {
        var ctx = CreateSimpleContext();
        ctx.AddSubjectMappingEvent("orders.mapped");

        ctx.Event.Events.Count.ShouldBe(2);
        var mapping = ctx.Event.Events[1].ShouldBeOfType<MsgTraceSubjectMapping>();
        mapping.Type.ShouldBe(MsgTraceTypes.SubjectMapping);
        mapping.MappedTo.ShouldBe("orders.mapped");
    }

    [Fact]
    public void AddEgressEvent_appends_egress_with_subscription_and_queue()
    {
        var ctx = CreateSimpleContext();
        ctx.AddEgressEvent(
            clientId: 99,
            clientName: "subscriber",
            clientKind: MsgTraceContext.KindClient,
            subscriptionSubject: "orders.>",
            queue: "workers");

        ctx.Event.Events.Count.ShouldBe(2);
        var egress = ctx.Event.Events[1].ShouldBeOfType<MsgTraceEgress>();
        egress.Type.ShouldBe(MsgTraceTypes.Egress);
        egress.Kind.ShouldBe(MsgTraceContext.KindClient);
        egress.Cid.ShouldBe(99UL);
        egress.Name.ShouldBe("subscriber");
        egress.Subscription.ShouldBe("orders.>");
        egress.Queue.ShouldBe("workers");
    }

    [Fact]
    public void AddEgressEvent_records_account_when_different_from_ingress()
    {
        var ctx = CreateSimpleContext(accountName: "acctA");
        ctx.AddEgressEvent(
            clientId: 99,
            clientName: "subscriber",
            clientKind: MsgTraceContext.KindClient,
            subscriptionSubject: "api.>",
            account: "acctB");

        var egress = ctx.Event.Events[1].ShouldBeOfType<MsgTraceEgress>();
        egress.Account.ShouldBe("acctB");
    }

    [Fact]
    public void AddEgressEvent_omits_account_when_same_as_ingress()
    {
        var ctx = CreateSimpleContext(accountName: "$G");
        ctx.AddEgressEvent(
            clientId: 99,
            clientName: "subscriber",
            clientKind: MsgTraceContext.KindClient,
            subscriptionSubject: "test",
            account: "$G");

        var egress = ctx.Event.Events[1].ShouldBeOfType<MsgTraceEgress>();
        egress.Account.ShouldBeNull();
    }

    [Fact]
    public void AddEgressEvent_for_router_omits_subscription_and_queue()
    {
        var ctx = CreateSimpleContext();
        ctx.AddEgressEvent(
            clientId: 1,
            clientName: "route-1",
            clientKind: MsgTraceContext.KindRouter,
            subscriptionSubject: "should.not.appear",
            queue: "should.not.appear");

        var egress = ctx.Event.Events[1].ShouldBeOfType<MsgTraceEgress>();
        egress.Subscription.ShouldBeNull();
        egress.Queue.ShouldBeNull();
    }

    [Fact]
    public void AddEgressEvent_with_error()
    {
        var ctx = CreateSimpleContext();
        ctx.AddEgressEvent(
            clientId: 50,
            clientName: "slow-client",
            clientKind: MsgTraceContext.KindClient,
            error: MsgTraceErrors.ClientClosed);

        var egress = ctx.Event.Events[1].ShouldBeOfType<MsgTraceEgress>();
        egress.Error.ShouldBe(MsgTraceErrors.ClientClosed);
    }

    [Fact]
    public void AddStreamExportEvent_records_account_and_target()
    {
        var ctx = CreateSimpleContext();
        ctx.AddStreamExportEvent("exportAccount", "export.subject");

        ctx.Event.Events.Count.ShouldBe(2);
        var se = ctx.Event.Events[1].ShouldBeOfType<MsgTraceStreamExport>();
        se.Type.ShouldBe(MsgTraceTypes.StreamExport);
        se.Account.ShouldBe("exportAccount");
        se.To.ShouldBe("export.subject");
    }

    [Fact]
    public void AddServiceImportEvent_records_from_and_to()
    {
        var ctx = CreateSimpleContext();
        ctx.AddServiceImportEvent("importAccount", "from.subject", "to.subject");

        ctx.Event.Events.Count.ShouldBe(2);
        var si = ctx.Event.Events[1].ShouldBeOfType<MsgTraceServiceImport>();
        si.Type.ShouldBe(MsgTraceTypes.ServiceImport);
        si.Account.ShouldBe("importAccount");
        si.From.ShouldBe("from.subject");
        si.To.ShouldBe("to.subject");
    }

    // --- JetStream events ---

    [Fact]
    public void AddJetStreamEvent_records_stream_name()
    {
        var ctx = CreateSimpleContext();
        ctx.AddJetStreamEvent("ORDERS");

        ctx.Event.Events.Count.ShouldBe(2);
        var js = ctx.Event.Events[1].ShouldBeOfType<MsgTraceJetStreamEntry>();
        js.Type.ShouldBe(MsgTraceTypes.JetStream);
        js.Stream.ShouldBe("ORDERS");
    }

    [Fact]
    public void UpdateJetStreamEvent_sets_subject_and_nointerest()
    {
        var ctx = CreateSimpleContext();
        ctx.AddJetStreamEvent("ORDERS");
        ctx.UpdateJetStreamEvent("orders.new", noInterest: true);

        var js = ctx.Event.Events[1].ShouldBeOfType<MsgTraceJetStreamEntry>();
        js.Subject.ShouldBe("orders.new");
        js.NoInterest.ShouldBeTrue();
    }

    [Fact]
    public void SendEventFromJetStream_requires_both_phases()
    {
        var ctx = CreateSimpleContext();
        ctx.AddJetStreamEvent("ORDERS");

        bool published = false;
        ctx.PublishCallback = (dest, reply, body) => { published = true; };

        // Phase 1: message path calls SendEvent — should not publish yet
        ctx.SendEvent();
        published.ShouldBeFalse();

        // Phase 2: JetStream path calls SendEventFromJetStream — now publishes
        ctx.SendEventFromJetStream();
        published.ShouldBeTrue();
    }

    [Fact]
    public void SendEventFromJetStream_with_error()
    {
        var ctx = CreateSimpleContext();
        ctx.AddJetStreamEvent("ORDERS");

        object? publishedBody = null;
        ctx.PublishCallback = (dest, reply, body) => { publishedBody = body; };

        ctx.SendEvent(); // Phase 1
        ctx.SendEventFromJetStream("stream full"); // Phase 2

        publishedBody.ShouldNotBeNull();
        var js = ctx.Event.Events[1].ShouldBeOfType<MsgTraceJetStreamEntry>();
        js.Error.ShouldBe("stream full");
    }

    // --- Hop tracking ---

    [Fact]
    public void SetHopHeader_increments_and_builds_hop_id()
    {
        var ctx = CreateSimpleContext();

        ctx.SetHopHeader();
        ctx.Event.Hops.ShouldBe(1);
        ctx.NextHop.ShouldBe("1");

        ctx.SetHopHeader();
        ctx.Event.Hops.ShouldBe(2);
        ctx.NextHop.ShouldBe("2");
    }

    [Fact]
    public void SetHopHeader_chains_from_existing_hop()
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, "trace.dest"),
            (MsgTraceHeaders.TraceHop, "1"));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "router",
            accountName: "$G",
            subject: "test",
            msgSize: 0,
            clientKind: MsgTraceContext.KindRouter);

        ctx.ShouldNotBeNull();
        ctx.Hop.ShouldBe("1");

        ctx.SetHopHeader();
        ctx.NextHop.ShouldBe("1.1");

        ctx.SetHopHeader();
        ctx.NextHop.ShouldBe("1.2");
    }

    [Fact]
    public void AddEgressEvent_captures_and_clears_next_hop()
    {
        var ctx = CreateSimpleContext();
        ctx.SetHopHeader();
        ctx.NextHop.ShouldBe("1");

        ctx.AddEgressEvent(1, "route-1", MsgTraceContext.KindRouter);

        var egress = ctx.Event.Events[1].ShouldBeOfType<MsgTraceEgress>();
        egress.Hop.ShouldBe("1");

        // NextHop should be cleared after adding egress
        ctx.NextHop.ShouldBe("");
    }

    // --- SendEvent (non-JetStream) ---

    [Fact]
    public void SendEvent_publishes_immediately_without_jetstream()
    {
        var ctx = CreateSimpleContext();
        string? publishedDest = null;
        ctx.PublishCallback = (dest, reply, body) => { publishedDest = dest; };

        ctx.SendEvent();
        publishedDest.ShouldBe("trace.dest");
    }

    // --- JSON serialization ---

    [Fact]
    public void MsgTraceEvent_serializes_to_valid_json()
    {
        var ctx = CreateSimpleContext();
        ctx.Event.Server = new EventServerInfo { Name = "srv", Id = "SRV1" };
        ctx.AddSubjectMappingEvent("mapped.subject");
        ctx.AddEgressEvent(99, "subscriber", MsgTraceContext.KindClient, "test.>", "q1");
        ctx.AddStreamExportEvent("exportAcc", "export.subject");

        var json = JsonSerializer.Serialize(ctx.Event);
        var doc = JsonDocument.Parse(json);
        var root = doc.RootElement;

        root.GetProperty("server").GetProperty("name").GetString().ShouldBe("srv");
        root.GetProperty("request").GetProperty("msgsize").GetInt32().ShouldBe(64);
        root.GetProperty("events").GetArrayLength().ShouldBe(4);

        var events = root.GetProperty("events");
        events[0].GetProperty("type").GetString().ShouldBe(MsgTraceTypes.Ingress);
        events[1].GetProperty("type").GetString().ShouldBe(MsgTraceTypes.SubjectMapping);
        events[2].GetProperty("type").GetString().ShouldBe(MsgTraceTypes.Egress);
        events[3].GetProperty("type").GetString().ShouldBe(MsgTraceTypes.StreamExport);
    }

    [Fact]
    public void MsgTraceIngress_json_omits_null_error()
    {
        var ingress = new MsgTraceIngress
        {
            Type = MsgTraceTypes.Ingress,
            Cid = 1,
            Account = "$G",
            Subject = "test",
        };

        var json = JsonSerializer.Serialize<MsgTraceEntry>(ingress);
        var doc = JsonDocument.Parse(json);
        doc.RootElement.TryGetProperty("error", out _).ShouldBeFalse();
    }

    [Fact]
    public void MsgTraceEgress_json_omits_null_optional_fields()
    {
        var egress = new MsgTraceEgress
        {
            Type = MsgTraceTypes.Egress,
            Kind = MsgTraceContext.KindRouter,
            Cid = 5,
        };

        var json = JsonSerializer.Serialize<MsgTraceEntry>(egress);
        var doc = JsonDocument.Parse(json);
        var root = doc.RootElement;

        root.TryGetProperty("hop", out _).ShouldBeFalse();
        root.TryGetProperty("acc", out _).ShouldBeFalse();
        root.TryGetProperty("sub", out _).ShouldBeFalse();
        root.TryGetProperty("queue", out _).ShouldBeFalse();
        root.TryGetProperty("error", out _).ShouldBeFalse();
    }

    [Fact]
    public void Full_trace_event_with_all_event_types_serializes_correctly()
    {
        var ctx = CreateSimpleContext();
        ctx.Event.Server = new EventServerInfo { Name = "test-srv", Id = "ABC123" };
        ctx.AddSubjectMappingEvent("mapped");
        ctx.AddServiceImportEvent("importAcc", "from.sub", "to.sub");
        ctx.AddStreamExportEvent("exportAcc", "export.sub");
        ctx.AddJetStreamEvent("ORDERS");
        ctx.UpdateJetStreamEvent("orders.new", false);
        ctx.AddEgressEvent(100, "sub-1", MsgTraceContext.KindClient, "orders.>", "workers");
        ctx.AddEgressEvent(200, "route-east", MsgTraceContext.KindRouter, error: MsgTraceErrors.NoSupport);

        var json = JsonSerializer.Serialize(ctx.Event);
        var doc = JsonDocument.Parse(json);
        var events = doc.RootElement.GetProperty("events");

        events.GetArrayLength().ShouldBe(7);
        events[0].GetProperty("type").GetString().ShouldBe("in");
        events[1].GetProperty("type").GetString().ShouldBe("sm");
        events[2].GetProperty("type").GetString().ShouldBe("si");
        events[3].GetProperty("type").GetString().ShouldBe("se");
        events[4].GetProperty("type").GetString().ShouldBe("js");
        events[5].GetProperty("type").GetString().ShouldBe("eg");
        events[6].GetProperty("type").GetString().ShouldBe("eg");
    }

    // --- Helper ---

    private static MsgTraceContext CreateSimpleContext(string destination = "trace.dest", string accountName = "$G")
    {
        var headers = BuildHeaders(
            (MsgTraceHeaders.TraceDest, destination));

        var ctx = MsgTraceContext.Create(
            headers,
            clientId: 1,
            clientName: "publisher",
            accountName: accountName,
            subject: "test.subject",
            msgSize: 64);

        ctx.ShouldNotBeNull();
        return ctx;
    }
}