// Copyright 2025 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/scheduler.go in the NATS server Go source.

using System.Buffers.Binary;
using ZB.MOM.NatsNet.Server;
using ZB.MOM.NatsNet.Server.Internal.DataStructures;

namespace ZB.MOM.NatsNet.Server.Internal;

/// <summary>
/// Error for when we try to decode a binary-encoded message schedule with an unknown version number.
/// Mirrors <c>ErrMsgScheduleInvalidVersion</c>.
/// </summary>
public static class MsgSchedulingErrors
{
    public static readonly Exception ErrMsgScheduleInvalidVersion =
        new InvalidOperationException("msg scheduling: encoded version not known");
}

/// <summary>
/// A single scheduled message entry.
/// Mirrors the unnamed struct in the <c>schedules</c> map in scheduler.go.
/// </summary>
internal sealed class MsgSchedule
{
    public ulong Seq;
    public long Ts;
}

/// <summary>
/// Tracks per-subject scheduled messages using a hash wheel for TTL management.
/// Mirrors <c>MsgScheduling</c> in server/scheduler.go.
/// Note: <c>getScheduledMessages</c> is deferred to session 08/19 (requires JetStream types).
/// </summary>
public sealed class MsgScheduling
{
    private const int HeaderLen = 17; // 1 magic + 2 × uint64

    private readonly Action _run;
    private readonly HashWheel _ttls;
    private Timer? _timer;
    // _running is set to true by the run callback when getScheduledMessages is active (session 08/19).
#pragma warning disable CS0649
    private bool _running;
#pragma warning restore CS0649
    private long _deadline;
    private readonly Dictionary<string, MsgSchedule> _schedules;
    private readonly Dictionary<ulong, string> _seqToSubj;
    private readonly HashSet<string> _inflight;

    /// <summary>
    /// Creates a new <see cref="MsgScheduling"/> with the given callback.
    /// Mirrors <c>newMsgScheduling</c>.
    /// </summary>
    public MsgScheduling(Action run)
    {
        _run = run;
        _ttls = HashWheel.NewHashWheel();
        _schedules = new Dictionary<string, MsgSchedule>();
        _seqToSubj = new Dictionary<ulong, string>();
        _inflight = new HashSet<string>();
    }

    /// <summary>
    /// Adds a schedule entry and resets the timer.
    /// Mirrors <c>MsgScheduling.add</c>.
    /// </summary>
    public void Add(ulong seq, string subj, long ts)
    {
        Init(seq, subj, ts);
        ResetTimer();
    }

    /// <summary>
    /// Inserts or updates the schedule entry for the given subject.
    /// Mirrors <c>MsgScheduling.init</c>.
    /// </summary>
    public void Init(ulong seq, string subj, long ts)
    {
        if (_schedules.TryGetValue(subj, out var sched))
        {
            _seqToSubj.Remove(sched.Seq);
            _ttls.Remove(sched.Seq, sched.Ts);
            _ttls.Add(seq, ts);
            sched.Ts = ts;
            sched.Seq = seq;
        }
        else
        {
            _ttls.Add(seq, ts);
            _schedules[subj] = new MsgSchedule { Seq = seq, Ts = ts };
        }
        _seqToSubj[seq] = subj;
        _inflight.Remove(subj);
    }

    /// <summary>
    /// Updates the timestamp for an existing schedule without changing the sequence.
    /// Mirrors <c>MsgScheduling.update</c>.
    /// </summary>
    public void Update(string subj, long ts)
    {
        if (!_schedules.TryGetValue(subj, out var sched)) return;
        _ttls.Remove(sched.Seq, sched.Ts);
        _ttls.Add(sched.Seq, ts);
        sched.Ts = ts;
        _inflight.Remove(subj);
        ResetTimer();
    }

    /// <summary>
    /// Marks a subject as in-flight (being processed).
    /// Mirrors <c>MsgScheduling.markInflight</c>.
    /// </summary>
    public void MarkInflight(string subj)
    {
        if (_schedules.ContainsKey(subj))
            _inflight.Add(subj);
    }

    /// <summary>
    /// Returns true if the subject is currently in-flight.
    /// Mirrors <c>MsgScheduling.isInflight</c>.
    /// </summary>
    public bool IsInflight(string subj) => _inflight.Contains(subj);

    /// <summary>
    /// Removes the schedule entry for the given sequence number.
    /// Mirrors <c>MsgScheduling.remove</c>.
    /// </summary>
    public void Remove(ulong seq)
    {
        if (!_seqToSubj.TryGetValue(seq, out var subj)) return;
        _seqToSubj.Remove(seq);
        _schedules.Remove(subj);
    }

    /// <summary>
    /// Removes the schedule entry for the given subject.
    /// Mirrors <c>MsgScheduling.removeSubject</c>.
    /// </summary>
    public void RemoveSubject(string subj)
    {
        if (!_schedules.TryGetValue(subj, out var sched)) return;
        _ttls.Remove(sched.Seq, sched.Ts);
        _schedules.Remove(subj);
        _seqToSubj.Remove(sched.Seq);
    }

    /// <summary>
    /// Clears all in-flight markers.
    /// Mirrors <c>MsgScheduling.clearInflight</c>.
    /// </summary>
    public void ClearInflight() => _inflight.Clear();

    /// <summary>
    /// Arms or resets the internal timer to fire at the next scheduled expiration.
    /// Mirrors <c>MsgScheduling.resetTimer</c>.
    /// </summary>
    public void ResetTimer()
    {
        if (_running) return;

        var next = _ttls.GetNextExpiration(long.MaxValue);
        if (next == long.MaxValue)
        {
            ClearTimer(ref _timer);
            return;
        }

        // Convert nanosecond timestamp to DateTime (1 tick = 100 ns).
        var nextTicks = DateTime.UnixEpoch.Ticks + next / 100L;
        var nextUtc = new DateTime(nextTicks, DateTimeKind.Utc);
        var fireIn = nextUtc - DateTime.UtcNow;

        // Clamp minimum interval.
        if (fireIn < TimeSpan.FromMilliseconds(250))
            fireIn = TimeSpan.FromMilliseconds(250);

        var deadline = DateTime.UtcNow.Ticks + fireIn.Ticks;
        if (_deadline > 0 && deadline > _deadline) return;

        _deadline = deadline;
        if (_timer != null)
            _timer.Change(fireIn, Timeout.InfiniteTimeSpan);
        else
            _timer = new Timer(_ => _run(), null, fireIn, Timeout.InfiniteTimeSpan);
    }

    /// <summary>
    /// Processes expired schedule entries and returns the set of messages to be delivered.
    /// Each message is retrieved from storage, headers are cleaned and augmented, and the
    /// subject is replaced with the schedule target. Messages are returned sorted by
    /// sequence number.
    /// Mirrors Go <c>MsgScheduling.getScheduledMessages</c> in server/scheduler.go.
    /// </summary>
    /// <param name="loadMsg">
    /// Callback that loads a stored message by sequence number.
    /// The <c>StoreMsg</c> reuse buffer may be passed; returns <c>null</c> if not found.
    /// </param>
    /// <param name="loadLast">
    /// Callback that loads the last stored message for a given subject.
    /// Returns <c>null</c> if not found.
    /// </param>
    public List<InMsg> GetScheduledMessages(
        Func<ulong, ZB.MOM.NatsNet.Server.StoreMsg, ZB.MOM.NatsNet.Server.StoreMsg?> loadMsg,
        Func<string, ZB.MOM.NatsNet.Server.StoreMsg, ZB.MOM.NatsNet.Server.StoreMsg?> loadLast)
    {
        var smv  = new ZB.MOM.NatsNet.Server.StoreMsg();
        List<InMsg>? msgs = null;

        _ttls.ExpireTasks((seq, ts) =>
        {
            var sm = loadMsg(seq, smv);
            if (sm != null)
            {
                // Already in-flight for this subject — skip.
                var subj = sm.Subject;
                if (IsInflight(subj))
                    return false;

                // Validate the schedule pattern header.
                var patternBytes = NatsMessageHeaders.GetHeader(
                    NatsHeaderConstants.JsSchedulePattern, sm.Hdr);
                if (patternBytes == null || patternBytes.Length == 0)
                {
                    Remove(seq);
                    return true;
                }
                var pattern = System.Text.Encoding.ASCII.GetString(patternBytes);

                var (next, repeat, ok) = ParseMsgSchedule(pattern, ts);
                if (!ok)
                {
                    Remove(seq);
                    return true;
                }

                var (ttl, ttlOk) = ZB.MOM.NatsNet.Server.NatsStream.GetMessageScheduleTTL(sm.Hdr);
                if (!ttlOk)
                {
                    Remove(seq);
                    return true;
                }

                var target = ZB.MOM.NatsNet.Server.NatsStream.GetMessageScheduleTarget(sm.Hdr);
                if (string.IsNullOrEmpty(target))
                {
                    Remove(seq);
                    return true;
                }

                var source = ZB.MOM.NatsNet.Server.NatsStream.GetMessageScheduleSource(sm.Hdr);
                if (!string.IsNullOrEmpty(source))
                {
                    sm = loadLast(source, smv);
                    if (sm == null)
                    {
                        Remove(seq);
                        return true;
                    }
                }

                // Copy headers and body — message lives beyond this callback.
                var hdr = sm.Hdr.Length > 0 ? (byte[])sm.Hdr.Clone() : [];
                var msg = sm.Msg.Length > 0 ? (byte[])sm.Msg.Clone() : [];

                // Strip schedule-specific headers.
                hdr = NatsMessageHeaders.RemoveHeaderIfPresent(hdr, NatsHeaderConstants.JsSchedulePattern) ?? [];
                hdr = NatsMessageHeaders.RemoveHeaderIfPrefixPresent(hdr, "Nats-Schedule-") ?? [];
                hdr = NatsMessageHeaders.RemoveHeaderIfPrefixPresent(hdr, "Nats-Expected-") ?? [];
                hdr = NatsMessageHeaders.RemoveHeaderIfPresent(hdr, NatsHeaderConstants.JsMsgId) ?? [];
                hdr = NatsMessageHeaders.RemoveHeaderIfPresent(hdr, NatsHeaderConstants.JsMessageTtl) ?? [];
                hdr = NatsMessageHeaders.RemoveHeaderIfPresent(hdr, NatsHeaderConstants.JsMsgRollup) ?? [];

                // Add scheduler-specific headers.
                hdr = NatsMessageHeaders.GenHeader(hdr, NatsHeaderConstants.JsScheduler, subj);
                if (!repeat)
                {
                    hdr = NatsMessageHeaders.GenHeader(hdr, NatsHeaderConstants.JsScheduleNext,
                        NatsHeaderConstants.JsScheduleNextPurge);
                }
                else
                {
                    hdr = NatsMessageHeaders.GenHeader(hdr, NatsHeaderConstants.JsScheduleNext,
                        next.ToString("yyyy-MM-ddTHH:mm:ssK"));
                }
                if (!string.IsNullOrEmpty(ttl))
                    hdr = NatsMessageHeaders.GenHeader(hdr, NatsHeaderConstants.JsMessageTtl, ttl);

                msgs ??= [];
                msgs.Add(new InMsg { Seq = seq, Subject = target, Hdr = hdr, Msg = msg });
                MarkInflight(subj);
                return false;
            }

            Remove(seq);
            return true;
        });

        if (msgs == null)
            return [];

        // THW is unordered — sort by sequence before returning.
        msgs.Sort((a, b) => a.Seq.CompareTo(b.Seq));
        return msgs;
    }

    /// <summary>
    /// Encodes the current schedule state to a binary snapshot.
    /// Mirrors <c>MsgScheduling.encode</c>.
    /// </summary>
    public byte[] Encode(ulong highSeq)
    {
        var count = (ulong)_schedules.Count;
        var buf = new List<byte>(HeaderLen + (int)(count * 20));

        buf.Add(1); // magic version
        AppendUInt64(buf, count);
        AppendUInt64(buf, highSeq);

        foreach (var (subj, sched) in _schedules)
        {
            var slen = (ushort)Math.Min((ulong)subj.Length, ushort.MaxValue);
            AppendUInt16(buf, slen);
            buf.AddRange(System.Text.Encoding.Latin1.GetBytes(subj[..slen]));
            AppendVarint(buf, sched.Ts);
            AppendUvarint(buf, sched.Seq);
        }

        return [.. buf];
    }

    /// <summary>
    /// Decodes a binary snapshot into the current schedule.
    /// Returns the high-sequence stamp or throws on error.
    /// Mirrors <c>MsgScheduling.decode</c>.
    /// </summary>
    public (ulong highSeq, Exception? err) Decode(byte[] b)
    {
        if (b.Length < HeaderLen)
            return (0, new System.IO.EndOfStreamException("short buffer"));

        if (b[0] != 1)
            return (0, MsgSchedulingErrors.ErrMsgScheduleInvalidVersion);

        var count = BinaryPrimitives.ReadUInt64LittleEndian(b.AsSpan(1));
        var stamp = BinaryPrimitives.ReadUInt64LittleEndian(b.AsSpan(9));
        var offset = HeaderLen;

        for (ulong i = 0; i < count; i++)
        {
            if (offset + 2 > b.Length)
                return (0, new System.IO.EndOfStreamException("unexpected EOF"));

            var sl = BinaryPrimitives.ReadUInt16LittleEndian(b.AsSpan(offset));
            offset += 2;

            if (offset + sl > b.Length)
                return (0, new System.IO.EndOfStreamException("unexpected EOF"));

            var subj = System.Text.Encoding.Latin1.GetString(b, offset, sl);
            offset += sl;

            var (ts, tn) = ReadVarint(b, offset);
            if (tn < 0) return (0, new System.IO.EndOfStreamException("unexpected EOF"));
            offset += tn;

            var (seq, vn) = ReadUvarint(b, offset);
            if (vn < 0) return (0, new System.IO.EndOfStreamException("unexpected EOF"));
            offset += vn;

            Init(seq, subj, ts);
        }

        return (stamp, null);
    }

    /// <summary>
    /// Parses a message schedule pattern and returns the next fire time,
    /// whether it repeats, and whether the pattern was valid.
    /// Mirrors <c>parseMsgSchedule</c>.
    /// </summary>
    public static (DateTime next, bool repeat, bool ok) ParseMsgSchedule(string pattern, long ts)
    {
        if (pattern == string.Empty)
            return (default, false, true);

        if (pattern.StartsWith("@at ", StringComparison.Ordinal))
        {
            if (DateTime.TryParseExact(
                    pattern[4..],
                    "yyyy-MM-ddTHH:mm:ssK",
                    System.Globalization.CultureInfo.InvariantCulture,
                    System.Globalization.DateTimeStyles.AdjustToUniversal,
                    out var t))
                return (t, false, true);
            return (default, false, false);
        }

        if (pattern.StartsWith("@every ", StringComparison.Ordinal))
        {
            if (!TryParseDuration(pattern[7..], out var dur))
                return (default, false, false);

            if (dur.TotalSeconds < 1)
                return (default, false, false);

            // Advance past a stale next tick the same way Go does.
            var prev = DateTimeOffset.FromUnixTimeMilliseconds(ts / 1_000_000).UtcDateTime;
            var next = RoundToSecond(prev).Add(dur);
            var now = RoundToSecond(DateTime.UtcNow);
            if (next < now)
                next = now.Add(dur);

            return (next, true, true);
        }

        return (default, false, false);
    }

    // -------------------------------------------------------------------------
    // Private helpers
    // -------------------------------------------------------------------------

    private static void ClearTimer(ref Timer? timer)
    {
        var t = timer;
        if (t == null) return;
        t.Dispose();
        timer = null;
    }

    private static DateTime RoundToSecond(DateTime dt) =>
        new(dt.Year, dt.Month, dt.Day, dt.Hour, dt.Minute, dt.Second, DateTimeKind.Utc);

    // Naive duration parser for strings like "1s", "500ms", "2m", "1h30m".
    private static bool TryParseDuration(string s, out TimeSpan result)
    {
        result = default;
        if (s.EndsWith("ms", StringComparison.Ordinal) &&
            double.TryParse(s[..^2], out var ms))
        {
            result = TimeSpan.FromMilliseconds(ms);
            return true;
        }
        if (s.EndsWith('s') && double.TryParse(s[..^1], out var sec))
        {
            result = TimeSpan.FromSeconds(sec);
            return true;
        }
        if (s.EndsWith('m') && double.TryParse(s[..^1], out var min))
        {
            result = TimeSpan.FromMinutes(min);
            return true;
        }
        if (s.EndsWith('h') && double.TryParse(s[..^1], out var hr))
        {
            result = TimeSpan.FromHours(hr);
            return true;
        }
        // Try .NET TimeSpan.Parse as a fallback.
        return TimeSpan.TryParse(s, out result);
    }

    // -------------------------------------------------------------------------
    // Binary encoding helpers (mirrors encoding/binary in Go)
    // -------------------------------------------------------------------------

    private static void AppendUInt64(List<byte> buf, ulong v)
    {
        Span<byte> tmp = stackalloc byte[8];
        BinaryPrimitives.WriteUInt64LittleEndian(tmp, v);
        buf.AddRange(tmp.ToArray());
    }

    private static void AppendUInt16(List<byte> buf, ushort v)
    {
        Span<byte> tmp = stackalloc byte[2];
        BinaryPrimitives.WriteUInt16LittleEndian(tmp, v);
        buf.AddRange(tmp.ToArray());
    }

    /// <summary>Appends a zigzag-encoded signed varint (mirrors binary.AppendVarint).</summary>
    private static void AppendVarint(List<byte> buf, long x)
    {
        var ux = (ulong)(x << 1);
        if (x < 0) ux = ~ux;
        AppendUvarint(buf, ux);
    }

    /// <summary>Appends an unsigned varint (mirrors binary.AppendUvarint).</summary>
    private static void AppendUvarint(List<byte> buf, ulong x)
    {
        while (x >= 0x80)
        {
            buf.Add((byte)(x | 0x80));
            x >>= 7;
        }
        buf.Add((byte)x);
    }

    /// <summary>
    /// Reads a zigzag signed varint from <paramref name="b"/> starting at <paramref name="offset"/>.
    /// Returns (value, bytesRead); bytesRead is negative on overflow.
    /// </summary>
    private static (long value, int n) ReadVarint(byte[] b, int offset)
    {
        var (ux, n) = ReadUvarint(b, offset);
        var x = (long)(ux >> 1);
        if ((ux & 1) != 0) x = ~x;
        return (x, n);
    }

    /// <summary>
    /// Reads an unsigned varint from <paramref name="b"/> starting at <paramref name="offset"/>.
    /// Returns (value, bytesRead); bytesRead is negative on overflow.
    /// </summary>
    private static (ulong value, int n) ReadUvarint(byte[] b, int offset)
    {
        ulong x = 0;
        var s = 0;
        for (var i = offset; i < b.Length; i++)
        {
            var by = b[i];
            if (i - offset == 10) return (0, -(i - offset + 1)); // overflow
            if (by < 0x80)
            {
                if (i - offset == 9 && by > 1) return (0, -(i - offset + 1));
                return (x | ((ulong)by << s), i - offset + 1);
            }
            x |= (ulong)(by & 0x7F) << s;
            s += 7;
        }
        return (0, 0); // short buffer
    }
}