package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.execute.FederationFragmentMetrics;
import com.easyagents.federation.sql.execute.FederationExecutionObserver;
import com.easyagents.federation.sql.execute.FederationLocalOperatorMetrics;
import com.easyagents.federation.sql.execute.FederationQueryMetricsSnapshot;
import com.easyagents.federation.sql.execute.QueryId;
import com.easyagents.federation.sql.federation.FederationFragmentPlan;
import com.easyagents.federation.sql.federation.FederationQueryMode;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.calcite.avatica.util.ByteString;

/**
 * 维护查询级和分片级有界、低开销指标。
 */
final class FederationQueryMetricsTracker {

    private final QueryId queryId;
    private final FederationQueryMode queryMode;
    private final boolean planCacheHit;
    private final long planningNanos;
    private final long executionStartedNanos = System.nanoTime();
    private final AtomicLong executionFinishedNanos = new AtomicLong(-1);
    private final AtomicLong firstRowNanos = new AtomicLong(-1);
    private final AtomicLong returnedRows = new AtomicLong();
    private final AtomicLong returnedBytes;
    private final AtomicLong intermediateRows = new AtomicLong();
    private final AtomicLong intermediateBytes;
    private final AtomicLong admissionWaitNanos = new AtomicLong();
    private final AtomicLong connectionAcquireNanos = new AtomicLong();
    private final AtomicLong databaseExecutionNanos = new AtomicLong();
    private final AtomicLong localExecutionNanos = new AtomicLong();
    private final AtomicBoolean complete = new AtomicBoolean();
    private final AtomicBoolean cancelled = new AtomicBoolean();
    private final AtomicBoolean timedOut = new AtomicBoolean();
    private final AtomicBoolean truncated = new AtomicBoolean();
    private final AtomicReference<String> terminalErrorCode = new AtomicReference<>("");
    private final Map<String, FragmentState> fragments = new LinkedHashMap<>();
    private final Map<String, LocalOperatorState> localOperators = new LinkedHashMap<>();

    /**
     * 创建指标跟踪器。
     *
     * @param queryId 查询标识
     * @param queryMode 查询模式
     * @param planCacheHit 是否命中计划缓存
     * @param planningNanos 编译耗时
     * @param fragmentPlans 分片计划
     */
    FederationQueryMetricsTracker(
        QueryId queryId,
        FederationQueryMode queryMode,
        boolean planCacheHit,
        long planningNanos,
        List<FederationFragmentPlan> fragmentPlans
    ) {
        this.queryId = queryId;
        this.queryMode = queryMode;
        this.planCacheHit = planCacheHit;
        this.planningNanos = planningNanos;
        boolean byteMetricsAvailable = queryMode != FederationQueryMode.SINGLE_SOURCE;
        this.returnedBytes = new AtomicLong(byteMetricsAvailable ? 0 : -1);
        this.intermediateBytes = new AtomicLong(byteMetricsAvailable ? 0 : -1);
        fragmentPlans.forEach(fragment -> fragments.put(
            fragment.fragmentId(),
            new FragmentState(fragment, byteMetricsAvailable)
        ));
    }

    /**
     * 记录调用方消费的一行最终结果。
     *
     * @param row 行值
     */
    void recordOutput(Object[] row) {
        firstRowNanos.compareAndSet(-1, System.nanoTime() - executionStartedNanos);
        returnedRows.incrementAndGet();
        returnedBytes.addAndGet(estimateRowBytes(row));
    }

    /**
     * 记录单源流式游标的一行，避免指标采集主动读取尚未被调用方访问的大字段。
     */
    void recordSingleSourceRow() {
        firstRowNanos.compareAndSet(-1, System.nanoTime() - executionStartedNanos);
        returnedRows.incrementAndGet();
        intermediateRows.incrementAndGet();
        fragment("fragment-1").rows.incrementAndGet();
    }

    /**
     * 记录查询级准入等待耗时。
     *
     * @param elapsedNanos 等待耗时
     */
    void recordAdmissionWait(long elapsedNanos) {
        admissionWaitNanos.addAndGet(Math.max(0, elapsedNanos));
    }

    /**
     * 返回指定分片的 Adapter 阶段观察器。
     *
     * @param fragmentId 分片标识
     * @return 阶段观察器
     */
    FederationExecutionObserver fragmentObserver(String fragmentId) {
        FragmentState state = fragment(fragmentId);
        state.start();
        return new FederationExecutionObserver() {
            @Override
            public void connectionAcquired(long elapsedNanos) {
                long safe = Math.max(0, elapsedNanos);
                state.connectionAcquireNanos.addAndGet(safe);
                connectionAcquireNanos.addAndGet(safe);
            }

            @Override
            public void databaseExecutionCompleted(long elapsedNanos) {
                long safe = Math.max(0, elapsedNanos);
                state.databaseExecutionNanos.addAndGet(safe);
                databaseExecutionNanos.addAndGet(safe);
            }

            @Override
            public void firstRowAvailable(long elapsedNanos) {
                state.firstRowNanos.compareAndSet(-1, Math.max(0, elapsedNanos));
            }
        };
    }

    /**
     * 记录物理分片读取的一行并返回累计中间结果计数。
     *
     * @param fragmentId 分片标识
     * @param row 行值
     * @return 当前累计行数和字节数
     */
    IntermediateUsage recordIntermediate(String fragmentId, Object[] row) {
        long bytes = estimateRowBytes(row);
        long rows = intermediateRows.incrementAndGet();
        long totalBytes = intermediateBytes.addAndGet(bytes);
        FragmentState fragment = fragment(fragmentId);
        fragment.rows.incrementAndGet();
        fragment.bytes.addAndGet(bytes);
        return new IntermediateUsage(rows, totalBytes);
    }

    /**
     * 记录 Calcite 本地算子交给下游的一行。
     *
     * @param row Calcite 内部行
     * @return 当前全局中间资源累计值
     */
    synchronized IntermediateUsage recordLocalIntermediate(
        String operatorName,
        Object[] row,
        long elapsedNanos
    ) {
        long bytes = estimateRowBytes(row);
        LocalOperatorState operator = localOperators.computeIfAbsent(
            operatorName,
            ignored -> new LocalOperatorState(operatorName)
        );
        operator.rows.incrementAndGet();
        operator.bytes.addAndGet(bytes);
        long safeElapsed = Math.max(0, elapsedNanos);
        operator.executionNanos.addAndGet(safeElapsed);
        localExecutionNanos.addAndGet(safeElapsed);
        return new IntermediateUsage(
            intermediateRows.incrementAndGet(),
            intermediateBytes.addAndGet(bytes)
        );
    }

    /**
     * 标记一个分片完成或关闭。
     *
     * @param fragmentId 分片标识
     */
    void finishFragment(String fragmentId) {
        fragment(fragmentId).finish();
    }

    /**
     * 标记最终结果已正常耗尽。
     */
    void finishSuccessfully() {
        complete.set(true);
        finishClosed();
    }

    /**
     * 标记查询收到取消请求。
     */
    void markCancelled() {
        if (terminalErrorCode.compareAndSet(
            "",
            FederationSqlErrorCode.QUERY_CANCELLED.name()
        )) {
            cancelled.set(true);
            finishClosed();
        }
    }

    /**
     * 标记查询因统一执行时限结束。
     */
    void markTimedOut() {
        if (terminalErrorCode.compareAndSet(
            "",
            FederationSqlErrorCode.QUERY_TIMEOUT.name()
        )) {
            timedOut.set(true);
            finishClosed();
        }
    }

    /**
     * 标记查询因最终结果行数上限停止继续消费。
     */
    void markTruncated() {
        truncated.set(true);
    }

    /**
     * 记录查询失败的稳定错误码。
     *
     * @param errorCode 错误码
     */
    void markFailed(FederationSqlErrorCode errorCode) {
        if (errorCode == FederationSqlErrorCode.QUERY_CANCELLED) {
            markCancelled();
            return;
        }
        if (errorCode == FederationSqlErrorCode.QUERY_TIMEOUT) {
            markTimedOut();
            return;
        }
        if (terminalErrorCode.compareAndSet("", errorCode.name())) {
            finishClosed();
        }
    }

    /**
     * 冻结查询和全部分片耗时；提前关闭不会标记查询正常完成。
     */
    void finishClosed() {
        long now = System.nanoTime();
        executionFinishedNanos.compareAndSet(-1, now);
        synchronized (this) {
            fragments.values().forEach(fragment -> fragment.finish(now));
        }
    }

    /**
     * 返回当前不可变指标快照。
     *
     * @return 指标快照
     */
    synchronized FederationQueryMetricsSnapshot snapshot() {
        long now = System.nanoTime();
        long executionEnd = executionFinishedNanos.get();
        long effectiveExecutionEnd = executionEnd < 0 ? now : executionEnd;
        List<FederationFragmentMetrics> fragmentSnapshots = new ArrayList<>(fragments.size());
        fragments.values().forEach(fragment -> fragmentSnapshots.add(
            new FederationFragmentMetrics(
                fragment.plan.fragmentId(),
                fragment.plan.sourceId(),
                fragment.rows.get(),
                fragment.bytes.get(),
                fragment.elapsedNanos(now),
                fragment.connectionAcquireNanos.get(),
                fragment.databaseExecutionNanos.get(),
                fragment.firstRowNanos.get(),
                fragment.complete.get()
            )
        ));
        List<FederationLocalOperatorMetrics> operatorSnapshots = new ArrayList<>(
            localOperators.size()
        );
        localOperators.values().forEach(operator -> operatorSnapshots.add(
            new FederationLocalOperatorMetrics(
                operator.name,
                operator.rows.get(),
                operator.bytes.get(),
                operator.executionNanos.get()
            )
        ));
        return new FederationQueryMetricsSnapshot(
            queryId,
            queryMode,
            planCacheHit,
            planningNanos,
            admissionWaitNanos.get(),
            connectionAcquireNanos.get(),
            databaseExecutionNanos.get(),
            localExecutionNanos.get(),
            effectiveExecutionEnd - executionStartedNanos,
            firstRowNanos.get(),
            returnedRows.get(),
            returnedBytes.get(),
            intermediateRows.get(),
            intermediateBytes.get(),
            complete.get(),
            cancelled.get(),
            timedOut.get(),
            truncated.get(),
            terminalErrorCode.get(),
            fragmentSnapshots,
            operatorSnapshots
        );
    }

    private synchronized FragmentState fragment(String fragmentId) {
        FragmentState state = fragments.get(fragmentId);
        if (state == null) {
            throw new IllegalArgumentException("unknown fragment: " + fragmentId);
        }
        return state;
    }

    /**
     * 对常见 JDBC 标量进行保守、常量时间的内存估算。
     */
    static long estimateRowBytes(Object[] row) {
        long bytes = 16L + 8L * row.length;
        for (Object value : row) {
            if (value == null) {
                continue;
            }
            if (value instanceof CharSequence text) {
                bytes += 40L + 2L * text.length();
            } else if (value instanceof byte[] binary) {
                bytes += 24L + binary.length;
            } else if (value instanceof ByteBuffer binary) {
                bytes += 24L + binary.remaining();
            } else if (value instanceof ByteString binary) {
                // Calcite Enumerable 使用 ByteString 承载 BINARY/VARBINARY。
                bytes += 24L + binary.length();
            } else if (value instanceof BigDecimal decimal) {
                bytes += 48L + decimal.precision() / 2L;
            } else if (value instanceof Number
                || value instanceof Boolean
                || value instanceof java.util.UUID
                || value instanceof java.time.OffsetTime
                || value instanceof java.time.OffsetDateTime
                || value instanceof java.time.LocalDate
                || value instanceof java.time.LocalTime
                || value instanceof java.time.LocalDateTime
                || value instanceof java.time.Instant
                || value instanceof java.sql.Date
                || value instanceof java.sql.Time
                || value instanceof java.sql.Timestamp) {
                bytes += 32L;
            } else {
                throw new FederationSqlException(
                    FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED,
                    "value cannot be safely measured for local federation: "
                        + value.getClass().getName()
                );
            }
        }
        return bytes;
    }

    /**
     * 累计中间结果资源用量。
     *
     * @param rows 行数
     * @param bytes 估算字节数
     */
    record IntermediateUsage(long rows, long bytes) {
    }

    private static final class FragmentState {

        private final FederationFragmentPlan plan;
        private final AtomicLong startedNanos = new AtomicLong(-1);
        private final AtomicLong finishedNanos = new AtomicLong(-1);
        private final AtomicLong rows = new AtomicLong();
        private final AtomicLong bytes;
        private final AtomicLong connectionAcquireNanos = new AtomicLong();
        private final AtomicLong databaseExecutionNanos = new AtomicLong();
        private final AtomicLong firstRowNanos = new AtomicLong(-1);
        private final AtomicBoolean complete = new AtomicBoolean();

        private FragmentState(FederationFragmentPlan plan, boolean byteMetricsAvailable) {
            this.plan = plan;
            this.bytes = new AtomicLong(byteMetricsAvailable ? 0 : -1);
        }

        private void finish() {
            finish(System.nanoTime());
        }

        private void start() {
            startedNanos.compareAndSet(-1, System.nanoTime());
        }

        private void finish(long now) {
            start();
            complete.set(true);
            finishedNanos.compareAndSet(-1, now);
        }

        private long elapsedNanos(long now) {
            long started = startedNanos.get();
            if (started < 0) {
                return 0;
            }
            long finished = finishedNanos.get();
            return (finished < 0 ? now : finished) - started;
        }
    }

    private static final class LocalOperatorState {

        private final String name;
        private final AtomicLong rows = new AtomicLong();
        private final AtomicLong bytes = new AtomicLong();
        private final AtomicLong executionNanos = new AtomicLong();

        private LocalOperatorState(String name) {
            this.name = name;
        }
    }
}
