package com.easyagents.federation.sql.runtime;

import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import org.apache.calcite.avatica.util.ByteString;
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 com.easyagents.federation.sql.source.SourceId;
import java.util.List;
import org.junit.Assert;
import org.junit.Test;

/**
 * 联邦查询指标估算边界测试。
 */
public class FederationQueryMetricsTrackerTest {

    /**
     * 验证 Calcite 二进制值按实际长度计入中间字节预算。
     */
    @Test
    public void shouldMeasureCalciteByteStringByActualLength() {
        long small = FederationQueryMetricsTracker.estimateRowBytes(
            new Object[] {new ByteString(new byte[8])}
        );
        long large = FederationQueryMetricsTracker.estimateRowBytes(
            new Object[] {new ByteString(new byte[4096])}
        );

        Assert.assertEquals(4096 - 8, large - small);
    }

    /**
     * 验证无法保守计量的专有对象不会按固定小值放过预算。
     */
    @Test
    public void shouldRejectUnmeasurableIntermediateValue() {
        try {
            FederationQueryMetricsTracker.estimateRowBytes(new Object[] {new Object()});
            Assert.fail("unknown values must fail closed");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(
                FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED,
                exception.errorCode()
            );
        }
    }

    /**
     * 验证查询、分片和本地算子阶段指标会在终态冻结。
     */
    @Test
    public void shouldPublishDetailedTerminalMetrics() {
        FederationQueryMetricsTracker tracker = new FederationQueryMetricsTracker(
            new QueryId("metrics-query"),
            FederationQueryMode.FEDERATED,
            true,
            10,
            List.of(new FederationFragmentPlan(
                "fragment-1",
                "SOURCE",
                new SourceId("source"),
                "SELECT 1",
                List.of(),
                List.of()
            ))
        );
        tracker.recordAdmissionWait(11);
        tracker.fragmentObserver("fragment-1").connectionAcquired(12);
        tracker.fragmentObserver("fragment-1").databaseExecutionCompleted(13);
        tracker.recordIntermediate("fragment-1", new Object[] {1});
        tracker.recordLocalIntermediate("EnumerableHashJoin", new Object[] {1}, 14);
        tracker.recordOutput(new Object[] {1});
        tracker.finishFragment("fragment-1");
        tracker.finishSuccessfully();

        FederationQueryMetricsSnapshot snapshot = tracker.snapshot();
        long frozenExecution = snapshot.executionNanos();
        Assert.assertTrue(snapshot.complete());
        Assert.assertTrue(snapshot.planCacheHit());
        Assert.assertEquals(11, snapshot.admissionWaitNanos());
        Assert.assertEquals(12, snapshot.connectionAcquireNanos());
        Assert.assertEquals(13, snapshot.databaseExecutionNanos());
        Assert.assertEquals(14, snapshot.localExecutionNanos());
        Assert.assertEquals(1, snapshot.localOperators().size());
        Assert.assertEquals(1, snapshot.fragments().size());
        Assert.assertEquals(12, snapshot.fragments().get(0).connectionAcquireNanos());
        Assert.assertEquals(13, snapshot.fragments().get(0).databaseExecutionNanos());
        Assert.assertEquals(frozenExecution, tracker.snapshot().executionNanos());
    }
}
