package com.easyagents.engine.es;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch.core.*;
import co.elastic.clients.elasticsearch.core.bulk.BulkOperation;
import co.elastic.clients.elasticsearch.core.bulk.IndexOperation;
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.security.KeyManagementException;
import java.security.NoSuchAlgorithmException;
import java.security.cert.X509Certificate;
import java.util.*;

public class ElasticSearcher implements DocumentSearcher {

    private static final Logger LOG = LoggerFactory.getLogger(ElasticSearcher.class);

    private final ESConfig esConfig;

    public ElasticSearcher(ESConfig esConfig) {
        this.esConfig = esConfig;
    }

    // 忽略 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
     */
    @Override
    public boolean addDocument(Document document) {
        if (document == null || document.getContent() == null) {
            return false;
        }

        RestClient restClient = null;
        ElasticsearchTransport transport = null;
        try {
            restClient = buildRestClient();
            transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
            ElasticsearchClient client = new ElasticsearchClient(transport);

            Map<String, Object> source = buildSource(document);
            String documentId = document.getId().toString();
            IndexOperation<?> indexOp = IndexOperation.of(i -> i
                .index(esConfig.getIndexName())
                .id(documentId)
                .document(JsonData.of(source))
            );

            BulkOperation bulkOp = BulkOperation.of(b -> b.index(indexOp));
            BulkRequest request = BulkRequest.of(b -> b.operations(Collections.singletonList(bulkOp)));
            BulkResponse response = client.bulk(request);
            return !response.errors();

        } catch (Exception e) {
            LOG.error(e.getMessage(), e);
            return false;
        } finally {
            closeResources(transport, restClient);
        }
    }

    @Override
    public List<Document> searchDocuments(KeywordSearchRequest request) {
        RestClient restClient = null;
        ElasticsearchTransport transport = null;

        try {
            restClient = buildRestClient();
            transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
            ElasticsearchClient client = new ElasticsearchClient(transport);

            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();
        } finally {
            closeResources(transport, restClient);
        }
    }

    @Override
    public boolean deleteDocument(Object id) {
        if (id == null) {
            return false;
        }

        RestClient restClient = null;
        ElasticsearchTransport transport = null;
        try {
            restClient = buildRestClient();
            transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
            ElasticsearchClient client = new ElasticsearchClient(transport);

            DeleteRequest request = DeleteRequest.of(d -> d
                .index(esConfig.getIndexName())
                .id(id.toString())
            );

            DeleteResponse response = client.delete(request);
            return response.result() == co.elastic.clients.elasticsearch._types.Result.Deleted;

        } catch (Exception e) {
            LOG.error("Error deleting document with id: " + id, e);
            return false;
        } finally {
            closeResources(transport, restClient);
        }
    }

    @Override
    public boolean updateDocument(Document document) {
        if (document == null || document.getId() == null) {
            return false;
        }

        RestClient restClient = null;
        ElasticsearchTransport transport = null;

        try {
            restClient = buildRestClient();
            transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
            ElasticsearchClient client = new ElasticsearchClient(transport);

            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;
        } finally {
            closeResources(transport, restClient);
        }
    }


    private void closeResources(AutoCloseable... closeables) {
        for (AutoCloseable closeable : closeables) {
            try {
                if (closeable != null)
                    closeable.close();
            } catch (Exception e) {
                LOG.error("Error closing resource", e);
            }
        }
    }

    @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;
            }))
        );
    }

    public boolean checkAvailable() {
        RestClient restClient = null;
        ElasticsearchTransport transport = null;
        try {
            restClient = buildRestClient();
            transport = new RestClientTransport(restClient, new JacksonJsonpMapper());
            ElasticsearchClient client = new ElasticsearchClient(transport);
            return client.info() != null;
        } catch (Exception e) {
            LOG.error("Elasticsearch availability check failed", e);
            return false;
        } finally {
            closeResources(transport, restClient);
        }
    }
}
