package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 节点本地联邦中间结果内存的保守准入控制器。
 */
final class NodeMemoryAdmissionController {

    private final long maximumBytes;
    private final ReentrantLock lock = new ReentrantLock(true);
    private final Condition released = lock.newCondition();
    private long reservedBytes;

    /**
     * 创建节点内存准入控制器。
     *
     * @param maximumBytes 节点允许同时预留的最大字节数
     */
    NodeMemoryAdmissionController(long maximumBytes) {
        if (maximumBytes <= 0L) {
            throw new IllegalArgumentException("maximumBytes must be positive");
        }
        this.maximumBytes = maximumBytes;
    }

    /**
     * 为一次联邦查询预留最坏情况下的本地中间结果预算。
     *
     * @param requestedBytes 查询预算
     * @param deadline 统一查询截止时间
     * @return 内存许可
     */
    Permit acquire(long requestedBytes, QueryDeadline deadline) {
        if (requestedBytes <= 0L || requestedBytes > maximumBytes) {
            throw new FederationSqlException(
                FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
                "federation intermediate memory budget exceeds the node limit"
            );
        }
        lock.lock();
        try {
            while (maximumBytes - reservedBytes < requestedBytes) {
                ensureAdmissionAllowed(deadline);
                long waitNanos = Math.max(
                    1L,
                    Math.min(deadline.remainingNanos(), TimeUnit.MILLISECONDS.toNanos(50))
                );
                try {
                    released.awaitNanos(waitNanos);
                } catch (InterruptedException exception) {
                    Thread.currentThread().interrupt();
                    throw new FederationSqlException(
                        FederationSqlErrorCode.EXECUTION_FAILED,
                        "waiting for federation node memory admission was interrupted",
                        exception
                    );
                }
            }
            ensureAdmissionAllowed(deadline);
            reservedBytes += requestedBytes;
            return new Permit(requestedBytes, this);
        } finally {
            lock.unlock();
        }
    }

    private static void ensureAdmissionAllowed(QueryDeadline deadline) {
        try {
            deadline.ensureAllowed();
        } catch (FederationSqlException exception) {
            if (exception.errorCode() == FederationSqlErrorCode.QUERY_CANCELLED) {
                throw exception;
            }
            if (exception.errorCode() == FederationSqlErrorCode.QUERY_TIMEOUT
                || exception.errorCode() == FederationSqlErrorCode.SQL_COMPILE_TIMEOUT) {
                throw new FederationSqlException(
                    FederationSqlErrorCode.NODE_MEMORY_ADMISSION_TIMEOUT,
                    "timed out while waiting for federation node memory admission",
                    exception
                );
            }
            throw exception;
        }
    }

    /**
     * 返回无需释放资源的许可。
     *
     * @return 空许可
     */
    static Permit none() {
        return new Permit(0L, null);
    }

    /**
     * 节点内存预留许可。
     */
    static final class Permit implements AutoCloseable {

        private final long bytes;
        private final NodeMemoryAdmissionController owner;
        private final AtomicBoolean closed = new AtomicBoolean();

        private Permit(long bytes, NodeMemoryAdmissionController owner) {
            this.bytes = bytes;
            this.owner = owner;
        }

        /**
         * 释放节点内存预算。
         */
        @Override
        public void close() {
            if (owner == null || !closed.compareAndSet(false, true)) {
                return;
            }
            owner.lock.lock();
            try {
                owner.reservedBytes -= bytes;
                owner.released.signalAll();
            } finally {
                owner.lock.unlock();
            }
        }
    }
}
