// Copyright 2024-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Adapted from server/msgtrace.go in the NATS server Go source.

using System.Text.Json;
using System.Text.Json.Serialization;

namespace ZB.MOM.NatsNet.Server;

// ============================================================================
// Message-trace header name constants
// Mirrors Go const block at top of server/msgtrace.go.
// ============================================================================

/// <summary>
/// NATS message-trace header names and special sentinel values.
/// Mirrors Go const block in server/msgtrace.go.
/// </summary>
public static class MsgTraceHeaders
{
    /// <summary>Header that carries the trace destination subject. Mirrors Go <c>MsgTraceDest</c>.</summary>
    public const string MsgTraceDest = "Nats-Trace-Dest";

    /// <summary>
    /// Sentinel value placed in the trace-dest header to disable tracing
    /// (must be an invalid NATS subject). Mirrors Go <c>MsgTraceDestDisabled</c>.
    /// </summary>
    public const string MsgTraceDestDisabled = "trace disabled";

    /// <summary>Header used for hop-count tracking across servers. Mirrors Go <c>MsgTraceHop</c>.</summary>
    public const string MsgTraceHop = "Nats-Trace-Hop";

    /// <summary>Header that carries the originating account name. Mirrors Go <c>MsgTraceOriginAccount</c>.</summary>
    public const string MsgTraceOriginAccount = "Nats-Trace-Origin-Account";

    /// <summary>
    /// When set to a truthy value, the message is consumed only for tracing
    /// and not delivered to subscribers. Mirrors Go <c>MsgTraceOnly</c>.
    /// </summary>
    public const string MsgTraceOnly = "Nats-Trace-Only";

    /// <summary>
    /// W3C trace-context parent header. NATS no longer lower-cases this but
    /// accepts it in any case. Mirrors Go <c>traceParentHdr</c> (internal).
    /// </summary>
    public const string TraceParentHdr = "traceparent";
}

// ============================================================================
// MsgTraceType — discriminator string for polymorphic trace event lists
// Mirrors Go <c>MsgTraceType string</c> in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Discriminator string identifying the concrete type of a trace event
/// within a <see cref="MsgTraceEvents"/> list.
/// Mirrors Go <c>MsgTraceType string</c> and its constants in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceType
{
    private readonly string _value;
    private MsgTraceType(string value) => _value = value;

    /// <inheritdoc/>
    public override string ToString() => _value;

    public static implicit operator MsgTraceType(string value) => new(value);
    public static implicit operator string(MsgTraceType t) => t._value;

    public override bool Equals(object? obj) =>
        obj is MsgTraceType other && _value == other._value;

    public override int GetHashCode() => _value.GetHashCode();

    // ---- Well-known type constants (mirror Go const block) ----

    /// <summary>Ingress event. Mirrors Go <c>MsgTraceIngressType = "in"</c>.</summary>
    public static readonly MsgTraceType Ingress = new("in");

    /// <summary>Subject-mapping event. Mirrors Go <c>MsgTraceSubjectMappingType = "sm"</c>.</summary>
    public static readonly MsgTraceType SubjectMapping = new("sm");

    /// <summary>Stream-export event. Mirrors Go <c>MsgTraceStreamExportType = "se"</c>.</summary>
    public static readonly MsgTraceType StreamExport = new("se");

    /// <summary>Service-import event. Mirrors Go <c>MsgTraceServiceImportType = "si"</c>.</summary>
    public static readonly MsgTraceType ServiceImport = new("si");

    /// <summary>JetStream storage event. Mirrors Go <c>MsgTraceJetStreamType = "js"</c>.</summary>
    public static readonly MsgTraceType JetStream = new("js");

    /// <summary>Egress (delivery) event. Mirrors Go <c>MsgTraceEgressType = "eg"</c>.</summary>
    public static readonly MsgTraceType Egress = new("eg");
}

// ============================================================================
// IMsgTrace — interface for polymorphic trace events
// Mirrors Go <c>MsgTrace interface</c> in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Marker interface implemented by all concrete message-trace event types.
/// Enables polymorphic handling of the <see cref="MsgTraceEvents"/> list.
/// Mirrors Go <c>MsgTrace interface</c> in server/msgtrace.go.
/// </summary>
public interface IMsgTrace
{
    /// <summary>Returns the discriminator type string for this trace event.</summary>
    string Typ();
}

// ============================================================================
// MsgTraceBase — shared fields present in every trace event
// Mirrors Go <c>MsgTraceBase</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Common base fields shared by all concrete message-trace event types.
/// Mirrors Go <c>MsgTraceBase</c> struct in server/msgtrace.go.
/// </summary>
public class MsgTraceBase : IMsgTrace
{
    [JsonPropertyName("type")]
    public string Type { get; set; } = string.Empty;

    [JsonPropertyName("ts")]
    public DateTime Timestamp { get; set; }

    /// <inheritdoc/>
    public virtual string Typ() => Type;
}

// ============================================================================
// MsgTraceIngress — client / route / gateway / leaf connection ingress event
// Mirrors Go <c>MsgTraceIngress</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Records the point at which a message was received by the server from a
/// client, route, gateway, or leaf connection.
/// Mirrors Go <c>MsgTraceIngress</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceIngress : MsgTraceBase
{
    [JsonPropertyName("kind")]
    public int Kind { get; set; }

    [JsonPropertyName("cid")]
    public ulong Cid { get; set; }

    [JsonPropertyName("name")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Name { get; set; }

    [JsonPropertyName("acc")]
    public string Account { get; set; } = string.Empty;

    [JsonPropertyName("subj")]
    public string Subject { get; set; } = string.Empty;

    [JsonPropertyName("error")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Error { get; set; }
}

// ============================================================================
// MsgTraceSubjectMapping — subject-mapping rewrite event
// Mirrors Go <c>MsgTraceSubjectMapping</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Records a subject-mapping rewrite applied to an in-flight message.
/// Mirrors Go <c>MsgTraceSubjectMapping</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceSubjectMapping : MsgTraceBase
{
    [JsonPropertyName("to")]
    public string MappedTo { get; set; } = string.Empty;
}

// ============================================================================
// MsgTraceStreamExport — stream export / cross-account delivery event
// Mirrors Go <c>MsgTraceStreamExport</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Records delivery of a message to a stream-export destination account.
/// Mirrors Go <c>MsgTraceStreamExport</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceStreamExport : MsgTraceBase
{
    [JsonPropertyName("acc")]
    public string Account { get; set; } = string.Empty;

    [JsonPropertyName("to")]
    public string To { get; set; } = string.Empty;
}

// ============================================================================
// MsgTraceServiceImport — service import routing event
// Mirrors Go <c>MsgTraceServiceImport</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Records routing of a message via a service-import from one account to
/// another.
/// Mirrors Go <c>MsgTraceServiceImport</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceServiceImport : MsgTraceBase
{
    [JsonPropertyName("acc")]
    public string Account { get; set; } = string.Empty;

    [JsonPropertyName("from")]
    public string From { get; set; } = string.Empty;

    [JsonPropertyName("to")]
    public string To { get; set; } = string.Empty;
}

// ============================================================================
// MsgTraceJetStream — JetStream storage event
// Mirrors Go <c>MsgTraceJetStream</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Records the attempt (and outcome) of storing or delivering a message
/// to a JetStream stream.
/// Mirrors Go <c>MsgTraceJetStream</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceJetStream : MsgTraceBase
{
    [JsonPropertyName("stream")]
    public string Stream { get; set; } = string.Empty;

    [JsonPropertyName("subject")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Subject { get; set; }

    [JsonPropertyName("nointerest")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
    public bool NoInterest { get; set; }

    [JsonPropertyName("error")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Error { get; set; }
}

// ============================================================================
// MsgTraceEgress — outbound delivery event
// Mirrors Go <c>MsgTraceEgress</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Records the outbound delivery of a message to a subscriber, route,
/// gateway, or leaf connection.
/// Mirrors Go <c>MsgTraceEgress</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceEgress : MsgTraceBase
{
    [JsonPropertyName("kind")]
    public int Kind { get; set; }

    [JsonPropertyName("cid")]
    public ulong Cid { get; set; }

    [JsonPropertyName("name")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Name { get; set; }

    [JsonPropertyName("hop")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Hop { get; set; }

    [JsonPropertyName("acc")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Account { get; set; }

    [JsonPropertyName("sub")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Subscription { get; set; }

    [JsonPropertyName("queue")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Queue { get; set; }

    [JsonPropertyName("error")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public string? Error { get; set; }

    /// <summary>
    /// Optional link to the <see cref="MsgTraceEvent"/> produced by the remote
    /// server that received this egress message (route/leaf/gateway hop).
    /// Not serialised. Mirrors Go <c>Link *MsgTraceEvent</c>.
    /// </summary>
    [JsonIgnore]
    public MsgTraceEvent? Link { get; set; }
}

// ============================================================================
// MsgTraceEvents — polymorphic list with custom JSON deserialiser
// Mirrors Go <c>MsgTraceEvents []MsgTrace</c> and its UnmarshalJSON in msgtrace.go.
// ============================================================================

/// <summary>
/// Custom JSON converter that deserialises a <c>MsgTraceEvents</c> JSON array
/// into the correct concrete <see cref="IMsgTrace"/> subtype, using the
/// <c>"type"</c> discriminator field.
/// Mirrors Go <c>MsgTraceEvents.UnmarshalJSON</c> in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceEventsConverter : JsonConverter<List<IMsgTrace>>
{
    private static readonly Dictionary<string, Func<JsonElement, IMsgTrace>> Factories = new()
    {
        ["in"] = e => e.Deserialize<MsgTraceIngress>()!,
        ["sm"] = e => e.Deserialize<MsgTraceSubjectMapping>()!,
        ["se"] = e => e.Deserialize<MsgTraceStreamExport>()!,
        ["si"] = e => e.Deserialize<MsgTraceServiceImport>()!,
        ["js"] = e => e.Deserialize<MsgTraceJetStream>()!,
        ["eg"] = e => e.Deserialize<MsgTraceEgress>()!,
    };

    public override List<IMsgTrace> Read(
        ref Utf8JsonReader reader,
        Type typeToConvert,
        JsonSerializerOptions options)
    {
        var result = new List<IMsgTrace>();
        using var doc = JsonDocument.ParseValue(ref reader);

        foreach (var element in doc.RootElement.EnumerateArray())
        {
            if (!element.TryGetProperty("type", out var typeProp))
                throw new JsonException("MsgTrace element missing 'type' field.");

            var typeStr = typeProp.GetString() ?? string.Empty;

            if (!Factories.TryGetValue(typeStr, out var factory))
                throw new JsonException($"Unknown MsgTrace type '{typeStr}'.");

            result.Add(factory(element));
        }

        return result;
    }

    public override void Write(
        Utf8JsonWriter writer,
        List<IMsgTrace> value,
        JsonSerializerOptions options)
    {
        writer.WriteStartArray();
        foreach (var item in value)
            JsonSerializer.Serialize(writer, item, item.GetType(), options);
        writer.WriteEndArray();
    }
}

// ============================================================================
// MsgTraceRequest — the original request metadata included in a trace event
// Mirrors Go <c>MsgTraceRequest</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// Captures the headers and size of the original message that triggered a
/// trace event.
/// Mirrors Go <c>MsgTraceRequest</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceRequest
{
    /// <summary>
    /// Original message headers, preserving header-name casing.
    /// Mirrors Go <c>Header map[string][]string</c> (not http.Header, so casing is preserved).
    /// </summary>
    [JsonPropertyName("header")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
    public Dictionary<string, List<string>>? Header { get; set; }

    [JsonPropertyName("msgsize")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
    public int MsgSize { get; set; }
}

// ============================================================================
// MsgTraceEvent — top-level trace event published to the trace destination
// Mirrors Go <c>MsgTraceEvent</c> struct in server/msgtrace.go.
// ============================================================================

/// <summary>
/// The top-level message-trace advisory published to the trace destination
/// subject. Contains server identity, the original request metadata, the
/// hop count, and the ordered list of trace events.
/// Mirrors Go <c>MsgTraceEvent</c> struct in server/msgtrace.go.
/// </summary>
public sealed class MsgTraceEvent
{
    [JsonPropertyName("server")]
    public ServerInfo Server { get; set; } = new();

    [JsonPropertyName("request")]
    public MsgTraceRequest Request { get; set; } = new();

    [JsonPropertyName("hops")]
    [JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingDefault)]
    public int Hops { get; set; }

    [JsonPropertyName("events")]
    [JsonConverter(typeof(MsgTraceEventsConverter))]
    public List<IMsgTrace> Events { get; set; } = [];

    // ---- Convenience accessors (mirrors Go helper methods on MsgTraceEvent) ----

    /// <summary>
    /// Returns the first event if it is a <see cref="MsgTraceIngress"/>, else null.
    /// Mirrors Go <c>MsgTraceEvent.Ingress()</c>.
    /// </summary>
    public MsgTraceIngress? Ingress() =>
        Events.Count > 0 ? Events[0] as MsgTraceIngress : null;

    /// <summary>
    /// Returns the first <see cref="MsgTraceSubjectMapping"/> in the event list, or null.
    /// Mirrors Go <c>MsgTraceEvent.SubjectMapping()</c>.
    /// </summary>
    public MsgTraceSubjectMapping? SubjectMapping() =>
        Events.OfType<MsgTraceSubjectMapping>().FirstOrDefault();

    /// <summary>
    /// Returns all <see cref="MsgTraceStreamExport"/> events.
    /// Mirrors Go <c>MsgTraceEvent.StreamExports()</c>.
    /// </summary>
    public IReadOnlyList<MsgTraceStreamExport> StreamExports() =>
        Events.OfType<MsgTraceStreamExport>().ToList();

    /// <summary>
    /// Returns all <see cref="MsgTraceServiceImport"/> events.
    /// Mirrors Go <c>MsgTraceEvent.ServiceImports()</c>.
    /// </summary>
    public IReadOnlyList<MsgTraceServiceImport> ServiceImports() =>
        Events.OfType<MsgTraceServiceImport>().ToList();

    /// <summary>
    /// Returns the first <see cref="MsgTraceJetStream"/> event, or null.
    /// Mirrors Go <c>MsgTraceEvent.JetStream()</c>.
    /// </summary>
    public MsgTraceJetStream? JetStream() =>
        Events.OfType<MsgTraceJetStream>().FirstOrDefault();

    /// <summary>
    /// Returns all <see cref="MsgTraceEgress"/> events.
    /// Mirrors Go <c>MsgTraceEvent.Egresses()</c>.
    /// </summary>
    public IReadOnlyList<MsgTraceEgress> Egresses() =>
        Events.OfType<MsgTraceEgress>().ToList();
}
