/*
 * 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.append;

import org.apache.paimon.CoreOptions;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.BinaryVector;
import org.apache.paimon.data.BlobData;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.fileindex.FileIndexOptions;
import org.apache.paimon.fileindex.FileIndexPredicate;
import org.apache.paimon.fileindex.bitmap.BitmapIndexResult;
import org.apache.paimon.format.FileFormat;
import org.apache.paimon.format.blob.BlobFileFormat;
import org.apache.paimon.fs.ByteArraySeekableStream;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.local.LocalFileIO;
import org.apache.paimon.io.BundleRecords;
import org.apache.paimon.io.DataFileMeta;
import org.apache.paimon.io.DataFilePathFactory;
import org.apache.paimon.manifest.FileSource;
import org.apache.paimon.operation.BlobFileContext;
import org.apache.paimon.options.Options;
import org.apache.paimon.predicate.PredicateBuilder;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.LongCounter;
import org.apache.paimon.utils.RoaringBitmap32;
import org.apache.paimon.utils.StatsCollectorFactories;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import java.io.IOException;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.stream.Stream;

import static org.apache.paimon.types.VectorType.isVectorStoreFile;
import static org.assertj.core.api.Assertions.assertThat;

/** Tests for {@link DedicatedFormatRollingFileWriter} with vector-store. */
public class DedicatedFormatRollingFileWriterVectorTest {

    private static final int VECTOR_DIM = 12;
    private static final RowType SCHEMA =
            RowType.builder()
                    .field("f0", DataTypes.INT())
                    .field("f1", DataTypes.STRING())
                    .field("f2", DataTypes.BLOB())
                    .field("f3", DataTypes.VECTOR(VECTOR_DIM, DataTypes.FLOAT()))
                    .field("f4", DataTypes.INT())
                    .build();

    private static final long TARGET_FILE_SIZE = 2 * 1024 * 1024L; // 2 MB
    private static final long VECTOR_TARGET_FILE_SIZE = 4 * 1024 * 1024L; // 4 MB
    private static final long SCHEMA_ID = 1L;
    private static final String COMPRESSION = "none";
    private static final Random RANDOM = new Random(System.currentTimeMillis());

    @TempDir java.nio.file.Path tempDir;

    private DedicatedFormatRollingFileWriter writer;
    private DataFilePathFactory pathFactory;
    private LongCounter seqNumCounter;

    @BeforeEach
    public void setUp() throws IOException {
        // Setup file system and path factory
        LocalFileIO fileIO = LocalFileIO.create();
        pathFactory =
                new DataFilePathFactory(
                        new Path(tempDir + "/bucket-0"),
                        "parquet",
                        "data-", // dataFilePrefix should include the hyphen to match expected
                        // format: data-{uuid}-{count}
                        "changelog",
                        false,
                        null,
                        null);
        seqNumCounter = new LongCounter();

        // Initialize the writer
        writer =
                new DedicatedFormatRollingFileWriter(
                        fileIO,
                        SCHEMA_ID,
                        FileFormat.fromIdentifier("parquet", new Options()),
                        FileFormat.fromIdentifier("json", new Options()),
                        TARGET_FILE_SIZE,
                        TARGET_FILE_SIZE,
                        VECTOR_TARGET_FILE_SIZE,
                        Long.MAX_VALUE,
                        SCHEMA,
                        pathFactory,
                        () -> seqNumCounter,
                        COMPRESSION,
                        new StatsCollectorFactories(new CoreOptions(new Options())),
                        new FileIndexOptions(),
                        FileSource.APPEND,
                        false,
                        BlobFileContext.create(SCHEMA, new CoreOptions(new Options())));
    }

    @Test
    public void testBasicWriting() throws IOException {
        // Write a single row
        writer.write(makeRows(1, 10).get(0));
        assertThat(writer.recordCount()).isEqualTo(1);
    }

    @Test
    public void testMultipleWrites() throws Exception {
        // Write multiple rows
        int rowNum = RANDOM.nextInt(64) + 1;
        writer.write(makeRows(rowNum, 10).iterator());
        writer.close();
        List<DataFileMeta> metasResult = writer.result();

        assertThat(metasResult.size()).isEqualTo(3); // blob is small, normal/blob/vector 3 files
        assertThat(metasResult.get(0).fileFormat()).isEqualTo("parquet");
        assertThat(metasResult.get(1).fileFormat()).isEqualTo("blob");
        assertThat(metasResult.get(2).fileFormat()).isEqualTo("json");
        assertThat(writer.recordCount()).isEqualTo(rowNum);

        assertThat(metasResult.get(0).rowCount()).isEqualTo(metasResult.get(1).rowCount());
        assertThat(metasResult.get(0).rowCount()).isEqualTo(metasResult.get(2).rowCount());
    }

    @Test
    public void testAbortAfterRollingByRowsDeletesVectorFiles() throws IOException {
        RowType schema =
                RowType.builder()
                        .field("f0", DataTypes.INT())
                        .field("f1", DataTypes.VECTOR(VECTOR_DIM, DataTypes.FLOAT()))
                        .build();
        long hugeSize = 1024L * 1024 * 1024;
        writer =
                new DedicatedFormatRollingFileWriter(
                        LocalFileIO.create(),
                        SCHEMA_ID,
                        FileFormat.fromIdentifier("parquet", new Options()),
                        FileFormat.fromIdentifier("json", new Options()),
                        hugeSize,
                        hugeSize,
                        hugeSize,
                        1L,
                        schema,
                        pathFactory,
                        LongCounter::new,
                        COMPRESSION,
                        new StatsCollectorFactories(new CoreOptions(new Options())),
                        new FileIndexOptions(),
                        FileSource.APPEND,
                        false,
                        null);

        writer.write(GenericRow.of(1, BinaryVector.fromPrimitiveArray(new float[VECTOR_DIM])));
        writer.abort();

        try (Stream<java.nio.file.Path> files = Files.walk(tempDir)) {
            assertThat(files.filter(Files::isRegularFile).count()).isZero();
        }
    }

    @Test
    public void testVectorTargetFileSize() throws Exception {
        // 100k vector-store data would create 1 normal, 1 blob, and 3 vector-store files
        int rowNum = 100 * 1000;
        writer.write(makeRows(rowNum, 1).iterator());
        writer.close();
        List<DataFileMeta> results = writer.result();

        // Verify that we have multiple files due to rolling
        assertThat(results.size()).isGreaterThan(1);

        // Check that vector-store files meet the target size requirement
        List<DataFileMeta> vectorStoreFiles =
                results.stream()
                        .filter(file -> isVectorStoreFile(file.fileName()))
                        .collect(java.util.stream.Collectors.toList());

        assertThat(vectorStoreFiles.size()).isEqualTo(3);

        // Verify that vector-store files are close to the target size (within reasonable tolerance)
        for (DataFileMeta file : vectorStoreFiles.subList(0, vectorStoreFiles.size() - 1)) {
            long fileSize = file.fileSize();
            assertThat(fileSize)
                    .as("Vector-store file size should be close to target size")
                    .isGreaterThanOrEqualTo(VECTOR_TARGET_FILE_SIZE)
                    .isLessThanOrEqualTo(VECTOR_TARGET_FILE_SIZE + 256 * 1024);
        }

        // Verify total record count
        assertThat(writer.recordCount()).isEqualTo(rowNum);
    }

    @Test
    public void testBundleWriting() throws Exception {
        List<InternalRow> rows = makeRows(8, 10);
        writer.writeBundle(new SingleUseBundleRecords(rows));
        writer.close();
        List<DataFileMeta> metasResult = writer.result();

        assertThat(writer.recordCount()).isEqualTo(rows.size());
        assertThat(metasResult).hasSize(3);
        assertThat(metasResult.get(0).rowCount()).isEqualTo(rows.size());
        assertThat(metasResult.get(1).rowCount()).isEqualTo(rows.size());
        assertThat(metasResult.get(2).rowCount()).isEqualTo(rows.size());
        assertBundleSequenceSideEffects(metasResult, rows.size());
    }

    @Test
    public void testBundleWritingRespectsVectorTargetFileSize() throws Exception {
        long vectorTargetFileSize = 16 * 1024L;
        writer =
                new DedicatedFormatRollingFileWriter(
                        LocalFileIO.create(),
                        SCHEMA_ID,
                        FileFormat.fromIdentifier("parquet", new Options()),
                        FileFormat.fromIdentifier("json", new Options()),
                        128 * 1024 * 1024,
                        128 * 1024 * 1024,
                        vectorTargetFileSize,
                        Long.MAX_VALUE,
                        SCHEMA,
                        new DataFilePathFactory(
                                new Path(tempDir + "/bundle-vector-size-test"),
                                "parquet",
                                "data-",
                                "changelog",
                                false,
                                null,
                                null),
                        () -> new LongCounter(),
                        COMPRESSION,
                        new StatsCollectorFactories(new CoreOptions(new Options())),
                        new FileIndexOptions(),
                        FileSource.APPEND,
                        false,
                        BlobFileContext.create(SCHEMA, new CoreOptions(new Options())));

        List<InternalRow> rows = makeRows(2000, 1);
        writer.writeBundle(new SingleUseBundleRecords(rows));
        writer.close();

        List<DataFileMeta> vectorStoreFiles =
                writer.result().stream()
                        .filter(file -> isVectorStoreFile(file.fileName()))
                        .collect(java.util.stream.Collectors.toList());
        assertThat(vectorStoreFiles)
                .as("Bundle writes should still roll vector files inside the bundle.")
                .hasSizeGreaterThan(1);
        for (DataFileMeta file : vectorStoreFiles.subList(0, vectorStoreFiles.size() - 1)) {
            assertThat(file.fileSize()).isGreaterThanOrEqualTo(vectorTargetFileSize);
        }
        assertThat(writer.recordCount()).isEqualTo(rows.size());
    }

    @Test
    public void testBundleWritingPreservesMainFileIndexSideEffects() throws Exception {
        Options options = new Options();
        options.set("file-index.bitmap.columns", "f0");
        options.set("file-index.in-manifest-threshold", "1 MB");
        CoreOptions coreOptions = new CoreOptions(options);
        writer =
                new DedicatedFormatRollingFileWriter(
                        LocalFileIO.create(),
                        SCHEMA_ID,
                        FileFormat.fromIdentifier("parquet", new Options()),
                        FileFormat.fromIdentifier("json", new Options()),
                        TARGET_FILE_SIZE,
                        TARGET_FILE_SIZE,
                        VECTOR_TARGET_FILE_SIZE,
                        Long.MAX_VALUE,
                        SCHEMA,
                        pathFactory,
                        () -> seqNumCounter,
                        COMPRESSION,
                        new StatsCollectorFactories(coreOptions),
                        new FileIndexOptions(coreOptions),
                        FileSource.APPEND,
                        false,
                        BlobFileContext.create(SCHEMA, coreOptions));

        List<InternalRow> rows = makeRows(4, 10);
        writer.writeBundle(new SingleUseBundleRecords(rows));
        writer.close();

        DataFileMeta mainFile =
                writer.result().stream()
                        .filter(file -> "parquet".equals(file.fileFormat()))
                        .findFirst()
                        .get();

        assertThat(mainFile.rowCount()).isEqualTo(rows.size());
        assertThat(mainFile.embeddedIndex()).isNotEmpty();
        assertThat(mainFile.extraFiles()).isEmpty();

        RowType normalFileSchema =
                RowType.builder()
                        .field("f0", DataTypes.INT())
                        .field("f1", DataTypes.STRING())
                        .field("f4", DataTypes.INT())
                        .build();
        PredicateBuilder predicateBuilder = new PredicateBuilder(normalFileSchema);
        try (FileIndexPredicate index =
                new FileIndexPredicate(
                        new ByteArraySeekableStream(mainFile.embeddedIndex()), normalFileSchema)) {
            assertThat(((BitmapIndexResult) index.evaluate(predicateBuilder.equal(0, 2))).get())
                    .isEqualTo(RoaringBitmap32.bitmapOf(2));
            assertThat(index.evaluate(predicateBuilder.equal(0, 99)).remain()).isFalse();
        }
    }

    @Test
    void testVectorStoreFileNameFormatWithSharedUuid() throws Exception {
        // 100k vector-store data would create 1 normal, 1 blob, and 3 vector-store files
        int rowNum = 100 * 1000;
        writer.write(makeRows(rowNum, 1).iterator());
        writer.close();
        List<DataFileMeta> results = writer.result();

        // Get uuid from vector-store files. The pattern is data-{uuid}-{count}.vector.json
        DataFileMeta oneVectorStoreFile =
                results.stream().filter(file -> isVectorStoreFile(file.fileName())).findAny().get();
        String uuidAndCnt = oneVectorStoreFile.fileName().split(".vector.")[0];
        String prefix = uuidAndCnt.substring(0, uuidAndCnt.lastIndexOf('-') + 1); // data-{uuid}-

        // Verify all files use the same UUID and have sequential counters
        for (int i = 0; i < results.size(); ++i) {
            String fileName = results.get(i).fileName();
            assertThat(fileName).as("All files should use the same UUID").startsWith(prefix);
            int counter = Integer.parseInt(fileName.substring(prefix.length()).split("\\.")[0]);
            assertThat(counter).as("File counter should be sequential").isEqualTo(i);
        }
    }

    @Test
    void testVectorStoreStatsMainPart() throws Exception {
        // Write multiple rows
        int rowNum = RANDOM.nextInt(64) + 1;
        writer.write(makeRows(rowNum, 10).iterator());
        writer.close();
        List<DataFileMeta> metasResult = writer.result();

        // Check row count
        for (DataFileMeta file : metasResult) {
            assertThat(file.rowCount()).isEqualTo(rowNum);
            assertThat(file.deleteRowCount().get()).isEqualTo(0); // There is no deleted rows
        }

        // Check statistics
        for (DataFileMeta file : metasResult) {
            if (BlobFileFormat.isBlobFile(file.fileName())) {
                assertThat(file.writeCols()).isEqualTo(Collections.singletonList("f2"));
            } else if (isVectorStoreFile(file.fileName())) {
                assertThat(file.writeCols()).isEqualTo(Arrays.asList("f3"));
            } else {
                assertThat(file.writeCols()).isEqualTo(Arrays.asList("f0", "f1", "f4"));
                assertThat(file.valueStats().minValues().getInt(0)).isEqualTo(0);
                assertThat(file.valueStats().maxValues().getInt(0)).isEqualTo(rowNum - 1);
            }
        }

        // Verify total record count
        assertThat(writer.recordCount()).isEqualTo(rowNum);
    }

    @Test
    public void testVectorStoreNoBlob() throws Exception {
        RowType schema =
                RowType.builder()
                        .field("f0", DataTypes.INT())
                        .field("f1", DataTypes.STRING())
                        .field("f2", DataTypes.VECTOR(VECTOR_DIM, DataTypes.FLOAT()))
                        .field("f3", DataTypes.INT())
                        .build();
        writer =
                new DedicatedFormatRollingFileWriter(
                        LocalFileIO.create(),
                        SCHEMA_ID,
                        FileFormat.fromIdentifier("parquet", new Options()),
                        FileFormat.fromIdentifier("json", new Options()),
                        TARGET_FILE_SIZE,
                        TARGET_FILE_SIZE,
                        VECTOR_TARGET_FILE_SIZE,
                        Long.MAX_VALUE,
                        schema,
                        pathFactory,
                        () -> seqNumCounter,
                        COMPRESSION,
                        new StatsCollectorFactories(new CoreOptions(new Options())),
                        new FileIndexOptions(),
                        FileSource.APPEND,
                        false,
                        null);

        // 100k vector-store data would create 1 normal and 3 vector-store files
        int rowNum = 100 * 1000;
        List<InternalRow> rows = makeRows(rowNum, 1);
        for (InternalRow row : rows) {
            writer.write(
                    GenericRow.of(
                            row.getInt(0), row.getString(1), row.getVector(3), row.getInt(4)));
        }
        writer.close();
        List<DataFileMeta> results = writer.result();

        // Check normal, blob, and vector-store files
        List<DataFileMeta> normalFiles = new ArrayList<>();
        List<DataFileMeta> blobFiles = new ArrayList<>();
        List<DataFileMeta> vectorStoreFiles = new ArrayList<>();
        for (DataFileMeta file : results) {
            if (BlobFileFormat.isBlobFile(file.fileName())) {
                blobFiles.add(file);
            } else if (isVectorStoreFile(file.fileName())) {
                vectorStoreFiles.add(file);
            } else {
                normalFiles.add(file);
            }
        }
        assertThat(normalFiles.size()).isEqualTo(1);
        assertThat(blobFiles.size()).isEqualTo(0);
        assertThat(vectorStoreFiles.size()).isEqualTo(3);

        // Verify total record count
        assertThat(writer.recordCount()).isEqualTo(rowNum);
    }

    private List<InternalRow> makeRows(int rowNum, int blobDataSize) {
        List<InternalRow> rows = new ArrayList<>(rowNum);
        byte[] blobData = new byte[blobDataSize];
        RANDOM.nextBytes(blobData);
        for (int i = 0; i < rowNum; ++i) {
            byte[] string = new byte[1];
            RANDOM.nextBytes(string);
            byte[] buf = new byte[VECTOR_DIM];
            RANDOM.nextBytes(buf);
            float[] vector = new float[VECTOR_DIM];
            for (int j = 0; j < VECTOR_DIM; ++j) {
                vector[j] = buf[j];
            }
            int label = RANDOM.nextInt(32) + 1;
            rows.add(
                    GenericRow.of(
                            i,
                            BinaryString.fromBytes(string),
                            new BlobData(blobData),
                            BinaryVector.fromPrimitiveArray(vector),
                            label));
        }
        return rows;
    }

    private void assertBundleSequenceSideEffects(List<DataFileMeta> metas, long rowCount) {
        long expectedCounter = rowCount * metas.size();
        assertThat(seqNumCounter.getValue()).isEqualTo(expectedCounter);
        for (DataFileMeta meta : metas) {
            assertThat(meta.minSequenceNumber()).isEqualTo(expectedCounter - rowCount);
            assertThat(meta.maxSequenceNumber()).isEqualTo(expectedCounter - 1);
        }
    }

    private static class SingleUseBundleRecords implements BundleRecords {

        private final List<InternalRow> rows;
        private boolean iterated;

        private SingleUseBundleRecords(List<InternalRow> rows) {
            this.rows = rows;
        }

        @Override
        public Iterator<InternalRow> iterator() {
            if (iterated) {
                throw new IllegalStateException("Bundle should only be consumed once.");
            }
            iterated = true;
            return rows.iterator();
        }

        @Override
        public long rowCount() {
            return rows.size();
        }
    }
}
