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.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
import com.easyagents.federation.sql.federation.FederationSourceRuntimeIdentity;
import com.easyagents.federation.sql.source.SourceId;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 编译或执行期间持有的查询范围 Runtime 快照。
 */
final class FederationQueryScopeSnapshot implements AutoCloseable {

    private final FederationQueryScopeDefinition definition;
    private final Map<String, SourceRuntime> runtimesByBinding;
    private final List<SourceRuntime.RuntimeLease> leases;
    private final AtomicBoolean closed = new AtomicBoolean();

    private FederationQueryScopeSnapshot(
        FederationQueryScopeDefinition definition,
        Map<String, SourceRuntime> runtimesByBinding,
        List<SourceRuntime.RuntimeLease> leases
    ) {
        this.definition = definition;
        this.runtimesByBinding = Collections.unmodifiableMap(
            new LinkedHashMap<>(runtimesByBinding)
        );
        this.leases = List.copyOf(leases);
    }

    /**
     * 仅供同包编译器兼容测试借用既有单源 Runtime，不接管其生命周期。
     */
    static FederationQueryScopeSnapshot borrowedSingle(
        FederationQueryScopeDefinition definition,
        SourceRuntime runtime
    ) {
        return borrowed(definition, Map.of(definition.defaultBinding(), runtime));
    }

    /**
     * 仅供同包编译器契约测试借用既有多源 Runtime，不接管其生命周期。
     *
     * @param definition 查询范围
     * @param runtimesByBinding Binding 到 Runtime 的完整映射
     * @return 不持有 Runtime lease 的快照
     */
    static FederationQueryScopeSnapshot borrowed(
        FederationQueryScopeDefinition definition,
        Map<String, SourceRuntime> runtimesByBinding
    ) {
        if (definition == null || runtimesByBinding == null
            || !runtimesByBinding.keySet().containsAll(definition.bindings().keySet())) {
            throw new IllegalArgumentException(
                "borrowed snapshot must provide every declared binding runtime"
            );
        }
        return new FederationQueryScopeSnapshot(definition, runtimesByBinding, List.of());
    }

    /**
     * 按稳定 SourceId 顺序获取指定 Binding 的 Runtime，并合并重复物理源。
     *
     * @param sourceManager 数据源管理器
     * @param definition 查询范围
     * @param bindingNames 需要加载的 Binding 名称
     * @return Runtime 快照
     */
    static FederationQueryScopeSnapshot acquire(
        DefaultFederationSourceManager sourceManager,
        FederationQueryScopeDefinition definition,
        Set<String> bindingNames
    ) {
        if (bindingNames == null || bindingNames.isEmpty()) {
            throw new FederationSqlException(
                FederationSqlErrorCode.INVALID_ARGUMENT,
                "at least one query binding must be selected"
            );
        }
        Map<SourceId, Long> revisions = new LinkedHashMap<>();
        for (String bindingName : bindingNames) {
            FederationSourceBindingDefinition binding = definition.bindings().get(bindingName);
            if (binding == null) {
                throw new FederationSqlException(
                    FederationSqlErrorCode.INVALID_ARGUMENT,
                    "query references an unknown binding: " + bindingName
                );
            }
            revisions.merge(binding.sourceId(), binding.minimumRevision(), Math::max);
        }
        List<Map.Entry<SourceId, Long>> orderedSources = new ArrayList<>(revisions.entrySet());
        orderedSources.sort(Comparator.comparing(entry -> entry.getKey().value()));

        List<SourceRuntime.RuntimeLease> acquired = new ArrayList<>();
        Map<SourceId, SourceRuntime> runtimesBySource = new LinkedHashMap<>();
        try {
            for (Map.Entry<SourceId, Long> source : orderedSources) {
                SourceRuntime.RuntimeLease lease = sourceManager.acquireRuntime(
                    source.getKey(),
                    source.getValue()
                );
                acquired.add(lease);
                runtimesBySource.put(source.getKey(), lease.runtime());
            }
            Map<String, SourceRuntime> runtimesByBinding = new LinkedHashMap<>();
            bindingNames.stream().sorted().forEach(bindingName -> {
                FederationSourceBindingDefinition binding = definition.bindings().get(bindingName);
                runtimesByBinding.put(bindingName, runtimesBySource.get(binding.sourceId()));
            });
            return new FederationQueryScopeSnapshot(definition, runtimesByBinding, acquired);
        } catch (RuntimeException exception) {
            closeReverseAndSuppress(acquired, exception);
            throw exception;
        }
    }

    /**
     * 返回查询范围定义。
     *
     * @return 查询范围
     */
    FederationQueryScopeDefinition definition() {
        return definition;
    }

    /**
     * 返回指定 Binding 的 Runtime。
     *
     * @param bindingName Binding 名称
     * @return Runtime
     */
    SourceRuntime runtime(String bindingName) {
        SourceRuntime runtime = runtimesByBinding.get(bindingName);
        if (runtime == null) {
            throw new FederationSqlException(
                FederationSqlErrorCode.INVALID_ARGUMENT,
                "binding runtime is not loaded: " + bindingName
            );
        }
        return runtime;
    }

    /**
     * 返回已加载的 Binding 到 Runtime 映射。
     *
     * @return 不可变映射
     */
    Map<String, SourceRuntime> runtimesByBinding() {
        return runtimesByBinding;
    }

    /**
     * 返回全部 Binding 的运行身份。
     *
     * @return 运行身份列表
     */
    List<FederationSourceRuntimeIdentity> identities() {
        return runtimesByBinding.entrySet().stream()
            .map(entry -> identity(entry.getKey(), entry.getValue()))
            .toList();
    }

    private static FederationSourceRuntimeIdentity identity(
        String bindingName,
        SourceRuntime runtime
    ) {
        return new FederationSourceRuntimeIdentity(
            bindingName,
            runtime.definition().sourceId(),
            runtime.definition().revision(),
            runtime.sourceChecksum(),
            runtime.adapter().adapterId(),
            runtime.runtimeFingerprint()
        );
    }

    /**
     * 逆序释放所有唯一物理 Runtime lease。
     */
    @Override
    public void close() {
        if (!closed.compareAndSet(false, true)) {
            return;
        }
        FederationSqlException failure = null;
        for (int index = leases.size() - 1; index >= 0; index--) {
            try {
                leases.get(index).close();
            } catch (RuntimeException exception) {
                FederationSqlException wrapped = new FederationSqlException(
                    FederationSqlErrorCode.RESOURCE_CLOSE_FAILED,
                    "failed to close query scope runtime lease",
                    exception
                );
                if (failure == null) {
                    failure = wrapped;
                } else {
                    failure.addSuppressed(wrapped);
                }
            }
        }
        if (failure != null) {
            throw failure;
        }
    }

    private static void closeReverseAndSuppress(
        List<SourceRuntime.RuntimeLease> acquired,
        RuntimeException original
    ) {
        for (int index = acquired.size() - 1; index >= 0; index--) {
            try {
                acquired.get(index).close();
            } catch (RuntimeException closeFailure) {
                original.addSuppressed(closeFailure);
            }
        }
    }
}
