using ZB.MOM.WW.CBDD.Core.Transactions;

namespace ZB.MOM.WW.CBDD.Core.Storage;

public sealed partial class StorageEngine
{
    /// <summary>
    ///     Gets the number of active transactions (diagnostics).
    /// </summary>
    public int ActiveTransactionCount => _walCache.Count;

    /// <summary>
    ///     Prepares a transaction: writes all changes to WAL but doesn't commit yet.
    ///     Part of 2-Phase Commit protocol.
    /// </summary>
    /// <param name="transactionId">Transaction ID</param>
    /// <param name="writeSet">All writes to record in WAL</param>
    /// <returns>True if preparation succeeded</returns>
    public bool PrepareTransaction(ulong transactionId)
    {
        try
        {
            _wal.WriteBeginRecord(transactionId);

            foreach (var walEntry in _walCache[transactionId])
                _wal.WriteDataRecord(transactionId, walEntry.Key, walEntry.Value);

            _wal.Flush(); // Ensure WAL is persisted
            return true;
        }
        catch
        {
            // TODO: Log error?
            return false;
        }
    }

    /// <summary>
    ///     Prepares a transaction asynchronously by writing pending changes to the WAL.
    /// </summary>
    /// <param name="transactionId">The transaction identifier.</param>
    /// <param name="ct">The cancellation token.</param>
    /// <returns><see langword="true" /> if preparation succeeds; otherwise, <see langword="false" />.</returns>
    public async Task<bool> PrepareTransactionAsync(ulong transactionId, CancellationToken ct = default)
    {
        try
        {
            await _wal.WriteBeginRecordAsync(transactionId, ct);

            if (_walCache.TryGetValue(transactionId, out var changes))
                foreach (var walEntry in changes)
                    await _wal.WriteDataRecordAsync(transactionId, walEntry.Key, walEntry.Value, ct);

            await _wal.FlushAsync(ct); // Ensure WAL is persisted
            return true;
        }
        catch
        {
            return false;
        }
    }

    /// <summary>
    ///     Commits a transaction:
    ///     1. Writes all changes to WAL (for durability)
    ///     2. Writes commit record
    ///     3. Flushes WAL to disk
    ///     4. Moves pages from cache to WAL index (for future reads)
    ///     5. Clears WAL cache
    /// </summary>
    /// <param name="transactionId">Transaction to commit</param>
    /// <param name="writeSet">All writes performed in this transaction (unused, kept for compatibility)</param>
    public void CommitTransaction(ulong transactionId)
    {
        _commitLock.Wait();
        try
        {
            CommitTransactionCore(transactionId);
        }
        finally
        {
            _commitLock.Release();
        }
    }

    private void CommitTransactionCore(ulong transactionId)
    {
        // Get ALL pages from WAL cache (includes both data and index pages)
        if (!_walCache.TryGetValue(transactionId, out var pages))
        {
            // No writes for this transaction, just write commit record
            _wal.WriteCommitRecord(transactionId);
            _wal.Flush();
            return;
        }

        // 1. Write all changes to WAL (from cache, not writeSet!)
        _wal.WriteBeginRecord(transactionId);

        foreach ((uint pageId, byte[] data) in pages) _wal.WriteDataRecord(transactionId, pageId, data);

        // 2. Write commit record and flush
        _wal.WriteCommitRecord(transactionId);
        _wal.Flush(); // Durability: ensure WAL is on disk

        // 3. Move pages from cache to WAL index (for reads)
        _walCache.TryRemove(transactionId, out _);
        foreach (var kvp in pages) _walIndex[kvp.Key] = kvp.Value;

        // Auto-checkpoint if WAL grows too large
        if (_wal.GetCurrentSize() > MaxWalSize) CheckpointInternal();
    }

    /// <summary>
    ///     Commits a prepared transaction asynchronously by identifier.
    /// </summary>
    /// <param name="transactionId">The transaction identifier.</param>
    /// <param name="ct">The cancellation token.</param>
    public async Task CommitTransactionAsync(ulong transactionId, CancellationToken ct = default)
    {
        await _commitLock.WaitAsync(ct);
        try
        {
            await CommitTransactionCoreAsync(transactionId, ct);
        }
        finally
        {
            _commitLock.Release();
        }
    }

    private async Task CommitTransactionCoreAsync(ulong transactionId, CancellationToken ct)
    {
        // Get ALL pages from WAL cache (includes both data and index pages)
        if (!_walCache.TryGetValue(transactionId, out var pages))
        {
            // No writes for this transaction, just write commit record
            await _wal.WriteCommitRecordAsync(transactionId, ct);
            await _wal.FlushAsync(ct);
            return;
        }

        // 1. Write all changes to WAL (from cache, not writeSet!)
        await _wal.WriteBeginRecordAsync(transactionId, ct);

        foreach ((uint pageId, byte[] data) in pages) await _wal.WriteDataRecordAsync(transactionId, pageId, data, ct);

        // 2. Write commit record and flush
        await _wal.WriteCommitRecordAsync(transactionId, ct);
        await _wal.FlushAsync(ct); // Durability: ensure WAL is on disk

        // 3. Move pages from cache to WAL index (for reads)
        _walCache.TryRemove(transactionId, out _);
        foreach (var kvp in pages) _walIndex[kvp.Key] = kvp.Value;

        // Auto-checkpoint if WAL grows too large
        if (_wal.GetCurrentSize() > MaxWalSize)
            // Checkpoint might be sync or async. For now sync inside the lock is "safe" but blocking.
            // Ideally this should be async too.
            CheckpointInternal();
    }

    /// <summary>
    ///     Marks a transaction as committed after WAL writes.
    ///     Used for 2PC: after Prepare() writes to WAL, this finalizes the commit.
    /// </summary>
    /// <param name="transactionId">Transaction to mark committed</param>
    public void MarkTransactionCommitted(ulong transactionId)
    {
        _commitLock.Wait();
        try
        {
            _wal.WriteCommitRecord(transactionId);
            _wal.Flush();

            // Move from cache to WAL index
            if (_walCache.TryRemove(transactionId, out var pages))
                foreach (var kvp in pages)
                    _walIndex[kvp.Key] = kvp.Value;

            // Auto-checkpoint if WAL grows too large
            if (_wal.GetCurrentSize() > MaxWalSize) CheckpointInternal();
        }
        finally
        {
            _commitLock.Release();
        }
    }

    /// <summary>
    ///     Rolls back a transaction: discards all uncommitted changes.
    /// </summary>
    /// <param name="transactionId">Transaction to rollback</param>
    public void RollbackTransaction(ulong transactionId)
    {
        _walCache.TryRemove(transactionId, out _);
        _wal.WriteAbortRecord(transactionId);
    }

    /// <summary>
    ///     Writes an abort record for the specified transaction.
    /// </summary>
    /// <param name="transactionId">The transaction identifier.</param>
    internal void WriteAbortRecord(ulong transactionId)
    {
        _wal.WriteAbortRecord(transactionId);
    }

    #region Transaction Management

    /// <summary>
    ///     Begins a new transaction.
    /// </summary>
    /// <param name="isolationLevel">The transaction isolation level.</param>
    /// <returns>The started transaction.</returns>
    public Transaction BeginTransaction(IsolationLevel isolationLevel = IsolationLevel.ReadCommitted)
    {
        _commitLock.Wait();
        try
        {
            ulong txnId = _nextTransactionId++;
            var transaction = new Transaction(txnId, this, isolationLevel);
            _activeTransactions[txnId] = transaction;
            return transaction;
        }
        finally
        {
            _commitLock.Release();
        }
    }

    /// <summary>
    ///     Begins a new transaction asynchronously.
    /// </summary>
    /// <param name="isolationLevel">The transaction isolation level.</param>
    /// <param name="ct">The cancellation token.</param>
    /// <returns>The started transaction.</returns>
    public async Task<Transaction> BeginTransactionAsync(IsolationLevel isolationLevel = IsolationLevel.ReadCommitted,
        CancellationToken ct = default)
    {
        await _commitLock.WaitAsync(ct);
        try
        {
            ulong txnId = _nextTransactionId++;
            var transaction = new Transaction(txnId, this, isolationLevel);
            _activeTransactions[txnId] = transaction;
            return transaction;
        }
        finally
        {
            _commitLock.Release();
        }
    }

    /// <summary>
    ///     Commits the specified transaction.
    /// </summary>
    /// <param name="transaction">The transaction to commit.</param>
    public void CommitTransaction(Transaction transaction)
    {
        _commitLock.Wait();
        try
        {
            if (!_activeTransactions.ContainsKey(transaction.TransactionId))
                throw new InvalidOperationException($"Transaction {transaction.TransactionId} is not active.");

            // 1. Prepare (Write to WAL)
            // In a fuller 2PC, this would be separate. Here we do it as part of commit.
            if (!PrepareTransaction(transaction.TransactionId))
                throw new IOException("Failed to write transaction to WAL");

            // 2. Commit (Write commit record, flush, move to cache)
            // Use core commit path to avoid re-entering _commitLock.
            CommitTransactionCore(transaction.TransactionId);

            _activeTransactions.TryRemove(transaction.TransactionId, out _);
        }
        finally
        {
            _commitLock.Release();
        }
    }

    /// <summary>
    ///     Commits the specified transaction asynchronously.
    /// </summary>
    /// <param name="transaction">The transaction to commit.</param>
    /// <param name="ct">The cancellation token.</param>
    public async Task CommitTransactionAsync(Transaction transaction, CancellationToken ct = default)
    {
        await _commitLock.WaitAsync(ct);
        try
        {
            if (!_activeTransactions.ContainsKey(transaction.TransactionId))
                throw new InvalidOperationException($"Transaction {transaction.TransactionId} is not active.");

            // 1. Prepare (Write to WAL)
            if (!await PrepareTransactionAsync(transaction.TransactionId, ct))
                throw new IOException("Failed to write transaction to WAL");

            // 2. Commit (Write commit record, flush, move to cache)
            // Use core commit path to avoid re-entering _commitLock.
            await CommitTransactionCoreAsync(transaction.TransactionId, ct);

            _activeTransactions.TryRemove(transaction.TransactionId, out _);
        }
        finally
        {
            _commitLock.Release();
        }
    }

    /// <summary>
    ///     Rolls back the specified transaction.
    /// </summary>
    /// <param name="transaction">The transaction to roll back.</param>
    public void RollbackTransaction(Transaction transaction)
    {
        RollbackTransaction(transaction.TransactionId);
        _activeTransactions.TryRemove(transaction.TransactionId, out _);
    }

    // Rollback doesn't usually require async logic unless logging abort record is async, 
    // but for consistency we might consider it. For now, sync is fine as it's not the happy path bottleneck.

    #endregion
}