/*
 *  Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
 *  <p>
 *  Licensed 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
 *  <p>
 *  http://www.apache.org/licenses/LICENSE-2.0
 *  <p>
 *  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 com.easyagents.store.chroma;

import com.easyagents.core.document.Document;
import com.easyagents.core.store.SearchWrapper;
import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreResult;
import com.easyagents.core.model.client.HttpClient;
import org.junit.AfterClass;
import org.junit.Assume;
import org.junit.BeforeClass;
import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import static org.junit.Assert.*;

/**
 * ChromaVectorStore的测试类，测试文档的存储、搜索、更新和删除功能
 * 包含连接检查和错误处理机制，支持在无真实Chroma服务器时跳过测试
 */
public class ChromaVectorStoreTest {

    private static ChromaVectorStore store;
    private static String testTenant = "default_tenant";
    private static String testDatabase = "default_database";
    private static String testCollectionName = "test_collection";
    private static boolean isChromaAvailable = false;
    private static boolean useMock = false; // 设置为true可以在没有真实Chroma服务器时使用模拟模式

    /**
     * 在测试开始前初始化ChromaVectorStore实例
     */
    @BeforeClass
    public static void setUp() {
        // 创建配置对象
        ChromaVectorStoreConfig config = new ChromaVectorStoreConfig();
        config.setHost("localhost");
        config.setPort(8000);
        config.setCollectionName(testCollectionName);
        config.setTenant(testTenant);
        config.setDatabase(testDatabase);
        config.setAutoCreateCollection(true);

        // 初始化存储实例
        try {
            store = new ChromaVectorStore(config);
            System.out.println("ChromaVectorStore initialized successfully.");

            // 检查连接是否可用
            isChromaAvailable = checkChromaConnection(config);

            if (!isChromaAvailable && !useMock) {
                System.out.println("Chroma server is not available. Tests will be skipped unless useMock is set to true.");
            }
        } catch (Exception e) {
            System.err.println("Failed to initialize ChromaVectorStore: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 检查Chroma服务器连接是否可用
     */
    private static boolean checkChromaConnection(ChromaVectorStoreConfig config) {
        try {
            String baseUrl = "http://" + config.getHost() + ":" + config.getPort();
            String healthCheckUrl = baseUrl + "/api/v2/heartbeat";

            HttpClient httpClient = new HttpClient();
            System.out.println("Checking Chroma server connection at: " + healthCheckUrl);

            // 使用较短的超时时间进行健康检查
            String response = httpClient.get(healthCheckUrl);
            if (response != null) {
                System.out.println("Chroma server connection successful! Response: " + response);
                return true;
            } else {
                System.out.println("Chroma server connection failed: Empty response");
                return false;
            }
        } catch (Exception e) {
            System.out.println("Chroma server connection failed: " + e.getMessage());
            System.out.println("Please ensure Chroma server is running on http://" + config.getHost() + ":" + config.getPort());
            System.out.println("To run tests without a real Chroma server, set 'useMock = true'");
            return false;
        }
    }

    /**
     * 检查是否应该运行测试
     */
    private void assumeChromaAvailable() {
        Assume.assumeTrue("Chroma server is not available and mock mode is disabled",
                         isChromaAvailable || useMock);
    }

    /**
     * 在所有测试完成后清理资源
     */
    @AfterClass
    public static void tearDown() {
        if (store != null) {
            try {
                store.close();
                System.out.println("ChromaVectorStore closed successfully.");
            } catch (Exception e) {
                System.err.println("Error closing ChromaVectorStore: " + e.getMessage());
            }
        }
    }

    /**
     * 测试存储文档功能
     */
    @Test
    public void testStoreDocuments() {
        assumeChromaAvailable();

        System.out.println("Starting testStoreDocuments...");

        // 创建测试文档
        List<Document> documents = createTestDocuments();

        // 如果使用模拟模式，直接返回成功结果
        if (useMock) {
            System.out.println("Running in mock mode. Simulating store operation.");
            StoreResult mockResult = StoreResult.successWithIds(documents);
            assertTrue("Mock store operation should be successful", mockResult.isSuccess());
            assertEquals("All document IDs should be returned in mock mode",
                         documents.size(), mockResult.ids().size());
            System.out.println("testStoreDocuments completed successfully in mock mode.");
            return;
        }

        // 存储文档
        try {
            StoreResult result = store.doStore(documents, StoreOptions.DEFAULT);
            System.out.println("Store result: " + result);

            // 验证存储是否成功
            assertTrue("Store operation should be successful", result.isSuccess());
            assertEquals("All document IDs should be returned", documents.size(), result.ids().size());

            System.out.println("testStoreDocuments completed successfully.");
        } catch (Exception e) {
            System.err.println("Failed to store documents: " + e.getMessage());
            e.printStackTrace();
            fail("Store operation failed with exception: " + e.getMessage());
        }
    }

    /**
     * 测试搜索文档功能
     */
    @Test
    public void testSearchDocuments() {
        assumeChromaAvailable();

        System.out.println("Starting testSearchDocuments...");

        // 创建测试文档
        List<Document> documents = createTestDocuments();

        // 如果使用模拟模式
        if (useMock) {
            System.out.println("Running in mock mode. Simulating search operation.");
            // 模拟搜索结果，返回前3个文档
            List<Document> mockResults = new ArrayList<>(documents.subList(0, Math.min(3, documents.size())));
            for (int i = 0; i < mockResults.size(); i++) {
                mockResults.get(i).setScore(1.0 - i * 0.1); // 模拟相似度分数
            }

            // 验证模拟结果
            assertNotNull("Mock search results should not be null", mockResults);
            assertFalse("Mock search results should not be empty", mockResults.isEmpty());
            assertTrue("Mock search results should have the correct maximum size", mockResults.size() <= 3);

            System.out.println("testSearchDocuments completed successfully in mock mode.");
            return;
        }

        try {
            // 首先存储一些测试文档
            store.doStore(documents, StoreOptions.DEFAULT);

            // 创建搜索包装器
            SearchWrapper searchWrapper = new SearchWrapper();
            // 使用第一个文档的向量进行搜索
            searchWrapper.setVector(documents.get(0).getVector());
            searchWrapper.setMaxResults(3);

            // 执行搜索
            List<Document> searchResults = store.doSearch(searchWrapper, StoreOptions.DEFAULT);

            // 验证搜索结果
            assertNotNull("Search results should not be null", searchResults);
            assertFalse("Search results should not be empty", searchResults.isEmpty());
            assertTrue("Search results should have the correct maximum size",
                searchResults.size() <= searchWrapper.getMaxResults());

            // 打印搜索结果
            System.out.println("Search results:");
            for (Document doc : searchResults) {
                System.out.printf("id=%s, content=%s, vector=%s, score=%s\n",
                    doc.getId(), doc.getContent(), Arrays.toString(doc.getVector()), doc.getScore());
            }

            System.out.println("testSearchDocuments completed successfully.");
        } catch (Exception e) {
            System.err.println("Failed to search documents: " + e.getMessage());
            e.printStackTrace();
            fail("Search operation failed with exception: " + e.getMessage());
        }
    }

    /**
     * 测试更新文档功能
     */
    @Test
    public void testUpdateDocuments() {
        assumeChromaAvailable();

        System.out.println("Starting testUpdateDocuments...");

        // 创建测试文档
        List<Document> documents = createTestDocuments();

        // 如果使用模拟模式
        if (useMock) {
            System.out.println("Running in mock mode. Simulating update operation.");

            // 修改文档内容
            Document updatedDoc = documents.get(0);
            String originalContent = updatedDoc.getContent();
            updatedDoc.setContent(originalContent + " [UPDATED]");

            // 模拟更新结果
            StoreResult mockResult = StoreResult.successWithIds(Arrays.asList(updatedDoc));
            assertTrue("Mock update operation should be successful", mockResult.isSuccess());

            System.out.println("testUpdateDocuments completed successfully in mock mode.");
            return;
        }

        try {
            // 首先存储一些测试文档
            store.doStore(documents, StoreOptions.DEFAULT);

            // 修改文档内容
            Document updatedDoc = documents.get(0);
            String originalContent = updatedDoc.getContent();
            updatedDoc.setContent(originalContent + " [UPDATED]");

            // 执行更新
            StoreResult result = store.doUpdate(Arrays.asList(updatedDoc), StoreOptions.DEFAULT);

            // 验证更新是否成功
            assertTrue("Update operation should be successful", result.isSuccess());

            // 搜索更新后的文档以验证更改
            SearchWrapper searchWrapper = new SearchWrapper();
            searchWrapper.setVector(updatedDoc.getVector());
            searchWrapper.setMaxResults(1);

            List<Document> searchResults = store.doSearch(searchWrapper, StoreOptions.DEFAULT);
            assertTrue("Should find the updated document", !searchResults.isEmpty());
            assertEquals("Document content should be updated",
                updatedDoc.getContent(), searchResults.get(0).getContent());

            System.out.println("testUpdateDocuments completed successfully.");
        } catch (Exception e) {
            System.err.println("Failed to update documents: " + e.getMessage());
            e.printStackTrace();
            fail("Update operation failed with exception: " + e.getMessage());
        }
    }

    /**
     * 测试删除文档功能
     */
    @Test
    public void testDeleteDocuments() {
        assumeChromaAvailable();

        System.out.println("Starting testDeleteDocuments...");

        // 创建测试文档
        List<Document> documents = createTestDocuments();

        // 如果使用模拟模式
        if (useMock) {
            System.out.println("Running in mock mode. Simulating delete operation.");

            // 获取要删除的文档ID
            List<Object> idsToDelete = new ArrayList<>();
            idsToDelete.add(documents.get(0).getId());

            // 模拟删除结果
            StoreResult mockResult = StoreResult.success();
            assertTrue("Mock delete operation should be successful", mockResult.isSuccess());

            System.out.println("testDeleteDocuments completed successfully in mock mode.");
            return;
        }

        try {
            // 首先存储一些测试文档
            store.doStore(documents, StoreOptions.DEFAULT);

            // 获取要删除的文档ID
            List<Object> idsToDelete = new ArrayList<>();
            idsToDelete.add(documents.get(0).getId());

            // 执行删除
            StoreResult result = store.doDelete(idsToDelete, StoreOptions.DEFAULT);

            // 验证删除是否成功
            assertTrue("Delete operation should be successful", result.isSuccess());

            // 尝试搜索已删除的文档
            SearchWrapper searchWrapper = new SearchWrapper();
            searchWrapper.setVector(documents.get(0).getVector());
            searchWrapper.setMaxResults(10);

            List<Document> searchResults = store.doSearch(searchWrapper, StoreOptions.DEFAULT);

            // 检查结果中是否包含已删除的文档
            boolean deletedDocFound = searchResults.stream()
                .anyMatch(doc -> doc.getId().equals(documents.get(0).getId()));

            assertFalse("Deleted document should not be found", deletedDocFound);

            System.out.println("testDeleteDocuments completed successfully.");
        } catch (Exception e) {
            System.err.println("Failed to delete documents: " + e.getMessage());
            e.printStackTrace();
            fail("Delete operation failed with exception: " + e.getMessage());
        }
    }

    /**
     * 创建测试文档
     */
    private List<Document> createTestDocuments() {
        List<Document> documents = new ArrayList<>();

        // 创建5个测试文档，每个文档都有不同的内容和向量
        for (int i = 0; i < 5; i++) {
            Document doc = new Document();
            doc.setId("doc_" + i);
            doc.setContent("This is test document content " + i);
            doc.setTitle("Test Document " + i);

            // 创建一个简单的向量，向量维度为10
            float[] vector = new float[10];
            for (int j = 0; j < vector.length; j++) {
                vector[j] = i + j * 0.1f;
            }
            doc.setVector(vector);

            documents.add(doc);
        }

        return documents;
    }
}
