package com.easyagents.engine.es;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.*;
import co.elastic.clients.elasticsearch.core.search.SourceConfig;
import co.elastic.clients.json.JsonData;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.ElasticsearchTransport;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.easyagents.core.document.Document;
import com.easyagents.search.engine.service.DocumentSearcher;
import com.easyagents.search.engine.service.KeywordSearchMetadataKeys;
import com.easyagents.search.engine.service.KeywordSearchRequest;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.client.RestClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.io.IOException;
import java.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.X509Certificate;
import java.util.*;

/**
 * 基于 Elasticsearch 的关键词搜索器。
 *
 * <p>底层客户端线程安全并与搜索器实例共享生命周期，避免逐文档重复创建网络连接。</p>
 */
public class ElasticSearcher implements DocumentSearcher, AutoCloseable {

    private static final Logger LOG = LoggerFactory.getLogger(ElasticSearcher.class);
    static final int MAX_BULK_OPERATIONS = 200;
    static final long MAX_BULK_ESTIMATED_BYTES = 5L * 1024L * 1024L;
    private static final long BULK_OPERATION_METADATA_BYTES = 128L;

    private final ESConfig esConfig;
    private final JacksonJsonpMapper jsonpMapper = new JacksonJsonpMapper();
    private final ElasticsearchTransport transport;
    private final ElasticsearchClient client;

    /**
     * 创建 Elasticsearch 搜索器及其共享客户端。
     *
     * @param esConfig Elasticsearch 配置
     * @throws IllegalStateException 客户端初始化失败时抛出
     */
    public ElasticSearcher(ESConfig esConfig) {
        this.esConfig = Objects.requireNonNull(esConfig, "ESConfig 不能为 null");
        RestClient openedRestClient = null;
        try {
            openedRestClient = buildRestClient();
            this.transport = new RestClientTransport(openedRestClient, jsonpMapper);
            this.client = new ElasticsearchClient(transport);
        } catch (Exception e) {
            closeQuietly(openedRestClient);
            throw new IllegalStateException("初始化 Elasticsearch 客户端失败", e);
        }
    }

    // 忽略 SSL 的 client 构建逻辑
    private RestClient buildRestClient() throws NoSuchAlgorithmException, KeyManagementException {
        TrustManager[] trustAllCerts = new TrustManager[]{
            new X509TrustManager() {
                public X509Certificate[] getAcceptedIssuers() {
                    return null;
                }

                public void checkClientTrusted(X509Certificate[] certs, String authType) {
                }

                public void checkServerTrusted(X509Certificate[] certs, String authType) {
                }
            }
        };

        SSLContext sslContext = SSLContext.getInstance("TLS");
        sslContext.init(null, trustAllCerts, new java.security.SecureRandom());

        CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(
            AuthScope.ANY,
            new UsernamePasswordCredentials(esConfig.getUserName(), esConfig.getPassword()));

        return RestClient.builder(HttpHost.create(esConfig.getHost()))
            .setHttpClientConfigCallback(httpClientBuilder -> {
                httpClientBuilder.setSSLContext(sslContext);
                httpClientBuilder.setSSLHostnameVerifier((hostname, session) -> true);
                httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
                return httpClientBuilder;
            })
            .build();
    }


    /**
     * 添加文档到 Elasticsearch。
     *
     * @param document 待写入文档
     * @return 写入成功时返回 {@code true}
     */
    @Override
    public boolean addDocument(Document document) {
        return addDocuments(Collections.singletonList(document));
    }

    /**
     * 使用有界 Bulk 请求批量写入文档。
     *
     * @param documents 待写入文档
     * @return 全部文档写入成功时返回 {@code true}
     */
    @Override
    public boolean addDocuments(List<Document> documents) {
        if (!areValidDocuments(documents)) {
            return false;
        }
        try {
            List<List<Document>> batches = partitionDocuments(documents);
            for (int batchIndex = 0; batchIndex < batches.size(); batchIndex++) {
                List<Document> batch = batches.get(batchIndex);
                BulkResponse response = client.bulk(buildBulkRequest(batch));
                if (!isBulkSuccessful(response, "写入", batchIndex, batch.size())) {
                    return false;
                }
            }
            return true;
        } catch (Exception e) {
            LOG.error("Elasticsearch 批量写入异常: count={}", documents.size(), e);
            return false;
        }
    }

    /**
     * 按操作数量和序列化后的估算字节数拆分 Bulk 批次。
     *
     * @param documents 已完成校验的文档
     * @return 有界文档批次
     * @throws IOException 文档序列化失败时抛出
     * @throws IllegalArgumentException 单个文档超过批次字节限制时抛出
     */
    List<List<Document>> partitionDocuments(List<Document> documents) throws IOException {
        List<List<Document>> batches = new ArrayList<>();
        List<Document> currentBatch = new ArrayList<>(Math.min(documents.size(), MAX_BULK_OPERATIONS));
        long currentBatchBytes = 0L;
        for (Document document : documents) {
            long documentBytes = estimateBulkOperationBytes(document);
            if (documentBytes > MAX_BULK_ESTIMATED_BYTES) {
                throw new IllegalArgumentException(
                    "单个 Elasticsearch 索引文档超过 Bulk 字节限制: id=" + document.getId());
            }
            if (!currentBatch.isEmpty()
                && (currentBatch.size() >= MAX_BULK_OPERATIONS
                || currentBatchBytes + documentBytes > MAX_BULK_ESTIMATED_BYTES)) {
                batches.add(currentBatch);
                currentBatch = new ArrayList<>(Math.min(documents.size(), MAX_BULK_OPERATIONS));
                currentBatchBytes = 0L;
            }
            currentBatch.add(document);
            currentBatchBytes += documentBytes;
        }
        if (!currentBatch.isEmpty()) {
            batches.add(currentBatch);
        }
        return batches;
    }

    private long estimateBulkOperationBytes(Document document) throws IOException {
        long sourceBytes = jsonpMapper.objectMapper().writeValueAsBytes(buildSource(document)).length;
        return sourceBytes
            + utf8Length(esConfig.getIndexName())
            + utf8Length(document.getId().toString())
            + BULK_OPERATION_METADATA_BYTES;
    }

    private int utf8Length(String value) {
        return value == null ? 0 : value.getBytes(java.nio.charset.StandardCharsets.UTF_8).length;
    }

    /**
     * 构建包含全部文档的 Bulk 请求。
     *
     * @param documents 已完成校验的文档
     * @return Bulk 请求
     */
    BulkRequest buildBulkRequest(List<Document> documents) {
        BulkRequest.Builder builder = new BulkRequest.Builder();
        for (Document document : documents) {
            builder.operations(operation -> operation.index(index -> index
                .index(esConfig.getIndexName())
                .id(document.getId().toString())
                .document(JsonData.of(buildSource(document)))
            ));
        }
        return builder.build();
    }

    private boolean areValidDocuments(List<Document> documents) {
        if (documents == null || documents.isEmpty()) {
            return false;
        }
        for (Document document : documents) {
            if (document == null || document.getId() == null || document.getContent() == null) {
                return false;
            }
        }
        return true;
    }

    /**
     * 按请求条件搜索文档。
     *
     * @param request 搜索请求
     * @return 命中文档
     */
    @Override
    public List<Document> searchDocuments(KeywordSearchRequest request) {
        try {
            SearchResponse<Map> response = client.search(buildSearchRequest(request), Map.class);
            List<Document> results = new ArrayList<>();
            response.hits().hits().forEach(hit -> {
                Map source = hit.source();
                Document document = toDocument(hit.id(), source, hit.score());
                if (document == null) {
                    return;
                }
                results.add(document);
            });
            return results;

        } catch (Exception e) {
            LOG.error(e.getMessage(), e);
            return Collections.emptyList();
        }
    }

    /**
     * 删除指定文档。
     *
     * @param id 文档 ID
     * @return 删除成功时返回 {@code true}
     */
    @Override
    public boolean deleteDocument(Object id) {
        return deleteDocuments(Collections.singletonList(id));
    }

    /**
     * 使用有界 Bulk 请求批量删除文档。
     *
     * <p>删除不存在的文档按幂等成功处理；仅 Bulk 条目包含实际错误时返回失败。</p>
     *
     * @param ids 文档 ID 集合
     * @return 全部删除操作成功时返回 {@code true}
     */
    @Override
    public boolean deleteDocuments(Collection<?> ids) {
        if (ids == null || ids.isEmpty()) {
            return false;
        }
        List<String> documentIds = new ArrayList<>(ids.size());
        for (Object id : ids) {
            if (id == null) {
                return false;
            }
            documentIds.add(id.toString());
        }
        try {
            for (int start = 0, batchIndex = 0;
                 start < documentIds.size();
                 start += MAX_BULK_OPERATIONS, batchIndex++) {
                int end = Math.min(start + MAX_BULK_OPERATIONS, documentIds.size());
                List<String> batch = documentIds.subList(start, end);
                BulkResponse response = client.bulk(buildDeleteBulkRequest(batch));
                if (!isBulkSuccessful(response, "删除", batchIndex, batch.size())) {
                    return false;
                }
            }
            return true;
        } catch (Exception e) {
            LOG.error("Elasticsearch 批量删除异常: count={}", documentIds.size(), e);
            return false;
        }
    }

    /**
     * 构建批量删除请求。
     *
     * @param ids 文档 ID
     * @return Bulk 删除请求
     */
    BulkRequest buildDeleteBulkRequest(List<String> ids) {
        BulkRequest.Builder builder = new BulkRequest.Builder();
        for (String id : ids) {
            builder.operations(operation -> operation.delete(delete -> delete
                .index(esConfig.getIndexName())
                .id(id)
            ));
        }
        return builder.build();
    }

    private boolean isBulkSuccessful(BulkResponse response,
                                     String operation,
                                     int batchIndex,
                                     int batchSize) {
        if (response != null && !response.errors()) {
            return true;
        }
        if (response == null) {
            LOG.error("Elasticsearch 批量{}未返回结果: batchIndex={}, batchSize={}",
                operation, batchIndex, batchSize);
            return false;
        }
        response.items().stream()
            .filter(item -> item.error() != null)
            .forEach(item -> LOG.error(
                "Elasticsearch 批量{}失败: batchIndex={}, index={}, id={}, status={}, reason={}",
                operation,
                batchIndex,
                item.index(),
                item.id(),
                item.status(),
                item.error().reason()
            ));
        return false;
    }

    /**
     * 更新指定文档。
     *
     * @param document 待更新文档
     * @return 更新成功时返回 {@code true}
     */
    @Override
    public boolean updateDocument(Document document) {
        if (document == null || document.getId() == null) {
            return false;
        }

        try {
            UpdateRequest<Map<String, Object>, Map<String, Object>> request = UpdateRequest.of(u -> u
                .index(esConfig.getIndexName())
                .id(document.getId().toString())
                .doc(buildSource(document))
            );

            @SuppressWarnings("unchecked")
            Class<Map<String, Object>> documentClass = (Class<Map<String, Object>>) (Class<?>) Map.class;
            UpdateResponse<Map<String, Object>> response = client.update(request, documentClass);
            return response.result() == co.elastic.clients.elasticsearch._types.Result.Updated
                || response.result() == co.elastic.clients.elasticsearch._types.Result.NoOp;
        } catch (Exception e) {
            LOG.error("Error updating document with id: " + document.getId(), e);
            return false;
        }
    }

    private static void closeQuietly(AutoCloseable closeable) {
        if (closeable == null) {
            return;
        }
        try {
            closeable.close();
        } catch (Exception ignored) {
            // 初始化失败时仅执行尽力清理，原始异常由构造方法继续抛出。
        }
    }

    @SuppressWarnings("unchecked")
    private Document toDocument(String hitId, Map source, Double score) {
        if (source == null || source.isEmpty()) {
            return null;
        }

        Document document = new Document();
        Object id = source.get("id");
        document.setId(id != null ? id : hitId);

        Object title = source.get("title");
        if (title != null) {
            document.setTitle(String.valueOf(title));
        }

        Object content = source.get("content");
        if (content != null) {
            document.setContent(String.valueOf(content));
        }

        Object metadataMap = source.get("metadataMap");
        if (metadataMap instanceof Map<?, ?>) {
            document.setMetadataMap(new HashMap<>((Map<String, Object>) metadataMap));
        }

        document.setScore(score);
        return document;
    }

    Map<String, Object> buildSource(Document document) {
        Map<String, Object> source = new HashMap<String, Object>();
        source.put("id", document.getId());
        source.put("content", document.getContent());
        if (document.getTitle() != null) {
            source.put("title", document.getTitle());
        }
        if (document.getMetadataMap() != null && !document.getMetadataMap().isEmpty()) {
            source.put("metadataMap", new HashMap<String, Object>(document.getMetadataMap()));
            Object knowledgeId = document.getMetadata(KeywordSearchMetadataKeys.KNOWLEDGE_ID);
            if (knowledgeId != null) {
                source.put(KeywordSearchMetadataKeys.KNOWLEDGE_ID, String.valueOf(knowledgeId));
            }
        }
        return source;
    }

    SearchRequest buildSearchRequest(KeywordSearchRequest request) {
        KeywordSearchRequest effectiveRequest = request == null ? new KeywordSearchRequest() : request;
        return SearchRequest.of(s -> s
            .index(esConfig.getIndexName())
            .size(effectiveRequest.getCount())
            .source(SourceConfig.of(sc -> sc.filter(f -> f.includes("id", "title", "content", "metadataMap"))))
            .query(q -> q.bool(b -> {
                b.must(m -> m.multiMatch(mm -> mm
                    .query(effectiveRequest.getKeyword())
                    .fields("title", "content")
                ));
                if (effectiveRequest.getKnowledgeId() != null && !effectiveRequest.getKnowledgeId().trim().isEmpty()) {
                    b.filter(f -> f.term(t -> t
                        .field(KeywordSearchMetadataKeys.KNOWLEDGE_ID)
                        .value(v -> v.stringValue(effectiveRequest.getKnowledgeId().trim()))
                    ));
                }
                return b;
            }))
        );
    }

    /**
     * 检查 Elasticsearch 服务是否可用。
     *
     * @return 服务可访问时返回 {@code true}
     */
    public boolean checkAvailable() {
        try {
            return client.info() != null;
        } catch (Exception e) {
            LOG.error("Elasticsearch availability check failed", e);
            return false;
        }
    }

    /**
     * 关闭共享传输层及其底层 RestClient。
     *
     * @throws IOException 关闭客户端资源失败时抛出
     */
    @Override
    public void close() throws IOException {
        transport.close();
    }
}
