// 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.index.vector;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

public class VectorIndexJavaApiTest {

    public static void main(String[] args) {
        testNativeLibraryResourcePaths();
        testSingleResultCopiesArrays();
        testBatchResultCopiesArraysAndSlicesRows();
        testMetadata();
        testSearchParametersRemainAlgorithmSpecific();
        testReaderRejectsNegativeAdjacencyCacheBudget();
        testClosedReaderRejectsOperations();
        testClosedTrainerRejectsOperations();
        testClosedTrainingRejectsOperations();
        testClosedWriterRejectsOperations();
        testReaderAndWriterApiCompile();
    }

    private static void testNativeLibraryResourcePaths() {
        assertEquals(
                "/native/linux/x86_64/libpaimon_vindex_jni.so",
                NativeLibraryLoader.resourcePath("Linux", "amd64"));
        assertEquals(
                "/native/linux/aarch64/libpaimon_vindex_jni.so",
                NativeLibraryLoader.resourcePath("Linux", "arm64"));
        assertEquals(
                "/native/macos/aarch64/libpaimon_vindex_jni.dylib",
                NativeLibraryLoader.resourcePath("Mac OS X", "aarch64"));
        assertEquals(
                "/native/windows/x86_64/paimon_vindex_jni.dll",
                NativeLibraryLoader.resourcePath("Windows 11", "x86_64"));
        assertThrows(
                UnsatisfiedLinkError.class,
                new ThrowingRunnable() {
                    @Override
                    public void run() {
                        NativeLibraryLoader.resourcePath("Mac OS X", "x86_64");
                    }
                });
    }

    private static void testSearchParametersRemainAlgorithmSpecific() {
        VectorSearchParams params = new VectorSearchParams(10, 4).withLSearch(200);

        assertEquals(VectorSearchParams.SEARCH_WIDTH_DISKANN_L_SEARCH, params.searchWidth());
        assertEquals(200, params.width());
        assertEquals(
                VectorSearchParams.SEARCH_WIDTH_IVF_NPROBE,
                new VectorSearchParams(10, 4).searchWidth());
    }

    private static void testSingleResultCopiesArrays() {
        long[] ids = new long[] {11L, 7L};
        float[] distances = new float[] {0.1f, 0.3f};

        VectorSearchResult result = new VectorSearchResult(ids, distances);
        ids[0] = 99L;
        distances[0] = 9.0f;

        assertArrayEquals(new long[] {11L, 7L}, result.ids());
        assertArrayEquals(new float[] {0.1f, 0.3f}, result.distances());

        long[] resultIds = result.ids();
        resultIds[0] = 99L;
        assertArrayEquals(new long[] {11L, 7L}, result.ids());
    }

    private static void testBatchResultCopiesArraysAndSlicesRows() {
        long[] ids = new long[] {1L, 2L, 3L, 4L, 5L, 6L};
        float[] distances = new float[] {0.1f, 0.2f, 0.3f, 1.1f, 1.2f, 1.3f};

        VectorSearchBatchResult result = new VectorSearchBatchResult(ids, distances, 2, 3);
        ids[0] = 99L;
        distances[0] = 9.0f;

        assertEquals(2, result.queryCount());
        assertEquals(3, result.topK());
        assertArrayEquals(new long[] {1L, 2L, 3L, 4L, 5L, 6L}, result.ids());
        assertArrayEquals(new float[] {0.1f, 0.2f, 0.3f, 1.1f, 1.2f, 1.3f}, result.distances());
        assertArrayEquals(new long[] {4L, 5L, 6L}, result.idsForQuery(1));
        assertArrayEquals(new float[] {1.1f, 1.2f, 1.3f}, result.distancesForQuery(1));

        assertThrows(IllegalArgumentException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                new VectorSearchBatchResult(new long[] {1L}, new float[] {1.0f}, 2, 3);
            }
        });
        assertThrows(IndexOutOfBoundsException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                result.idsForQuery(2);
            }
        });
    }

    private static void testMetadata() {
        VectorIndexMetadata metadata =
                new VectorIndexMetadata(
                        "ivf_sq", 16, 4, "cosine", 123L, 0, 8, 0, 0, 0, 0.0f);
        assertEquals("ivf_sq", metadata.indexType());
        assertEquals(16, metadata.dimension());
        assertEquals(4, metadata.nlist());
        assertEquals("cosine", metadata.metric());
        assertEquals(123L, metadata.totalVectors());
        assertEquals(8, metadata.pqBits());
        assertEquals(0, metadata.rqBits());

        VectorIndexMetadata diskAnnMetadata =
                new VectorIndexMetadata(
                        "diskann", 128, 1, "l2", 1000L, 16, 4, 0, 64, 100, 1.2f);
        assertEquals(4, diskAnnMetadata.pqBits());
        assertEquals(0, diskAnnMetadata.rqBits());
        assertEquals(64, diskAnnMetadata.diskAnnMaxDegree());
        assertEquals(100, diskAnnMetadata.diskAnnBuildSearchListSize());
        assertEquals(Float.valueOf(1.2f), Float.valueOf(diskAnnMetadata.diskAnnAlpha()));
    }

    private static void testReaderRejectsNegativeAdjacencyCacheBudget() {
        assertThrows(IllegalArgumentException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                new VectorIndexReader(
                        new EmptyVectorIndexInput(),
                        -1);
            }
        });
    }

    private static void testClosedReaderRejectsOperations() {
        final VectorIndexReader reader = VectorIndexReader.fromNativePointerForTesting(0L);
        reader.close();
        reader.close();

        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.metadata();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.indexType();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.dimension();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.totalVectors();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.optimizeForSearch();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.warmupQueries(new float[] {0.0f}, 1, 0);
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.search(new float[] {0.0f}, new VectorSearchParams(1, 1));
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                reader.searchBatch(new float[] {0.0f}, 1, new VectorSearchParams(1, 1));
            }
        });
    }

    private static void testClosedWriterRejectsOperations() {
        final VectorIndexWriter writer = VectorIndexWriter.fromNativePointerForTesting(0L);
        writer.close();
        writer.close();

        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                writer.dimension();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                writer.addVectors(new long[] {1L}, new float[] {0.0f, 1.0f}, 1);
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                writer.writeIndex(new Object());
            }
        });
    }

    private static void testClosedTrainerRejectsOperations() {
        final VectorIndexTrainer trainer = VectorIndexTrainer.fromNativePointerForTesting(0L);
        trainer.close();
        trainer.close();

        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                trainer.dimension();
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                trainer.addTrainingVectors(new float[] {0.0f, 1.0f}, 1);
            }
        });
        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                trainer.finishTraining();
            }
        });
    }

    private static void testClosedTrainingRejectsOperations() {
        final VectorIndexTraining training = VectorIndexTraining.fromNativePointerForTesting(0L);
        training.close();
        training.close();

        assertThrows(IllegalStateException.class, new ThrowingRunnable() {
            @Override
            public void run() {
                new VectorIndexWriter(training);
            }
        });
    }

    private static void testReaderAndWriterApiCompile() {
        Map<String, String> options = ivfPqOptions(2, 4);
        VectorIndexReader closedReader = VectorIndexReader.fromNativePointerForTesting(0L);
        closedReader.close();
        closedReader.close();

        VectorIndexWriter closedWriter = VectorIndexWriter.fromNativePointerForTesting(0L);
        closedWriter.close();
        closedWriter.close();

        if (System.currentTimeMillis() < 0) {
            VectorIndexReader reader = new VectorIndexReader(new EmptyVectorIndexInput());
            VectorIndexReader coalescedReader =
                    new VectorIndexReader(
                            new EmptyVectorIndexInput(),
                            4L * 1024 * 1024 * 1024);
            coalescedReader.close();
            reader.metadata();
            reader.indexType();
            reader.dimension();
            reader.totalVectors();
            reader.optimizeForSearch();
            reader.warmupQueries(new float[] {0.0f, 1.0f}, 1, 100);
            reader.readPlan();
            VectorSearchParams params = new VectorSearchParams(10, 4);
            VectorSearchParams diskAnnParams = params.withLSearch(100);
            reader.search(new float[] {0.0f, 1.0f}, params);
            reader.search(new float[] {0.0f, 1.0f}, diskAnnParams);
            reader.search(new float[] {0.0f, 1.0f}, params, new byte[] {1, 2});
            reader.search(new float[] {0.0f, 1.0f}, diskAnnParams, new byte[] {1, 2});
            reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, params);
            reader.searchBatch(
                    new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, diskAnnParams);
            reader.searchBatch(new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2, params, new byte[] {1, 2});
            reader.searchBatch(
                    new float[] {0.0f, 1.0f, 2.0f, 3.0f},
                    2,
                    diskAnnParams,
                    new byte[] {1, 2});
            VectorIndexTraining training =
                    VectorIndexTrainer.train(options, new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2);
            VectorIndexWriter writer = new VectorIndexWriter(training);
            writer.dimension();
            writer.addVectors(new long[] {1L, 2L}, new float[] {0.0f, 1.0f, 2.0f, 3.0f}, 2);
            writer.writeIndex(new Object());

            VectorIndexTrainer trainer = VectorIndexTrainer.create(options);
            trainer.dimension();
            VectorIndexTraining stagedTraining =
                    trainer.addTrainingVectors(new float[] {0.0f, 1.0f}, 1)
                            .addTrainingVectors(new float[] {2.0f, 3.0f}, 1)
                            .finishTraining();
            VectorIndexWriter stagedWriter = new VectorIndexWriter(stagedTraining);
            stagedWriter.addVectors(new long[] {1L}, new float[] {0.0f, 1.0f}, 1);
            stagedWriter.writeIndex(new Object());
        }
    }

    private static Map<String, String> ivfFlatOptions(int dimension, int nlist) {
        Map<String, String> options = new HashMap<String, String>();
        options.put("index.type", "ivf_flat");
        options.put("dimension", Integer.toString(dimension));
        options.put("nlist", Integer.toString(nlist));
        options.put("metric", "l2");
        return options;
    }

    private static Map<String, String> ivfPqOptions(int dimension, int nlist) {
        Map<String, String> options = ivfFlatOptions(dimension, nlist);
        options.put("index.type", "ivf_pq");
        options.put("use-opq", "false");
        return options;
    }

    private static void assertEquals(int expected, int actual) {
        if (expected != actual) {
            throw new AssertionError("expected " + expected + " but got " + actual);
        }
    }

    private static void assertEquals(long expected, long actual) {
        if (expected != actual) {
            throw new AssertionError("expected " + expected + " but got " + actual);
        }
    }

    private static void assertEquals(Object expected, Object actual) {
        if (!expected.equals(actual)) {
            throw new AssertionError("expected " + expected + " but got " + actual);
        }
    }

    private static void assertArrayEquals(long[] expected, long[] actual) {
        if (!Arrays.equals(expected, actual)) {
            throw new AssertionError(
                    "expected " + Arrays.toString(expected) + " but got " + Arrays.toString(actual));
        }
    }

    private static void assertArrayEquals(float[] expected, float[] actual) {
        if (!Arrays.equals(expected, actual)) {
            throw new AssertionError(
                    "expected " + Arrays.toString(expected) + " but got " + Arrays.toString(actual));
        }
    }

    private static void assertThrows(Class<? extends Throwable> expected, ThrowingRunnable runnable) {
        try {
            runnable.run();
        } catch (Throwable t) {
            if (expected.isInstance(t)) {
                return;
            }
            throw new AssertionError(
                    "expected " + expected.getName() + " but got " + t.getClass().getName(), t);
        }
        throw new AssertionError("expected " + expected.getName());
    }

    private interface ThrowingRunnable {
        void run() throws Throwable;
    }

    private static final class EmptyVectorIndexInput implements VectorIndexInput {
        @Override
        public void pread(long[] positions, byte[][] buffers) {}
    }
}
