// Copyright 2012-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0

using System.Buffers.Binary;
using System.Security.Cryptography;

namespace ZB.MOM.NatsNet.Server;

internal sealed partial class Raft
{
    private const string SnapshotFilePrefix = "snap";
    private const string SnapshotFileSuffix = ".bin";
    private const int SnapshotHeaderLength = 20;
    private const int SnapshotChecksumLength = 8;

    internal byte[] EncodeSnapshot(Snapshot? snap)
    {
        if (snap is null)
        {
            return [];
        }

        var peerStateLength = snap.PeerState.Length;
        var dataLength = snap.Data.Length;
        var payloadLength = SnapshotHeaderLength + peerStateLength + dataLength;
        var buffer = new byte[payloadLength + SnapshotChecksumLength];

        BinaryPrimitives.WriteUInt64LittleEndian(buffer.AsSpan(0, 8), snap.LastTerm);
        BinaryPrimitives.WriteUInt64LittleEndian(buffer.AsSpan(8, 8), snap.LastIndex);
        BinaryPrimitives.WriteUInt32LittleEndian(buffer.AsSpan(16, 4), (uint)peerStateLength);

        var writeOffset = SnapshotHeaderLength;
        snap.PeerState.CopyTo(buffer, writeOffset);
        writeOffset += peerStateLength;
        snap.Data.CopyTo(buffer, writeOffset);

        var checksum = SHA256.HashData(buffer.AsSpan(0, payloadLength));
        checksum.AsSpan(0, SnapshotChecksumLength).CopyTo(buffer.AsSpan(payloadLength, SnapshotChecksumLength));
        return buffer;
    }

    internal Exception? InstallSnapshotInternal(Snapshot snap)
    {
        ArgumentNullException.ThrowIfNull(snap);

        try
        {
            var snapshotsDirectory = GetSnapshotsDirectory();
            Directory.CreateDirectory(snapshotsDirectory);

            var newSnapshotFile = Path.Combine(snapshotsDirectory, FormatSnapshotFileName(snap.LastTerm, snap.LastIndex));
            File.WriteAllBytes(newSnapshotFile, EncodeSnapshot(snap));

            if (!string.IsNullOrWhiteSpace(SnapFile) &&
                !string.Equals(SnapFile, newSnapshotFile, StringComparison.Ordinal) &&
                File.Exists(SnapFile))
            {
                File.Delete(SnapFile);
            }

            SnapFile = newSnapshotFile;
            PApplied = snap.LastIndex;
            PIndex = snap.LastIndex;
            PTerm = snap.LastTerm;
            Commit = Math.Max(Commit, snap.LastIndex);
            WalBytes = (ulong)(new FileInfo(newSnapshotFile).Length);
            return null;
        }
        catch (Exception ex)
        {
            return ex;
        }
        finally
        {
            Snapshotting = false;
        }
    }

    internal Checkpoint? CreateSnapshotCheckpointLocked(bool force)
    {
        if (State() == RaftState.Closed)
        {
            return null;
        }

        if (Snapshotting)
        {
            return null;
        }

        if (!force && Progress_ is { Count: > 0 })
        {
            return null;
        }

        if (Applied_ == 0)
        {
            return null;
        }

        Snapshotting = true;
        var snapshotFile = Path.Combine(GetSnapshotsDirectory(), FormatSnapshotFileName(Term_, Applied_));
        return new Checkpoint
        {
            Node = this,
            Term = Term_,
            Applied = Applied_,
            PApplied = PApplied,
            SnapFile = snapshotFile,
            PeerState = [.. Wps],
        };
    }

    internal (ulong Term, ulong Index, Exception? Error) TermAndIndexFromSnapFile(string snapFileName)
    {
        if (string.IsNullOrWhiteSpace(snapFileName))
        {
            return (0, 0, new InvalidOperationException("bad snapshot file name"));
        }

        var fileName = Path.GetFileNameWithoutExtension(snapFileName);
        var segments = fileName.Split('-', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);
        if (segments.Length != 3 || !string.Equals(segments[0], SnapshotFilePrefix, StringComparison.Ordinal))
        {
            return (0, 0, new InvalidOperationException("bad snapshot file name"));
        }

        if (!ulong.TryParse(segments[1], out var term) || !ulong.TryParse(segments[2], out var index))
        {
            return (0, 0, new InvalidOperationException("bad snapshot file name"));
        }

        return (term, index, null);
    }

    internal Exception? SetupLastSnapshot()
    {
        var snapshotDirectory = GetSnapshotsDirectory();
        if (!Directory.Exists(snapshotDirectory))
        {
            return new InvalidOperationException("no snapshot available");
        }

        var snapshotFiles = Directory.GetFiles(snapshotDirectory, $"*{SnapshotFileSuffix}");
        if (snapshotFiles.Length == 0)
        {
            return new InvalidOperationException("no snapshot available");
        }

        ulong latestTerm = 0;
        ulong latestIndex = 0;
        string? latestSnapshot = null;
        foreach (var candidate in snapshotFiles)
        {
            var (term, index, parseError) = TermAndIndexFromSnapFile(candidate);
            if (parseError is not null)
            {
                File.Delete(candidate);
                continue;
            }

            if (term > latestTerm || (term == latestTerm && index > latestIndex))
            {
                latestTerm = term;
                latestIndex = index;
                latestSnapshot = candidate;
            }
        }

        if (latestSnapshot is null)
        {
            return null;
        }

        SnapFile = latestSnapshot;
        var (snapshot, loadError) = LoadLastSnapshot();
        if (loadError is not null || snapshot is null)
        {
            return loadError;
        }

        PIndex = snapshot.LastIndex;
        PTerm = snapshot.LastTerm;
        Commit = snapshot.LastIndex;
        PApplied = snapshot.LastIndex;
        Wps = [.. snapshot.PeerState];

        foreach (var oldFile in snapshotFiles)
        {
            if (!string.Equals(oldFile, latestSnapshot, StringComparison.Ordinal))
            {
                File.Delete(oldFile);
            }
        }

        return null;
    }

    internal (Snapshot? Snapshot, Exception? Error) LoadLastSnapshot()
    {
        if (string.IsNullOrWhiteSpace(SnapFile))
        {
            return (null, new InvalidOperationException("no snapshot available"));
        }

        try
        {
            var buffer = File.ReadAllBytes(SnapFile);
            if (buffer.Length < SnapshotHeaderLength + SnapshotChecksumLength)
            {
                return (null, new InvalidOperationException("snapshot corrupt"));
            }

            var payloadLength = buffer.Length - SnapshotChecksumLength;
            var expectedChecksum = buffer.AsSpan(payloadLength, SnapshotChecksumLength);
            var computedChecksum = SHA256.HashData(buffer.AsSpan(0, payloadLength));
            if (!expectedChecksum.SequenceEqual(computedChecksum.AsSpan(0, SnapshotChecksumLength)))
            {
                return (null, new InvalidOperationException("snapshot corrupt"));
            }

            var term = BinaryPrimitives.ReadUInt64LittleEndian(buffer.AsSpan(0, 8));
            var index = BinaryPrimitives.ReadUInt64LittleEndian(buffer.AsSpan(8, 8));
            var peerStateLength = (int)BinaryPrimitives.ReadUInt32LittleEndian(buffer.AsSpan(16, 4));
            if (SnapshotHeaderLength + peerStateLength > payloadLength)
            {
                return (null, new InvalidOperationException("snapshot corrupt"));
            }

            var peerState = buffer.AsSpan(SnapshotHeaderLength, peerStateLength).ToArray();
            var data = buffer.AsSpan(SnapshotHeaderLength + peerStateLength, payloadLength - SnapshotHeaderLength - peerStateLength).ToArray();
            if (index == 0)
            {
                File.Delete(SnapFile);
                SnapFile = string.Empty;
                return (null, null);
            }

            return (new Snapshot
            {
                LastTerm = term,
                LastIndex = index,
                PeerState = peerState,
                Data = data,
            }, null);
        }
        catch (Exception ex)
        {
            return (null, ex);
        }
    }

    internal void StepdownLocked(string newLeader)
    {
        StateValue = (int)RaftState.Follower;
        LeaderId = newLeader;
        Lsut = DateTime.UtcNow;
    }

    internal bool IsCatchingUp() => Catchup is not null;

    internal bool IsCurrent(bool includeForwardProgress = false)
    {
        if (State() == RaftState.Closed || Commit == 0 || Catchup is not null)
        {
            return false;
        }

        if (Paused && HCommit > Commit)
        {
            return false;
        }

        if (Commit == Applied_)
        {
            HcBehind = false;
            return true;
        }

        if (!includeForwardProgress)
        {
            return false;
        }

        var startDelta = Commit > Applied_ ? Commit - Applied_ : 0;
        return startDelta <= 1;
    }

    internal string SelectNextLeader()
    {
        var nextLeader = string.Empty;
        ulong highestIndex = 0;
        foreach (var (peer, peerState) in Peers_)
        {
            if (string.Equals(peer, Id, StringComparison.Ordinal) || peerState.Li <= highestIndex)
            {
                continue;
            }

            nextLeader = peer;
            highestIndex = peerState.Li;
        }

        return nextLeader;
    }

    internal TimeSpan RandCampaignTimeout()
    {
        var min = 150;
        var max = 300;
        var delta = Random.Shared.Next(min, max);
        return TimeSpan.FromMilliseconds(delta);
    }

    internal Exception? CampaignInternal(TimeSpan electionTimeout)
    {
        _ = electionTimeout;
        if (State() == RaftState.Leader)
        {
            return new InvalidOperationException("already leader");
        }

        Campaign();
        return null;
    }

    internal Exception? XferCampaign()
    {
        if (State() == RaftState.Leader)
        {
            Lxfer = false;
            return new InvalidOperationException("already leader");
        }

        Lxfer = true;
        Campaign();
        return null;
    }

    internal void UpdateKnownPeersLocked(IReadOnlyList<string> knownPeers)
    {
        ProposeKnownPeers(knownPeers);
    }

    private string GetSnapshotsDirectory()
    {
        if (string.IsNullOrWhiteSpace(StoreDir))
        {
            StoreDir = Path.Combine(Path.GetTempPath(), "natsnet-raft");
        }

        return Path.Combine(StoreDir, "snapshots");
    }

    private static string FormatSnapshotFileName(ulong term, ulong index) =>
        $"{SnapshotFilePrefix}-{term:D20}-{index:D20}{SnapshotFileSuffix}";
}
