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.api.SqlExecutionContext;
import com.easyagents.federation.sql.execute.FederationFragmentExecutionContext;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.execute.SqlExecutionOptions;
import com.easyagents.federation.sql.execute.SqlParameter;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import com.easyagents.federation.sql.federation.FederationFragmentPlan;
import java.math.BigDecimal;
import java.sql.Date;
import java.sql.Time;
import java.sql.Timestamp;
import java.sql.Types;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.time.OffsetDateTime;
import java.time.ZoneOffset;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.calcite.avatica.util.ByteString;
import org.apache.calcite.linq4j.AbstractEnumerable;
import org.apache.calcite.linq4j.Enumerable;
import org.apache.calcite.linq4j.Enumerator;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.sql.type.SqlTypeName;

/**
 * 一次联邦查询的分片执行、类型桥接、预算和取消上下文。
 */
final class FederationExecutionSession implements AutoCloseable {

    static final String DATA_CONTEXT_KEY = "easyagents.federation.session";

    private final DefaultFederationSqlPlan plan;
    private final SqlExecutionContext context;
    private final FederationQueryScopeSnapshot runtimeSnapshot;
    private final QueryCancellationRegistry.QueryRegistration registration;
    private final FederationExecutionPolicy policy;
    private final FederationQueryMetricsTracker metrics;
    private final QueryDeadline deadline;
    private final Semaphore fragmentSlots;
    private final Set<FederationResultCursor> openFragmentCursors = ConcurrentHashMap.newKeySet();
    private final AtomicBoolean closed = new AtomicBoolean();

    /**
     * 创建查询执行会话。
     *
     * @param plan 联邦计划
     * @param context 执行上下文
     * @param runtimeSnapshot 实际引用源 Runtime 快照
     * @param registration 查询取消登记
     * @param policy 有效资源策略
     * @param metrics 指标跟踪器
     * @param deadline 请求级统一截止时间
     */
    FederationExecutionSession(
        DefaultFederationSqlPlan plan,
        SqlExecutionContext context,
        FederationQueryScopeSnapshot runtimeSnapshot,
        QueryCancellationRegistry.QueryRegistration registration,
        FederationExecutionPolicy policy,
        FederationQueryMetricsTracker metrics,
        QueryDeadline deadline
    ) {
        this.plan = plan;
        this.context = context;
        this.runtimeSnapshot = runtimeSnapshot;
        this.registration = registration;
        this.policy = policy;
        this.metrics = metrics;
        this.deadline = deadline;
        this.fragmentSlots = new Semaphore(policy.maximumConcurrentFragments(), true);
    }

    /**
     * 为 Calcite ScannableTable 创建一个按需执行物理分片的 Enumerable。
     *
     * @param fragmentId 分片标识
     * @return 分片行 Enumerable
     */
    Enumerable<Object[]> scan(String fragmentId) {
        CompiledFederationFragment compiled = plan.compiledFragments().get(fragmentId);
        if (compiled == null) {
            throw new FederationSqlException(
                FederationSqlErrorCode.EXECUTION_FAILED,
                "compiled fragment is missing: " + fragmentId
            );
        }
        return new AbstractEnumerable<>() {
            @Override
            public Enumerator<Object[]> enumerator() {
                return openFragment(compiled);
            }
        };
    }

    /**
     * 返回 Calcite DataContext 使用的查询参数值。
     *
     * @return 参数键值
     */
    Map<String, Object> dataContextValues() {
        Map<String, Object> values = new java.util.LinkedHashMap<>();
        values.put(DATA_CONTEXT_KEY, this);
        for (int index = 0; index < context.parameters().size(); index++) {
            values.put("?" + index, normalizeParameter(context.parameters().get(index)));
        }
        return java.util.Collections.unmodifiableMap(values);
    }

    /**
     * 在本地算子或分片读取前检查取消、关闭和总时限。
     */
    void ensureExecutionAllowed() {
        if (closed.get()) {
            throw new FederationSqlException(
                FederationSqlErrorCode.EXECUTION_FAILED,
                "federation execution session is closed"
            );
        }
        if (registration.cancellationRequested()) {
            if (registration.timeoutRequested()) {
                metrics.markTimedOut();
            } else {
                metrics.markCancelled();
            }
            registration.ensureNotCancelled();
        }
        deadline.ensureAllowed();
    }

    /**
     * 计量一个 Calcite 本地算子输出行并执行统一硬预算检查。
     *
     * @param operatorName 本地算子名称
     * @param value Calcite 内部行值
     * @param elapsedNanos 算子产生当前行的耗时
     */
    void recordLocalIntermediate(String operatorName, Object value, long elapsedNanos) {
        ensureExecutionAllowed();
        Object[] row = value instanceof Object[] values ? values : new Object[] {value};
        FederationQueryMetricsTracker.IntermediateUsage usage =
            metrics.recordLocalIntermediate(operatorName, row, elapsedNanos);
        if (usage.rows() > policy.maximumIntermediateRows()
            || usage.bytes() > policy.maximumIntermediateBytes()) {
            throw new FederationSqlException(
                FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
                "federation local operator " + operatorName
                    + " exceeded the intermediate result limit"
            );
        }
    }

    private Enumerator<Object[]> openFragment(CompiledFederationFragment compiled) {
        ensureExecutionAllowed();
        acquireFragmentSlot();
        FederationFragmentPlan fragment = compiled.plan();
        SourceRuntime runtime = runtimeSnapshot.runtime(fragment.bindingName());
        try {
            FederationResultCursor cursor = runtime.adapter().fragmentExecutor().execute(
                new FederationFragmentExecutionContext(
                    context.queryId(),
                    fragment.executableSql(),
                    remapParameters(fragment, context.parameters()),
                    fragmentOptions(context.options()),
                    runtime.handle().dataSource(),
                    runtime.compatibility(),
                    runtime.definition().adapterOptions(),
                    registration.statementLifecycle(),
                    metrics.fragmentObserver(fragment.fragmentId()),
                    deadline
                )
            );
            if (cursor == null || !context.queryId().equals(cursor.queryId())) {
                closeQuietly(cursor);
                throw new FederationSqlException(
                    FederationSqlErrorCode.EXECUTION_FAILED,
                    "adapter returned an invalid fragment cursor"
                );
            }
            openFragmentCursors.add(cursor);
            return new FragmentEnumerator(compiled, cursor);
        } catch (RuntimeException exception) {
            fragmentSlots.release();
            throw exception;
        }
    }

    private void acquireFragmentSlot() {
        while (true) {
            ensureExecutionAllowed();
            long remaining = deadline.remainingNanos();
            if (remaining <= 0L) {
                ensureExecutionAllowed();
            }
            try {
                if (fragmentSlots.tryAcquire(
                    Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(50)),
                    TimeUnit.NANOSECONDS
                )) {
                    return;
                }
            } catch (InterruptedException exception) {
                Thread.currentThread().interrupt();
                throw new FederationSqlException(
                    FederationSqlErrorCode.EXECUTION_FAILED,
                    "waiting for a federation fragment slot was interrupted",
                    exception
                );
            }
        }
    }

    private SqlExecutionOptions fragmentOptions(SqlExecutionOptions requested) {
        // 最终 maxRows 不能截断 Join/Aggregate 输入，否则会改变查询结果。
        return new SqlExecutionOptions(
            requested.fetchSize(),
            0,
            effectiveJdbcTimeoutSeconds(requested.queryTimeoutSeconds()),
            requested.readOnly()
        );
    }

    private int effectiveJdbcTimeoutSeconds(int requestedSeconds) {
        return deadline.boundedQueryTimeoutSeconds(requestedSeconds);
    }

    private static List<SqlParameter> remapParameters(
        FederationFragmentPlan fragment,
        List<SqlParameter> parameters
    ) {
        if (fragment.parameterMapping().isEmpty()) {
            return List.of();
        }
        List<SqlParameter> remapped = new ArrayList<>(fragment.parameterMapping().size());
        for (Integer index : fragment.parameterMapping()) {
            remapped.add(parameters.get(index));
        }
        return List.copyOf(remapped);
    }

    /**
     * 关闭所有仍活动的物理分片游标。
     */
    @Override
    public void close() {
        if (!closed.compareAndSet(false, true)) {
            return;
        }
        RuntimeException failure = null;
        for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) {
            try {
                cursor.close();
            } catch (RuntimeException exception) {
                if (failure == null) {
                    failure = exception;
                } else {
                    failure.addSuppressed(exception);
                }
            } finally {
                openFragmentCursors.remove(cursor);
            }
        }
        if (failure != null) {
            throw failure;
        }
    }

    private final class FragmentEnumerator implements Enumerator<Object[]> {

        private final CompiledFederationFragment compiled;
        private final FederationResultCursor cursor;
        private final AtomicBoolean released = new AtomicBoolean();
        private Object[] current;

        private FragmentEnumerator(
            CompiledFederationFragment compiled,
            FederationResultCursor cursor
        ) {
            this.compiled = compiled;
            this.cursor = cursor;
        }

        /** {@inheritDoc} */
        @Override
        public Object[] current() {
            return current;
        }

        /** {@inheritDoc} */
        @Override
        public boolean moveNext() {
            ensureExecutionAllowed();
            try {
                if (!cursor.next()) {
                    release(true);
                    return false;
                }
                Object[] row = normalizeRow(cursor, compiled.rowType());
                FederationQueryMetricsTracker.IntermediateUsage usage =
                    metrics.recordIntermediate(compiled.plan().fragmentId(), row);
                if (usage.rows() > policy.maximumIntermediateRows()
                    || usage.bytes() > policy.maximumIntermediateBytes()) {
                    throw new FederationSqlException(
                        FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
                        "federation intermediate result limit was exceeded"
                    );
                }
                current = row;
                return true;
            } catch (RuntimeException exception) {
                release(false);
                throw exception;
            }
        }

        /** {@inheritDoc} */
        @Override
        public void reset() {
            throw new UnsupportedOperationException("fragment cursor cannot be reset");
        }

        /** {@inheritDoc} */
        @Override
        public void close() {
            release(false);
        }

        private void release(boolean exhausted) {
            if (!released.compareAndSet(false, true)) {
                return;
            }
            openFragmentCursors.remove(cursor);
            try {
                cursor.close();
            } finally {
                fragmentSlots.release();
                if (exhausted) {
                    metrics.finishFragment(compiled.plan().fragmentId());
                }
            }
        }
    }

    private static Object[] normalizeRow(
        FederationResultCursor cursor,
        RelDataType rowType
    ) {
        List<RelDataTypeField> fields = rowType.getFieldList();
        Object[] values = new Object[fields.size()];
        for (int index = 0; index < fields.size(); index++) {
            values[index] = normalizeValue(
                cursor.getObject(index + 1),
                fields.get(index).getType().getSqlTypeName()
            );
        }
        return values;
    }

    private static Object normalizeValue(Object value, SqlTypeName typeName) {
        if (value == null) {
            return null;
        }
        return switch (typeName) {
            case BOOLEAN -> (Boolean) value;
            case TINYINT -> ((Number) value).byteValue();
            case SMALLINT -> ((Number) value).shortValue();
            case INTEGER -> ((Number) value).intValue();
            case BIGINT -> ((Number) value).longValue();
            case REAL -> ((Number) value).floatValue();
            case FLOAT, DOUBLE -> ((Number) value).doubleValue();
            case DECIMAL -> value instanceof BigDecimal
                ? value
                : new BigDecimal(value.toString());
            case DATE -> normalizeDate(value);
            case TIME, TIME_WITH_LOCAL_TIME_ZONE, TIME_TZ -> normalizeTime(value);
            case TIMESTAMP, TIMESTAMP_WITH_LOCAL_TIME_ZONE, TIMESTAMP_TZ ->
                normalizeTimestamp(value);
            case BINARY, VARBINARY -> value instanceof ByteString
                ? value
                : new ByteString((byte[]) value);
            case CHAR, VARCHAR -> value.toString();
            case UUID -> value instanceof java.util.UUID
                ? value
                : java.util.UUID.fromString(value.toString());
            case ANY, OTHER -> normalizeKnownUnknownScalar(value);
            case ARRAY, MULTISET, MAP, ROW -> throw new FederationSqlException(
                FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED,
                "complex JDBC values are not supported by local federation operators"
            );
            default -> throw new FederationSqlException(
                FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED,
                "JDBC type cannot be safely normalized for local federation: " + typeName
            );
        };
    }

    /**
     * 接受驱动元数据无法精确声明、但仍可稳定序列化和计量的标准 JDBC 标量。
     *
     * @param value JDBC 返回值
     * @return 可供本地 Enumerable 透传的值
     * @throws FederationSqlException 值不在安全白名单内
     */
    private static Object normalizeKnownUnknownScalar(Object value) {
        if (value instanceof java.time.OffsetTime offsetTime) {
            return offsetTime.withOffsetSameInstant(ZoneOffset.UTC);
        }
        if (value instanceof java.time.OffsetDateTime offsetDateTime) {
            return offsetDateTime.withOffsetSameInstant(ZoneOffset.UTC);
        }
        if (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
            || value instanceof java.util.UUID) {
            return value;
        }
        throw new FederationSqlException(
            FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED,
            "JDBC value cannot be safely normalized for local federation: "
                + value.getClass().getName()
        );
    }

    private static Integer normalizeDate(Object value) {
        LocalDate date = value instanceof Date sqlDate
            ? sqlDate.toLocalDate()
            : (LocalDate) value;
        return Math.toIntExact(date.toEpochDay());
    }

    private static Integer normalizeTime(Object value) {
        LocalTime time;
        if (value instanceof Time sqlTime) {
            time = sqlTime.toLocalTime();
        } else if (value instanceof java.time.OffsetTime offsetTime) {
            time = offsetTime.withOffsetSameInstant(ZoneOffset.UTC).toLocalTime();
        } else {
            time = (LocalTime) value;
        }
        requireMillisecondPrecision(time.getNano(), "TIME");
        return Math.toIntExact(time.toNanoOfDay() / 1_000_000L);
    }

    private static Long normalizeTimestamp(Object value) {
        if (value instanceof Timestamp timestamp) {
            requireMillisecondPrecision(timestamp.getNanos(), "TIMESTAMP");
            return timestamp.getTime();
        }
        if (value instanceof Instant instant) {
            requireMillisecondPrecision(instant.getNano(), "TIMESTAMP");
            return instant.toEpochMilli();
        }
        if (value instanceof OffsetDateTime offset) {
            requireMillisecondPrecision(offset.getNano(), "TIMESTAMP");
            return offset.toInstant().toEpochMilli();
        }
        LocalDateTime local = (LocalDateTime) value;
        requireMillisecondPrecision(local.getNano(), "TIMESTAMP");
        return local.toInstant(ZoneOffset.UTC).toEpochMilli();
    }

    /**
     * 拒绝 Calcite 毫秒内部表示无法保真的时间值。
     *
     * @param nanos 秒内纳秒值
     * @param typeName SQL 时间类型名
     * @throws FederationSqlException 时间值包含亚毫秒精度
     */
    private static void requireMillisecondPrecision(int nanos, String typeName) {
        if (nanos % 1_000_000 != 0) {
            throw new FederationSqlException(
                FederationSqlErrorCode.FEDERATION_OPERATOR_UNSUPPORTED,
                typeName + " values with precision above milliseconds are not supported "
                    + "by local federation"
            );
        }
    }

    private static Object normalizeParameter(SqlParameter parameter) {
        Object value = parameter.value();
        if (value == null) {
            return null;
        }
        return switch (parameter.jdbcType()) {
            case Types.DATE -> normalizeDate(value);
            case Types.TIME, Types.TIME_WITH_TIMEZONE -> normalizeTime(value);
            case Types.TIMESTAMP, Types.TIMESTAMP_WITH_TIMEZONE -> normalizeTimestamp(value);
            case Types.BINARY, Types.VARBINARY, Types.LONGVARBINARY ->
                value instanceof ByteString ? value : new ByteString((byte[]) value);
            default -> value;
        };
    }

    private static void closeQuietly(FederationResultCursor cursor) {
        if (cursor != null) {
            cursor.close();
        }
    }
}
