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.FederationColumn;
import com.easyagents.federation.sql.execute.FederationQueryMetricsSnapshot;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.execute.QueryId;
import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.io.Reader;
import java.io.StringReader;
import java.sql.Date;
import java.sql.Time;
import java.sql.Timestamp;
import java.time.Instant;
import java.time.LocalTime;
import java.time.OffsetDateTime;
import java.time.OffsetTime;
import java.time.ZoneOffset;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.calcite.avatica.util.ByteString;
import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.linq4j.Enumerator;

/**
 * 将 Calcite 本地 Bindable 暴露为统一流式结果游标。
 */
final class FederatedResultCursor
    implements FederationResultCursor, CancellationAwareFederationCursor {

    private final DefaultFederationSqlPlan plan;
    private final QueryId queryId;
    private final FederationExecutionSession session;
    private final FederationQueryMetricsTracker metrics;
    private final int maximumRows;
    private final Enumerator<?> enumerator;
    private final AtomicBoolean closed = new AtomicBoolean();
    private long rows;
    private Object[] current;

    /**
     * 创建联邦结果游标并绑定本次查询 DataContext。
     *
     * @param plan 联邦计划
     * @param queryId 查询标识
     * @param session 执行会话
     * @param metrics 指标跟踪器
     * @param maximumRows 最终结果上限，0 表示不额外限制
     */
    FederatedResultCursor(
        DefaultFederationSqlPlan plan,
        QueryId queryId,
        FederationExecutionSession session,
        FederationQueryMetricsTracker metrics,
        int maximumRows
    ) {
        this.plan = plan;
        this.queryId = queryId;
        this.session = session;
        this.metrics = metrics;
        this.maximumRows = maximumRows;
        this.enumerator = plan.localBindable()
            .bind(new FederationDataContext(
                plan.localRootSchema(),
                (JavaTypeFactory) plan.relRoot().rel.getCluster().getTypeFactory(),
                session.dataContextValues()
            ))
            .enumerator();
    }

    /** {@inheritDoc} */
    @Override
    public QueryId queryId() {
        return queryId;
    }

    /** {@inheritDoc} */
    @Override
    public List<FederationColumn> columns() {
        return plan.columns();
    }

    /** {@inheritDoc} */
    @Override
    public FederationQueryMetricsSnapshot metrics() {
        return metrics.snapshot();
    }

    /** {@inheritDoc} */
    @Override
    public void markCancelled() {
        metrics.markCancelled();
    }

    /** {@inheritDoc} */
    @Override
    public void markTimedOut() {
        metrics.markTimedOut();
    }

    /** {@inheritDoc} */
    @Override
    public void markFailed(FederationSqlErrorCode errorCode) {
        metrics.markFailed(errorCode);
    }

    /** {@inheritDoc} */
    @Override
    public boolean next() {
        ensureOpen();
        session.ensureExecutionAllowed();
        if (maximumRows > 0 && rows >= maximumRows) {
            metrics.markTruncated();
            finish();
            return false;
        }
        try {
            if (!enumerator.moveNext()) {
                finish();
                return false;
            }
            Object value = enumerator.current();
            Object[] row = value instanceof Object[] values ? values : new Object[] {value};
            current = denormalize(row, plan.columns());
            rows++;
            metrics.recordOutput(current);
            return true;
        } catch (RuntimeException exception) {
            if (exception instanceof FederationSqlException sqlException
                && sqlException.errorCode() == FederationSqlErrorCode.QUERY_CANCELLED) {
                metrics.markCancelled();
            }
            closeAndSuppress(exception);
            throw exception;
        }
    }

    /** {@inheritDoc} */
    @Override
    public Object getObject(int columnIndex) {
        ensureCurrent(columnIndex);
        return current[columnIndex - 1];
    }

    /** {@inheritDoc} */
    @Override
    public InputStream getBinaryStream(int columnIndex) {
        Object value = getObject(columnIndex);
        if (value == null) {
            return null;
        }
        if (value instanceof byte[] bytes) {
            return new ByteArrayInputStream(bytes);
        }
        throw new UnsupportedOperationException("column is not a binary value");
    }

    /** {@inheritDoc} */
    @Override
    public Reader getCharacterStream(int columnIndex) {
        Object value = getObject(columnIndex);
        return value == null ? null : new StringReader(value.toString());
    }

    /** {@inheritDoc} */
    @Override
    public List<Object> row() {
        if (current == null) {
            throw new IllegalStateException("cursor is not positioned on a row");
        }
        return java.util.Collections.unmodifiableList(Arrays.asList(current.clone()));
    }

    /**
     * 关闭本地 Enumerator 和全部分片游标。
     */
    @Override
    public void close() {
        if (!closed.compareAndSet(false, true)) {
            return;
        }
        RuntimeException failure = null;
        try {
            try {
                enumerator.close();
            } catch (RuntimeException exception) {
                failure = exception;
            }
            try {
                session.close();
            } catch (RuntimeException exception) {
                if (failure == null) {
                    failure = exception;
                } else {
                    failure.addSuppressed(exception);
                }
            }
        } finally {
            metrics.finishClosed();
        }
        if (failure != null) {
            throw failure;
        }
    }

    private void finish() {
        metrics.finishSuccessfully();
        close();
    }

    private void ensureOpen() {
        if (closed.get()) {
            throw new IllegalStateException("cursor is closed");
        }
    }

    private void ensureCurrent(int columnIndex) {
        if (current == null) {
            throw new IllegalStateException("cursor is not positioned on a row");
        }
        if (columnIndex <= 0 || columnIndex > current.length) {
            throw new IndexOutOfBoundsException("column index out of range: " + columnIndex);
        }
    }

    private void closeAndSuppress(RuntimeException original) {
        try {
            close();
        } catch (RuntimeException closeFailure) {
            original.addSuppressed(closeFailure);
        }
    }

    private static Object[] denormalize(
        Object[] values,
        List<FederationColumn> columns
    ) {
        Object[] converted = values.clone();
        for (int index = 0; index < converted.length; index++) {
            Object value = converted[index];
            if (value == null) {
                continue;
            }
            converted[index] = switch (columns.get(index).jdbcType()) {
                case java.sql.Types.DATE -> value instanceof Integer days
                    ? Date.valueOf(java.time.LocalDate.ofEpochDay(days))
                    : value;
                case java.sql.Types.TIME -> value instanceof Integer millis
                    ? Time.valueOf(LocalTime.ofNanoOfDay(millis * 1_000_000L))
                    : value;
                case java.sql.Types.TIME_WITH_TIMEZONE -> value instanceof Integer millis
                    ? OffsetTime.of(
                        LocalTime.ofNanoOfDay(millis * 1_000_000L),
                        ZoneOffset.UTC
                    )
                    : value;
                case java.sql.Types.TIMESTAMP -> value instanceof Long millis
                    ? new Timestamp(millis)
                    : value;
                case java.sql.Types.TIMESTAMP_WITH_TIMEZONE -> value instanceof Long millis
                    ? OffsetDateTime.ofInstant(Instant.ofEpochMilli(millis), ZoneOffset.UTC)
                    : value;
                case java.sql.Types.BINARY, java.sql.Types.VARBINARY,
                     java.sql.Types.LONGVARBINARY -> value instanceof ByteString bytes
                    ? bytes.getBytes()
                    : value;
                default -> value;
            };
        }
        return converted;
    }
}
