package com.easyagents.engine.es;

import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.BulkRequest;
import com.easyagents.core.document.Document;
import com.easyagents.search.engine.service.KeywordSearchMetadataKeys;
import com.easyagents.search.engine.service.KeywordSearchRequest;
import org.junit.Assert;
import org.junit.Test;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * {@link ElasticSearcher} 请求构建回归测试。
 */
public class ElasticSearcherQueryBuilderTest {

    /**
     * 验证搜索请求同时包含多字段检索和知识库过滤。
     *
     * @throws Exception Elasticsearch 客户端关闭失败时抛出
     */
    @Test
    public void shouldBuildSearchRequestWithMultiMatchAndKnowledgeFilter() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            KeywordSearchRequest request = KeywordSearchRequest.of("客服", 5);
            request.setKnowledgeId("100");

            SearchRequest searchRequest = searcher.buildSearchRequest(request);

            Assert.assertEquals(5, searchRequest.size().intValue());
            Assert.assertNotNull(searchRequest.query().bool());
            Assert.assertEquals(1, searchRequest.query().bool().must().size());
            Assert.assertNotNull(searchRequest.query().bool().must().get(0).multiMatch());
            Assert.assertEquals(2, searchRequest.query().bool().must().get(0).multiMatch().fields().size());
            Assert.assertEquals(1, searchRequest.query().bool().filter().size());
            Assert.assertEquals("knowledgeId", searchRequest.query().bool().filter().get(0).term().field());
        }
    }

    /**
     * 验证知识库 ID 会写入顶层字段。
     *
     * @throws Exception Elasticsearch 客户端关闭失败时抛出
     */
    @Test
    public void shouldExtractKnowledgeIdToTopLevelSource() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            Document document = new Document();
            document.setId("1");
            document.setTitle("title");
            document.setContent("content");
            document.addMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID, "100");

            Map<String, Object> source = searcher.buildSource(document);

            Assert.assertEquals("100", source.get(KeywordSearchMetadataKeys.KNOWLEDGE_ID));
            Assert.assertTrue(source.get("metadataMap") instanceof Map);
        }
    }

    /**
     * 验证多个文档会进入同一个 Bulk 请求并保留稳定文档 ID。
     *
     * @throws Exception Elasticsearch 客户端关闭失败时抛出
     */
    @Test
    public void shouldBuildSingleBulkRequestForMultipleDocuments() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            Document first = new Document();
            first.setId("first");
            first.setContent("first content");
            Document second = new Document();
            second.setId("second");
            second.setContent("second content");

            BulkRequest request = searcher.buildBulkRequest(List.of(first, second));

            Assert.assertEquals(2, request.operations().size());
            Assert.assertEquals("first", request.operations().get(0).index().id());
            Assert.assertEquals("second", request.operations().get(1).index().id());
        }
    }

    /**
     * 验证 Bulk 写入会按最大操作数拆分，避免单次请求无界增长。
     *
     * @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出
     */
    @Test
    public void shouldPartitionBulkRequestsByOperationCount() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            List<Document> documents = new ArrayList<>();
            for (int index = 0; index <= ElasticSearcher.MAX_BULK_OPERATIONS; index++) {
                documents.add(document("count-" + index, "content"));
            }

            List<List<Document>> batches = searcher.partitionDocuments(documents);

            Assert.assertEquals(2, batches.size());
            Assert.assertEquals(ElasticSearcher.MAX_BULK_OPERATIONS, batches.get(0).size());
            Assert.assertEquals(1, batches.get(1).size());
        }
    }

    /**
     * 验证 Bulk 写入会按序列化字节数拆分。
     *
     * @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出
     */
    @Test
    public void shouldPartitionBulkRequestsByEstimatedBytes() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            int contentLength = (int) (ElasticSearcher.MAX_BULK_ESTIMATED_BYTES / 2L);
            Document first = document("bytes-first", "a".repeat(contentLength));
            Document second = document("bytes-second", "b".repeat(contentLength));

            List<List<Document>> batches = searcher.partitionDocuments(List.of(first, second));

            Assert.assertEquals(2, batches.size());
            Assert.assertEquals(1, batches.get(0).size());
            Assert.assertEquals(1, batches.get(1).size());
        }
    }

    /**
     * 验证单个超限文档会在发送请求前失败。
     *
     * @throws Exception Elasticsearch 客户端关闭或文档序列化失败时抛出
     */
    @Test
    public void shouldRejectSingleDocumentOverByteLimit() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            int contentLength = (int) ElasticSearcher.MAX_BULK_ESTIMATED_BYTES;
            Document oversized = document("oversized", "a".repeat(contentLength));

            try {
                searcher.partitionDocuments(List.of(oversized));
                Assert.fail("单个超限文档应拒绝构建 Bulk 请求");
            } catch (IllegalArgumentException expected) {
                Assert.assertTrue(expected.getMessage().contains("oversized"));
            }
        }
    }

    /**
     * 验证批量删除请求包含全部稳定文档 ID。
     *
     * @throws Exception Elasticsearch 客户端关闭失败时抛出
     */
    @Test
    public void shouldBuildBulkDeleteRequest() throws Exception {
        try (ElasticSearcher searcher = new ElasticSearcher(config())) {
            BulkRequest request = searcher.buildDeleteBulkRequest(List.of("first", "second"));

            Assert.assertEquals(2, request.operations().size());
            Assert.assertEquals("first", request.operations().get(0).delete().id());
            Assert.assertEquals("second", request.operations().get(1).delete().id());
        }
    }

    /**
     * 创建测试文档。
     *
     * @param id 文档 ID
     * @param content 文档内容
     * @return 文档
     */
    private Document document(String id, String content) {
        Document document = new Document();
        document.setId(id);
        document.setContent(content);
        return document;
    }

    /**
     * 创建测试用 Elasticsearch 配置。
     *
     * @return Elasticsearch 配置
     */
    private ESConfig config() {
        ESConfig config = new ESConfig();
        config.setHost("http://127.0.0.1:9200");
        config.setUserName("elastic");
        config.setPassword("elastic");
        config.setIndexName("easyflow");
        return config;
    }
}
