/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.paimon.lumina.index;

import org.apache.paimon.data.InternalArray;
import org.apache.paimon.data.InternalVector;
import org.apache.paimon.fs.PositionOutputStream;
import org.apache.paimon.globalindex.GlobalIndexSingleColumnWriter;
import org.apache.paimon.globalindex.ResultEntry;
import org.apache.paimon.globalindex.io.GlobalIndexFileWriter;
import org.apache.paimon.types.ArrayType;
import org.apache.paimon.types.DataType;
import org.apache.paimon.types.FloatType;
import org.apache.paimon.types.VectorType;

import org.aliyun.lumina.LuminaDataset;
import org.aliyun.lumina.LuminaFileOutput;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.lang.reflect.Field;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.FileChannel;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Vector global index writer using Lumina. Builds a single index file per shard.
 *
 * <p>Vectors are spilled to a temporary file on disk as they arrive via {@link #write(Object,
 * long)}, keeping Java heap usage constant (~8 MB buffer) regardless of dataset size. During index
 * build, the Lumina builder streams vectors from the temp file via the {@link LuminaDataset}
 * callback API.
 *
 * <p><b>Thread safety:</b> This class is <b>not</b> thread-safe. The underlying {@code
 * LuminaBuilder} must be used from a single thread or the caller must provide external
 * synchronization. The internal Lumina executor thread pool size is controlled globally by the
 * {@code LUMINA_EXECUTOR_THREAD_COUNT} environment variable and cannot be configured per-instance.
 */
public class LuminaVectorGlobalIndexWriter implements GlobalIndexSingleColumnWriter, Closeable {

    private static final String FILE_NAME_PREFIX = "lumina";

    private static final Logger LOG = LoggerFactory.getLogger(LuminaVectorGlobalIndexWriter.class);

    /**
     * Tracks which environment variables have already been set in this JVM. Used to avoid
     * concurrent reflection-based modification of {@code System.getenv()} when multiple shard tasks
     * run in parallel on the same TaskManager.
     */
    private static final ConcurrentHashMap<String, String> CONFIGURED_ENV_VARS =
            new ConcurrentHashMap<>();

    /** I/O buffer size for reading/writing the temp vector file (~8 MB). */
    private static final int IO_BUFFER_SIZE = 8 * 1024 * 1024;

    private final GlobalIndexFileWriter fileWriter;
    private final LuminaVectorIndexOptions options;
    private final Map<String, String> luminaOptions;
    private final int dim;

    /** Temporary file storing records as [rowId (long)][float * dim] in native byte order. */
    private File tempVectorFile;

    private FileChannel writeChannel;
    private ByteBuffer writeBuf;

    private final int recordSizeInBytes;
    private final float[] vectorBuf;
    private long count;
    private boolean closed;

    private long rowCount;

    public LuminaVectorGlobalIndexWriter(
            GlobalIndexFileWriter fileWriter,
            DataType fieldType,
            LuminaVectorIndexOptions options) {
        this.fileWriter = fileWriter;
        this.options = options;
        this.dim = validateAndResolveDimension(fieldType, options);
        this.luminaOptions = options.toLuminaOptions(dim);
        this.count = 0;
        this.closed = false;
        this.recordSizeInBytes = checkedRecordSize(dim, IO_BUFFER_SIZE);
        this.vectorBuf = new float[dim];

        try {
            this.tempVectorFile = File.createTempFile("lumina-vectors-", ".bin");
            this.tempVectorFile.deleteOnExit();
            @SuppressWarnings("resource")
            RandomAccessFile raf = new RandomAccessFile(tempVectorFile, "rw");
            this.writeChannel = raf.getChannel();
            this.writeBuf = ByteBuffer.allocateDirect(IO_BUFFER_SIZE);
            this.writeBuf.order(ByteOrder.nativeOrder());
        } catch (IOException e) {
            throw new RuntimeException("Failed to create temp vector file", e);
        }
    }

    private static int validateAndResolveDimension(
            DataType dataType, LuminaVectorIndexOptions options) {
        if (dataType instanceof VectorType) {
            VectorType vectorType = (VectorType) dataType;
            DataType elementType = vectorType.getElementType();
            if (!(elementType instanceof FloatType)) {
                throw new IllegalArgumentException(
                        "Lumina vector index requires float vector, but got: " + elementType);
            }
            int typeDimension = vectorType.getLength();
            if (options.isDimensionExplicitlyConfigured() && options.dimension() != typeDimension) {
                throw new IllegalArgumentException(
                        String.format(
                                "%s configured %d conflicts with VECTOR length %d.",
                                LuminaVectorIndexOptions.DIMENSION.key(),
                                options.dimension(),
                                typeDimension));
            }
            return typeDimension;
        }
        if (dataType instanceof ArrayType) {
            DataType elementType = ((ArrayType) dataType).getElementType();
            if (!(elementType instanceof FloatType)) {
                throw new IllegalArgumentException(
                        "Lumina vector index requires float array, but got: " + elementType);
            }
            return options.dimension();
        }
        throw new IllegalArgumentException(
                "Lumina vector index requires VectorType or ArrayType<FLOAT>, but got: "
                        + dataType);
    }

    @Override
    public void write(Object fieldData, long relativeRowId) {
        if (fieldData == null) {
            rowCount++;
            return;
        }

        // Validation must complete before any buffer/state mutation below
        float[] src = materializeAndValidate(fieldData, relativeRowId);

        if (writeBuf.remaining() < recordSizeInBytes) {
            flushWriteBuffer();
        }
        writeBuf.putLong(relativeRowId);
        for (int i = 0; i < dim; i++) {
            writeBuf.putFloat(src[i]);
        }
        rowCount++;
        count++;
    }

    /**
     * Validates the vector and returns a float[] ready to write. For float[] input, returns the
     * input directly (zero-copy). For InternalVector/InternalArray, reads into the reusable
     * vectorBuf field.
     */
    private float[] materializeAndValidate(Object fieldData, long relativeRowId) {
        if (fieldData instanceof float[]) {
            float[] vector = (float[]) fieldData;
            checkDimension(vector.length);
            for (int i = 0; i < dim; i++) {
                checkFinite(vector[i], relativeRowId, i);
            }
            return vector;
        } else if (fieldData instanceof InternalVector) {
            InternalVector vector = (InternalVector) fieldData;
            checkDimension(vector.size());
            for (int i = 0; i < dim; i++) {
                float v = vector.getFloat(i);
                checkFinite(v, relativeRowId, i);
                vectorBuf[i] = v;
            }
            return vectorBuf;
        } else if (fieldData instanceof InternalArray) {
            InternalArray array = (InternalArray) fieldData;
            checkDimension(array.size());
            for (int i = 0; i < dim; i++) {
                if (array.isNullAt(i)) {
                    throw new IllegalArgumentException("Vector element at index " + i + " is null");
                }
                float v = array.getFloat(i);
                checkFinite(v, relativeRowId, i);
                vectorBuf[i] = v;
            }
            return vectorBuf;
        } else {
            throw new RuntimeException(
                    "Unsupported vector type: " + fieldData.getClass().getName());
        }
    }

    private void flushWriteBuffer() {
        try {
            writeBuf.flip();
            while (writeBuf.hasRemaining()) {
                writeChannel.write(writeBuf);
            }
            writeBuf.clear();
        } catch (IOException e) {
            throw new RuntimeException("Failed to flush vector buffer to disk", e);
        }
    }

    @Override
    public List<ResultEntry> finish() {
        try {
            if (count == 0) {
                writeChannel.close();
                writeChannel = null;
                writeBuf = null;
                return Collections.emptyList();
            }
            // Flush remaining data and close the write channel
            flushWriteBuffer();
            writeChannel.close();
            writeChannel = null;
            writeBuf = null;
            return Collections.singletonList(buildIndex());
        } catch (IOException e) {
            throw new RuntimeException("Failed to write Lumina vector global index", e);
        } finally {
            // Delete the temp file eagerly so index builds don't exhaust disk.
            // Production callers (e.g. DefaultGlobalIndexBuilder) call finish()
            // but never close(), so we must clean up here.
            if (tempVectorFile != null) {
                tempVectorFile.delete();
                tempVectorFile = null;
            }
        }
    }

    private ResultEntry buildIndex() throws IOException {
        configureExecutorThreadCount();
        LOG.info(
                "Lumina index build started: {} vectors, dim={}, type={}, metric={}",
                count,
                dim,
                options.indexType(),
                options.metric());
        long buildStart = System.currentTimeMillis();

        try (LuminaIndex index =
                LuminaIndex.createForBuild(
                        options.indexType(), dim, options.metric(), luminaOptions)) {

            // Pretrain and insert via streaming file-backed Dataset API
            long phaseStart = System.currentTimeMillis();
            LOG.info("Lumina pretrain phase started");
            try (FileBackedDataset ds =
                    new FileBackedDataset(tempVectorFile, dim, count, "pretrain")) {
                index.pretrainFrom(ds);
            }
            LOG.info(
                    "Lumina pretrain phase done in {} ms", System.currentTimeMillis() - phaseStart);

            phaseStart = System.currentTimeMillis();
            LOG.info("Lumina insert phase started");
            try (FileBackedDataset ds =
                    new FileBackedDataset(tempVectorFile, dim, count, "insert")) {
                index.insertFrom(ds);
            }
            LOG.info("Lumina insert phase done in {} ms", System.currentTimeMillis() - phaseStart);

            phaseStart = System.currentTimeMillis();
            LOG.info("Lumina dump phase started");
            String fileName = fileWriter.newFileName(FILE_NAME_PREFIX);
            try (PositionOutputStream out = fileWriter.newOutputStream(fileName)) {
                index.dump(new OutputStreamFileOutput(out));
                out.flush();
            }
            LOG.info("Lumina dump phase done in {} ms", System.currentTimeMillis() - phaseStart);

            LOG.info(
                    "Lumina index build completed in {} ms",
                    System.currentTimeMillis() - buildStart);

            LuminaIndexMeta meta = new LuminaIndexMeta(luminaOptions);
            // rowCount = logical rows including nulls (not just indexed vectors)
            return new ResultEntry(fileName, rowCount, meta.serialize());
        }
    }

    /**
     * Sets the {@code LUMINA_EXECUTOR_THREAD_COUNT} environment variable from the {@link
     * LuminaVectorIndexOptions#DISKANN_BUILD_THREAD_COUNT} option so the native Lumina executor
     * thread pool is sized correctly before the builder is created.
     */
    private void configureExecutorThreadCount() {
        Map<String, String> luminaOpts = luminaOptions;
        String threadCountKey =
                LuminaVectorIndexOptions.toLuminaKey(
                        LuminaVectorIndexOptions.DISKANN_BUILD_THREAD_COUNT);
        String threadCount = luminaOpts.get(threadCountKey);
        if (threadCount != null && !threadCount.isEmpty()) {
            setEnv("LUMINA_EXECUTOR_THREAD_COUNT", threadCount);
        }
    }

    /**
     * Best-effort, thread-safe, idempotent attempt to set an environment variable in the current
     * process. Uses reflection to modify the JVM's process environment map. This is necessary
     * because Lumina's native executor thread pool reads {@code LUMINA_EXECUTOR_THREAD_COUNT} from
     * the OS environment.
     *
     * <p>When multiple shard tasks run in parallel on the same TaskManager JVM, this method ensures
     * the reflection-based map modification happens at most once per key-value pair.
     */
    @SuppressWarnings("unchecked")
    private static void setEnv(String key, String value) {
        // Skip if already set to the same value in this JVM
        if (value.equals(CONFIGURED_ENV_VARS.get(key))) {
            return;
        }
        synchronized (CONFIGURED_ENV_VARS) {
            // Double-check after acquiring lock
            if (value.equals(CONFIGURED_ENV_VARS.get(key))) {
                return;
            }
            try {
                Map<String, String> env = System.getenv();
                Field field = env.getClass().getDeclaredField("m");
                field.setAccessible(true);
                ((Map<String, String>) field.get(env)).put(key, value);
                CONFIGURED_ENV_VARS.put(key, value);
            } catch (Exception e) {
                LOG.warn(
                        "Failed to set environment variable '{}' for Lumina executor. "
                                + "Thread-count tuning may not take effect.",
                        key,
                        e);
            }
        }
    }

    private void checkDimension(int actualDim) {
        if (actualDim != dim) {
            throw new IllegalArgumentException(
                    String.format(
                            "Vector dimension mismatch: expected %d, but got %d", dim, actualDim));
        }
    }

    private void checkFinite(float value, long relativeRowId, int elementIndex) {
        if (!Float.isFinite(value)) {
            throw new IllegalArgumentException(
                    String.format(
                            "Vector element at rowId=%d, index=%d is %s",
                            relativeRowId, elementIndex, Float.toString(value)));
        }
    }

    private static int checkedRecordSize(int dim, int bufferCapacity) {
        long recordSize = Long.BYTES + (long) dim * Float.BYTES;
        if (recordSize > bufferCapacity || recordSize > Integer.MAX_VALUE) {
            throw new IllegalStateException(
                    "Vector record size "
                            + recordSize
                            + " exceeds buffer capacity "
                            + bufferCapacity);
        }
        return (int) recordSize;
    }

    @Override
    public void close() {
        if (!closed) {
            closed = true;
            try {
                if (writeChannel != null) {
                    writeChannel.close();
                }
            } catch (IOException ignored) {
            }
            writeBuf = null;
            if (tempVectorFile != null) {
                tempVectorFile.delete();
                tempVectorFile = null;
            }
        }
    }

    /**
     * A {@link LuminaDataset} that streams vectors from a temporary file. Reads are sequential and
     * buffered for efficiency. Each instance reads the file from the beginning independently.
     */
    static class FileBackedDataset implements LuminaDataset, Closeable {
        private final RandomAccessFile raf;
        private final FileChannel channel;
        private final int dim;
        private final long totalCount;
        private final int recordSizeInBytes;
        private long cursor;
        private final ByteBuffer readBuf;
        private final String phase;
        private int lastLoggedPercent;

        FileBackedDataset(File file, int dim, long totalCount, String phase) throws IOException {
            this(file, dim, totalCount, phase, IO_BUFFER_SIZE);
        }

        FileBackedDataset(File file, int dim, long totalCount, String phase, int bufferSize)
                throws IOException {
            this.raf = new RandomAccessFile(file, "r");
            this.channel = raf.getChannel();
            this.dim = dim;
            this.totalCount = totalCount;
            this.recordSizeInBytes = checkedRecordSize(dim, bufferSize);
            this.cursor = 0;
            this.readBuf = ByteBuffer.allocateDirect(bufferSize);
            this.readBuf.order(ByteOrder.nativeOrder());
            this.readBuf.limit(0); // empty initially
            this.phase = phase;
            this.lastLoggedPercent = -1;
        }

        @Override
        public int dim() {
            return dim;
        }

        @Override
        public long totalSize() {
            return totalCount;
        }

        @Override
        public long getNextBatch(float[] vectorBuf, long[] idBuf) {
            if (cursor >= totalCount) {
                return 0;
            }
            long remaining = totalCount - cursor;
            int maxByVectorBuf = vectorBuf.length / dim;
            int batchSize = (int) Math.min(Math.min(maxByVectorBuf, idBuf.length), remaining);

            try {
                for (int i = 0; i < batchSize; i++) {
                    ensureAvailable(recordSizeInBytes);
                    idBuf[i] = readBuf.getLong();
                    int base = i * dim;
                    for (int d = 0; d < dim; d++) {
                        vectorBuf[base + d] = readBuf.getFloat();
                    }
                }
            } catch (IOException e) {
                throw new RuntimeException("Failed to read vectors from temp file", e);
            }

            cursor += batchSize;

            int percent = (int) (cursor * 100 / totalCount);
            if (percent / 10 > lastLoggedPercent / 10) {
                LOG.info(
                        "Lumina {} progress: {}/{} vectors ({}%)",
                        phase, cursor, totalCount, percent);
                lastLoggedPercent = percent;
            }

            return batchSize;
        }

        private void ensureAvailable(int minBytes) throws IOException {
            int zeroReadCount = 0;
            while (readBuf.remaining() < minBytes) {
                readBuf.compact();
                int bytesRead = channel.read(readBuf);
                readBuf.flip();
                if (bytesRead == -1) {
                    throw new IOException("Unexpected end of temp file");
                }
                if (bytesRead == 0) {
                    if (++zeroReadCount > 100) {
                        throw new IOException(
                                "Unable to read from temp file: repeated zero-byte reads");
                    }
                } else {
                    zeroReadCount = 0;
                }
            }
        }

        @Override
        public void close() throws IOException {
            channel.close();
            raf.close();
        }
    }

    /** Adapts a {@link PositionOutputStream} to the {@link LuminaFileOutput} JNI callback API. */
    static class OutputStreamFileOutput implements LuminaFileOutput {
        private final PositionOutputStream out;

        OutputStreamFileOutput(PositionOutputStream out) {
            this.out = out;
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            out.write(b, off, len);
        }

        @Override
        public void flush() throws IOException {
            out.flush();
        }

        @Override
        public long getPos() throws IOException {
            return out.getPos();
        }

        @Override
        public void close() {}
    }
}
