package com.easyagents.federation.sql.runtime;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;

import com.easyagents.federation.sql.adapter.AdapterCompatibility;
import com.easyagents.federation.sql.adapter.AdapterDialectContext;
import com.easyagents.federation.sql.adapter.AdapterHints;
import com.easyagents.federation.sql.adapter.AdapterSchemaContext;
import com.easyagents.federation.sql.adapter.FederationSqlAdapterProvider;
import com.easyagents.federation.sql.adapter.FederationStatisticsCollector;
import com.easyagents.federation.sql.execute.FederationFragmentExecutor;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationStatisticsSnapshot;
import com.easyagents.federation.sql.federation.FederationStatisticsStatus;
import com.easyagents.federation.sql.federation.FederationTableStatistics;
import com.easyagents.federation.sql.source.ExternalSchemaDefinition;
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
import com.easyagents.federation.sql.source.FederationSourceDefinition;
import com.easyagents.federation.sql.source.RuntimeFingerprint;
import com.easyagents.federation.sql.source.SourceId;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import javax.sql.DataSource;
import org.apache.calcite.schema.Schema;
import org.apache.calcite.schema.impl.AbstractSchema;
import org.apache.calcite.sql.SqlDialect;
import org.junit.Test;

/**
 * Adapter 自动统计缓存的并发合并和失败降级测试。
 */
public class AdapterFederationStatisticsProviderTest {

    /**
     * 验证同一物理源 revision 的并发刷新只触发一次目录采集且不阻塞调用线程。
     *
     * @throws Exception 并发等待超时
     */
    @Test
    public void shouldMergeConcurrentRefreshForSameRevision() throws Exception {
        AtomicInteger collections = new AtomicInteger();
        CountDownLatch started = new CountDownLatch(1);
        CountDownLatch release = new CountDownLatch(1);
        FederationStatisticsCollector collector = context -> {
            collections.incrementAndGet();
            started.countDown();
            try {
                if (!release.await(5, TimeUnit.SECONDS)) {
                    throw new IllegalStateException("statistics test collector timed out");
                }
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
                throw new IllegalStateException("statistics test collector interrupted", exception);
            }
            FederationStatisticsSnapshot.TableKey key =
                new FederationStatisticsSnapshot.TableKey(
                    context.sourceDefinition().sourceId(),
                    "main",
                    "orders"
                );
            return Map.of(key, new FederationTableStatistics(
                42D,
                64L,
                context.collectedAt(),
                "test",
                Map.of(),
                List.of(),
                context.expiresAt(),
                FederationStatisticsStatus.PARTIAL
            ));
        };
        AdapterFederationStatisticsProvider provider = provider();
        SourceRuntime runtime = runtime(1L, collector);
        FederationQueryScopeSnapshot snapshot = snapshot(runtime);
        try {
            CompletableFuture<Void> first = provider.refreshIfNeeded(snapshot);
            if (!started.await(5, TimeUnit.SECONDS)) {
                throw new AssertionError("statistics collection did not start");
            }
            assertFalse("refresh should still be waiting on the collector", first.isDone());
            CompletableFuture<Void> second = provider.refreshIfNeeded(snapshot);
            release.countDown();
            CompletableFuture.allOf(first, second).get(5, TimeUnit.SECONDS);

            assertEquals(1, collections.get());
            assertNotNull(provider.snapshot().statistics(
                runtime.definition().sourceId(),
                "main",
                "orders"
            ));
        } finally {
            release.countDown();
            snapshot.close();
            provider.close();
        }
    }

    /**
     * 验证目录读取失败不阻断查询，并在失败限频窗口内避免重复访问数据库。
     */
    @Test
    public void shouldThrottleCollectionFailureWithoutFailingQuery() throws Exception {
        AtomicInteger collections = new AtomicInteger();
        FederationStatisticsCollector collector = context -> {
            collections.incrementAndGet();
            throw new java.sql.SQLException("catalog unavailable");
        };
        AdapterFederationStatisticsProvider provider = provider();
        SourceRuntime runtime = runtime(1L, collector);
        FederationQueryScopeSnapshot snapshot = snapshot(runtime);
        try {
            provider.refreshIfNeeded(snapshot).get(5, TimeUnit.SECONDS);
            provider.refreshIfNeeded(snapshot).get(5, TimeUnit.SECONDS);

            assertEquals(1, collections.get());
            assertEquals(null, provider.snapshot().statistics(
                runtime.definition().sourceId(),
                "main",
                "orders"
            ));
        } finally {
            snapshot.close();
            provider.close();
        }
    }

    /**
     * 验证异步采集持有独立 Runtime 租约，退休过程不会提前关闭正在使用的数据源。
     *
     * @throws Exception 并发等待超时
     */
    @Test
    public void shouldKeepRuntimeAliveDuringAsynchronousRefresh() throws Exception {
        CountDownLatch started = new CountDownLatch(1);
        CountDownLatch release = new CountDownLatch(1);
        AtomicInteger closedRuntimes = new AtomicInteger();
        FederationStatisticsCollector collector = context -> {
            started.countDown();
            try {
                if (!release.await(5, TimeUnit.SECONDS)) {
                    throw new IllegalStateException("statistics test collector timed out");
                }
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
                throw new IllegalStateException("statistics test collector interrupted", exception);
            }
            return Map.of();
        };
        AdapterFederationStatisticsProvider provider = provider();
        SourceRuntime runtime = runtime(1L, collector, ignored ->
            closedRuntimes.incrementAndGet()
        );
        FederationQueryScopeSnapshot snapshot = snapshot(runtime);
        try {
            CompletableFuture<Void> refresh = provider.refreshIfNeeded(snapshot);
            if (!started.await(5, TimeUnit.SECONDS)) {
                throw new AssertionError("statistics collection did not start");
            }

            runtime.retire();
            assertEquals(0, closedRuntimes.get());

            release.countDown();
            refresh.get(5, TimeUnit.SECONDS);
            assertEquals(1, closedRuntimes.get());
        } finally {
            release.countDown();
            snapshot.close();
            provider.close();
        }
    }

    /**
     * 创建短有效期的测试 Provider。
     *
     * @return 自动统计 Provider
     */
    private AdapterFederationStatisticsProvider provider() {
        return new AdapterFederationStatisticsProvider(
            Duration.ofMinutes(5),
            Duration.ofMinutes(1),
            1,
            2
        );
    }

    /**
     * 创建仅用于统计采集的最小 Source Runtime。
     *
     * @param revision Definition revision
     * @param collector 测试统计采集器
     * @return Source Runtime
     */
    private SourceRuntime runtime(long revision, FederationStatisticsCollector collector) {
        return runtime(revision, collector, ignored -> { });
    }

    /**
     * 创建可观察关闭事件的最小 Source Runtime。
     *
     * @param revision Definition revision
     * @param collector 测试统计采集器
     * @param closedListener Runtime 关闭监听器
     * @return Source Runtime
     */
    private SourceRuntime runtime(
        long revision,
        FederationStatisticsCollector collector,
        java.util.function.Consumer<SourceRuntime> closedListener
    ) {
        SourceId sourceId = new SourceId("source-a");
        FederationSourceDefinition definition = new FederationSourceDefinition(
            sourceId,
            revision,
            "test-adapter",
            List.of(new ExternalSchemaDefinition("main", "schema", 1)),
            Map.of()
        );
        return new SourceRuntime(
            definition,
            FederationDataSourceHandles.shared(
                dataSource(),
                new RuntimeFingerprint("test", "1", "test", "1", "1")
            ),
            new TestAdapter(collector),
            SqlDialect.DatabaseProduct.CALCITE.getDialect(),
            null,
            null,
            null,
            closedListener
        );
    }

    /**
     * 创建借用 Runtime 的单源查询快照。
     *
     * @param runtime Source Runtime
     * @return 查询快照
     */
    private FederationQueryScopeSnapshot snapshot(SourceRuntime runtime) {
        FederationQueryScopeDefinition scope = FederationQueryScopeDefinition.single(
            runtime.definition().sourceId(),
            runtime.definition().revision()
        );
        return FederationQueryScopeSnapshot.borrowedSingle(scope, runtime);
    }

    /**
     * 创建支持只读标记和关闭的最小 DataSource。
     *
     * @return 测试 DataSource
     */
    private DataSource dataSource() {
        Connection connection = (Connection) Proxy.newProxyInstance(
            getClass().getClassLoader(),
            new Class<?>[] {Connection.class},
            (proxy, method, arguments) -> switch (method.getName()) {
                case "isReadOnly" -> false;
                case "setReadOnly", "close" -> null;
                case "isClosed" -> false;
                case "isWrapperFor" -> false;
                case "unwrap" -> null;
                default -> defaultValue(method.getReturnType());
            }
        );
        return (DataSource) Proxy.newProxyInstance(
            getClass().getClassLoader(),
            new Class<?>[] {DataSource.class},
            (proxy, method, arguments) -> switch (method.getName()) {
                case "getConnection" -> connection;
                case "isWrapperFor" -> false;
                case "unwrap" -> null;
                default -> defaultValue(method.getReturnType());
            }
        );
    }

    /**
     * 返回代理方法原始类型的默认值。
     *
     * @param type 返回类型
     * @return 默认值
     */
    private Object defaultValue(Class<?> type) {
        if (!type.isPrimitive()) {
            return null;
        }
        if (type == boolean.class) {
            return false;
        }
        if (type == char.class) {
            return '\0';
        }
        return 0;
    }

    /**
     * 暴露测试统计采集器的最小 Adapter。
     */
    private static final class TestAdapter implements FederationSqlAdapterProvider {

        private final FederationStatisticsCollector collector;

        /**
         * 创建测试 Adapter。
         *
         * @param collector 统计采集器
         */
        private TestAdapter(FederationStatisticsCollector collector) {
            this.collector = collector;
        }

        /** {@inheritDoc} */
        @Override
        public String adapterId() {
            return "test-adapter";
        }

        /** {@inheritDoc} */
        @Override
        public boolean supports(DatabaseMetaData metadata, AdapterHints hints) {
            return true;
        }

        /** {@inheritDoc} */
        @Override
        public AdapterCompatibility compatibility(
            DatabaseMetaData metadata,
            AdapterHints hints
        ) {
            throw new UnsupportedOperationException();
        }

        /** {@inheritDoc} */
        @Override
        public Schema createSchema(AdapterSchemaContext context) {
            return new AbstractSchema();
        }

        /** {@inheritDoc} */
        @Override
        public SqlDialect createDialect(AdapterDialectContext context) {
            return SqlDialect.DatabaseProduct.CALCITE.getDialect();
        }

        /** {@inheritDoc} */
        @Override
        public FederationFragmentExecutor fragmentExecutor() {
            throw new UnsupportedOperationException();
        }

        /** {@inheritDoc} */
        @Override
        public Optional<FederationStatisticsCollector> statisticsCollector() {
            return Optional.of(collector);
        }
    }
}
