CBDDC/src/ZB.MOM.WW.CBDDC.Persistence/Surreal/SurrealPeerOplogConfirmationStore.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;
using SurrealDb.Net;
using SurrealDb.Net.Models;
using ZB.MOM.WW.CBDDC.Core;
using ZB.MOM.WW.CBDDC.Core.Network;

namespace ZB.MOM.WW.CBDDC.Persistence.Surreal;

public class SurrealPeerOplogConfirmationStore : PeerOplogConfirmationStore
{
    internal const string RegistrationSourceNodeId = "__peer_registration__";

    private readonly ILogger<SurrealPeerOplogConfirmationStore> _logger;
    private readonly ICBDDCSurrealSchemaInitializer _schemaInitializer;
    private readonly ISurrealDbClient _surrealClient;

    /// <summary>
    ///     Initializes a new instance of the <see cref="SurrealPeerOplogConfirmationStore"/> class.
    /// </summary>
    /// <param name="surrealEmbeddedClient">Embedded Surreal client wrapper.</param>
    /// <param name="schemaInitializer">Schema initializer.</param>
    /// <param name="logger">Optional logger.</param>
    public SurrealPeerOplogConfirmationStore(
        ICBDDCSurrealEmbeddedClient surrealEmbeddedClient,
        ICBDDCSurrealSchemaInitializer schemaInitializer,
        ILogger<SurrealPeerOplogConfirmationStore>? logger = null)
    {
        _ = surrealEmbeddedClient ?? throw new ArgumentNullException(nameof(surrealEmbeddedClient));
        _surrealClient = surrealEmbeddedClient.Client;
        _schemaInitializer = schemaInitializer ?? throw new ArgumentNullException(nameof(schemaInitializer));
        _logger = logger ?? NullLogger<SurrealPeerOplogConfirmationStore>.Instance;
    }

    /// <inheritdoc />
    public override async Task EnsurePeerRegisteredAsync(
        string peerNodeId,
        string address,
        PeerType type,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(peerNodeId))
            throw new ArgumentException("Peer node id is required.", nameof(peerNodeId));

        var existing =
            await FindByPairAsync(peerNodeId, RegistrationSourceNodeId, cancellationToken);

        if (existing == null)
        {
            var created = new PeerOplogConfirmation
            {
                PeerNodeId = peerNodeId,
                SourceNodeId = RegistrationSourceNodeId,
                ConfirmedWall = 0,
                ConfirmedLogic = 0,
                ConfirmedHash = "",
                LastConfirmedUtc = DateTimeOffset.UtcNow,
                IsActive = true
            };

            await UpsertAsync(created, SurrealStoreRecordIds.PeerOplogConfirmation(peerNodeId, RegistrationSourceNodeId),
                cancellationToken);

            _logger.LogDebug("Registered peer confirmation tracking for {PeerNodeId} ({Address}, {Type}).", peerNodeId,
                address, type);
            return;
        }

        if (existing.IsActive) return;

        existing.IsActive = true;
        existing.LastConfirmedUtcMs = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        RecordId recordId =
            existing.Id ?? SurrealStoreRecordIds.PeerOplogConfirmation(peerNodeId, RegistrationSourceNodeId);
        await UpsertAsync(existing, recordId, cancellationToken);
    }

    /// <inheritdoc />
    public override async Task UpdateConfirmationAsync(
        string peerNodeId,
        string sourceNodeId,
        HlcTimestamp timestamp,
        string hash,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(peerNodeId))
            throw new ArgumentException("Peer node id is required.", nameof(peerNodeId));

        if (string.IsNullOrWhiteSpace(sourceNodeId))
            throw new ArgumentException("Source node id is required.", nameof(sourceNodeId));

        var existing = await FindByPairAsync(peerNodeId, sourceNodeId, cancellationToken);
        long nowMs = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();

        if (existing == null)
        {
            var created = new PeerOplogConfirmation
            {
                PeerNodeId = peerNodeId,
                SourceNodeId = sourceNodeId,
                ConfirmedWall = timestamp.PhysicalTime,
                ConfirmedLogic = timestamp.LogicalCounter,
                ConfirmedHash = hash ?? "",
                LastConfirmedUtc = DateTimeOffset.FromUnixTimeMilliseconds(nowMs),
                IsActive = true
            };
            await UpsertAsync(created, SurrealStoreRecordIds.PeerOplogConfirmation(peerNodeId, sourceNodeId),
                cancellationToken);
            return;
        }

        bool isNewer = IsIncomingTimestampNewer(timestamp, existing);
        bool samePointHashChanged = timestamp.PhysicalTime == existing.ConfirmedWall &&
                                    timestamp.LogicalCounter == existing.ConfirmedLogic &&
                                    !string.Equals(existing.ConfirmedHash, hash, StringComparison.Ordinal);

        if (!isNewer && !samePointHashChanged && existing.IsActive) return;

        existing.ConfirmedWall = timestamp.PhysicalTime;
        existing.ConfirmedLogic = timestamp.LogicalCounter;
        existing.ConfirmedHash = hash ?? "";
        existing.LastConfirmedUtcMs = nowMs;
        existing.IsActive = true;

        RecordId recordId = existing.Id ?? SurrealStoreRecordIds.PeerOplogConfirmation(peerNodeId, sourceNodeId);
        await UpsertAsync(existing, recordId, cancellationToken);
    }

    /// <inheritdoc />
    public override async Task<IEnumerable<PeerOplogConfirmation>> GetConfirmationsAsync(
        CancellationToken cancellationToken = default)
    {
        var all = await SelectAllAsync(cancellationToken);
        return all
            .Where(c => !string.Equals(c.SourceNodeId, RegistrationSourceNodeId, StringComparison.Ordinal))
            .Select(c => c.ToDomain())
            .ToList();
    }

    /// <inheritdoc />
    public override async Task<IEnumerable<PeerOplogConfirmation>> GetConfirmationsForPeerAsync(
        string peerNodeId,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(peerNodeId))
            throw new ArgumentException("Peer node id is required.", nameof(peerNodeId));

        var all = await SelectAllAsync(cancellationToken);
        return all
            .Where(c => string.Equals(c.PeerNodeId, peerNodeId, StringComparison.Ordinal) &&
                        !string.Equals(c.SourceNodeId, RegistrationSourceNodeId, StringComparison.Ordinal))
            .Select(c => c.ToDomain())
            .ToList();
    }

    /// <inheritdoc />
    public override async Task RemovePeerTrackingAsync(string peerNodeId, CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(peerNodeId))
            throw new ArgumentException("Peer node id is required.", nameof(peerNodeId));

        var matches = (await SelectAllAsync(cancellationToken))
            .Where(c => string.Equals(c.PeerNodeId, peerNodeId, StringComparison.Ordinal))
            .ToList();

        if (matches.Count == 0) return;

        long nowMs = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        foreach (var match in matches)
        {
            if (!match.IsActive) continue;

            match.IsActive = false;
            match.LastConfirmedUtcMs = nowMs;

            RecordId recordId = match.Id ?? SurrealStoreRecordIds.PeerOplogConfirmation(match.PeerNodeId, match.SourceNodeId);
            await UpsertAsync(match, recordId, cancellationToken);
        }
    }

    /// <inheritdoc />
    public override async Task<IEnumerable<string>> GetActiveTrackedPeersAsync(
        CancellationToken cancellationToken = default)
    {
        var all = await SelectAllAsync(cancellationToken);
        return all
            .Where(c => c.IsActive)
            .Select(c => c.PeerNodeId)
            .Distinct(StringComparer.Ordinal)
            .ToList();
    }

    /// <inheritdoc />
    public override async Task DropAsync(CancellationToken cancellationToken = default)
    {
        await EnsureReadyAsync(cancellationToken);
        await _surrealClient.Delete(CBDDCSurrealSchemaNames.PeerOplogConfirmationsTable, cancellationToken);
    }

    /// <inheritdoc />
    public override async Task<IEnumerable<PeerOplogConfirmation>> ExportAsync(CancellationToken cancellationToken = default)
    {
        var all = await SelectAllAsync(cancellationToken);
        return all.Select(c => c.ToDomain()).ToList();
    }

    /// <inheritdoc />
    public override async Task ImportAsync(IEnumerable<PeerOplogConfirmation> items,
        CancellationToken cancellationToken = default)
    {
        foreach (var item in items)
        {
            var existing = await FindByPairAsync(item.PeerNodeId, item.SourceNodeId, cancellationToken);
            RecordId recordId =
                existing?.Id ?? SurrealStoreRecordIds.PeerOplogConfirmation(item.PeerNodeId, item.SourceNodeId);
            await UpsertAsync(item, recordId, cancellationToken);
        }
    }

    /// <inheritdoc />
    public override async Task MergeAsync(IEnumerable<PeerOplogConfirmation> items,
        CancellationToken cancellationToken = default)
    {
        foreach (var item in items)
        {
            var existing = await FindByPairAsync(item.PeerNodeId, item.SourceNodeId, cancellationToken);
            if (existing == null)
            {
                await UpsertAsync(item, SurrealStoreRecordIds.PeerOplogConfirmation(item.PeerNodeId, item.SourceNodeId),
                    cancellationToken);
                continue;
            }

            bool changed = false;
            var incomingTimestamp = new HlcTimestamp(item.ConfirmedWall, item.ConfirmedLogic, item.SourceNodeId);
            var existingTimestamp = new HlcTimestamp(existing.ConfirmedWall, existing.ConfirmedLogic, existing.SourceNodeId);

            if (incomingTimestamp > existingTimestamp)
            {
                existing.ConfirmedWall = item.ConfirmedWall;
                existing.ConfirmedLogic = item.ConfirmedLogic;
                existing.ConfirmedHash = item.ConfirmedHash;
                changed = true;
            }

            long incomingLastConfirmedMs = item.LastConfirmedUtc.ToUnixTimeMilliseconds();
            if (incomingLastConfirmedMs > existing.LastConfirmedUtcMs)
            {
                existing.LastConfirmedUtcMs = incomingLastConfirmedMs;
                changed = true;
            }

            if (existing.IsActive != item.IsActive)
            {
                existing.IsActive = item.IsActive;
                changed = true;
            }

            if (!changed) continue;

            RecordId recordId =
                existing.Id ?? SurrealStoreRecordIds.PeerOplogConfirmation(existing.PeerNodeId, existing.SourceNodeId);
            await UpsertAsync(existing, recordId, cancellationToken);
        }
    }

    private async Task UpsertAsync(PeerOplogConfirmation confirmation, RecordId recordId, CancellationToken cancellationToken)
    {
        await EnsureReadyAsync(cancellationToken);
        await _surrealClient.Upsert<SurrealPeerOplogConfirmationRecord, SurrealPeerOplogConfirmationRecord>(
            recordId,
            confirmation.ToSurrealRecord(),
            cancellationToken);
    }

    private async Task UpsertAsync(SurrealPeerOplogConfirmationRecord confirmation, RecordId recordId,
        CancellationToken cancellationToken)
    {
        await EnsureReadyAsync(cancellationToken);
        await _surrealClient.Upsert<SurrealPeerOplogConfirmationRecord, SurrealPeerOplogConfirmationRecord>(
            recordId,
            confirmation,
            cancellationToken);
    }

    private async Task EnsureReadyAsync(CancellationToken cancellationToken)
    {
        await _schemaInitializer.EnsureInitializedAsync(cancellationToken);
    }

    private async Task<List<SurrealPeerOplogConfirmationRecord>> SelectAllAsync(CancellationToken cancellationToken)
    {
        await EnsureReadyAsync(cancellationToken);
        var rows = await _surrealClient.Select<SurrealPeerOplogConfirmationRecord>(
            CBDDCSurrealSchemaNames.PeerOplogConfirmationsTable,
            cancellationToken);
        return rows?.ToList() ?? [];
    }

    private async Task<SurrealPeerOplogConfirmationRecord?> FindByPairAsync(string peerNodeId, string sourceNodeId,
        CancellationToken cancellationToken)
    {
        await EnsureReadyAsync(cancellationToken);
        RecordId deterministicId = SurrealStoreRecordIds.PeerOplogConfirmation(peerNodeId, sourceNodeId);
        var deterministic = await _surrealClient.Select<SurrealPeerOplogConfirmationRecord>(deterministicId, cancellationToken);
        if (deterministic != null &&
            string.Equals(deterministic.PeerNodeId, peerNodeId, StringComparison.Ordinal) &&
            string.Equals(deterministic.SourceNodeId, sourceNodeId, StringComparison.Ordinal))
            return deterministic;

        var all = await SelectAllAsync(cancellationToken);
        return all.FirstOrDefault(c =>
            string.Equals(c.PeerNodeId, peerNodeId, StringComparison.Ordinal) &&
            string.Equals(c.SourceNodeId, sourceNodeId, StringComparison.Ordinal));
    }

    private static bool IsIncomingTimestampNewer(HlcTimestamp incomingTimestamp, SurrealPeerOplogConfirmationRecord existing)
    {
        if (incomingTimestamp.PhysicalTime > existing.ConfirmedWall) return true;

        if (incomingTimestamp.PhysicalTime == existing.ConfirmedWall &&
            incomingTimestamp.LogicalCounter > existing.ConfirmedLogic)
            return true;

        return false;
    }
}