package top.tqx.demo_1.service.impl;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import top.tqx.demo_1.config.AiApiConfig;
import top.tqx.demo_1.dto.SearchRequest;
import top.tqx.demo_1.dto.SearchResponse;
import top.tqx.demo_1.service.RagService;
import top.tqx.demo_1.service.SearchService;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
public class SearchServiceImpl implements SearchService {

    @Autowired
    private AiApiConfig aiApiConfig;

    @Autowired
    private RagService ragService;

    @Override
    public SearchResponse search(SearchRequest request) throws IOException {
        String url = aiApiConfig.getFullUrl(aiApiConfig.getSearch().getSearch());
        log.info("调用RAG服务搜索: url={}", url);

        try {
            ObjectNode requestBody = ragService.createObjectNode();
            if (request.getQuery() != null) {
                requestBody.put("query", request.getQuery());
            }
            if (request.getTopK() != null) {
                requestBody.put("top_k", request.getTopK());
            }
            if (request.getCollections() != null && !request.getCollections().isEmpty()) {
                com.fasterxml.jackson.databind.node.ArrayNode arrayNode = requestBody.putArray("collections");
                for (String collection : request.getCollections()) {
                    arrayNode.add(collection);
                }
            }

            String jsonResponse = ragService.postJsonRequest(url, requestBody.toString(), 30000);
            JsonNode rootNode = ragService.parseJson(jsonResponse);
            log.info("RAG服务搜索成功");

            SearchResponse response = new SearchResponse();

            if (rootNode.has("contexts") && rootNode.get("contexts").isArray()) {
                List<String> contexts = new ArrayList<>();
                for (JsonNode contextNode : rootNode.get("contexts")) {
                    contexts.add(contextNode.asText());
                }
                response.setContexts(contexts);
            }

            if (rootNode.has("metadatas") && rootNode.get("metadatas").isArray()) {
                List<Map<String, Object>> metadatas = new ArrayList<>();
                for (JsonNode metadataNode : rootNode.get("metadatas")) {
                    Map<String, Object> metadata = new HashMap<>();
                    metadataNode.fieldNames().forEachRemaining(key -> {
                        JsonNode value = metadataNode.get(key);
                        if (value.isInt()) {
                            metadata.put(key, value.asInt());
                        } else if (value.isLong()) {
                            metadata.put(key, value.asLong());
                        } else if (value.isDouble()) {
                            metadata.put(key, value.asDouble());
                        } else if (value.isBoolean()) {
                            metadata.put(key, value.asBoolean());
                        } else {
                            metadata.put(key, value.asText());
                        }
                    });
                    metadatas.add(metadata);
                }
                response.setMetadatas(metadatas);
            }

            if (rootNode.has("scores") && rootNode.get("scores").isArray()) {
                List<Double> scores = new ArrayList<>();
                for (JsonNode scoreNode : rootNode.get("scores")) {
                    scores.add(scoreNode.asDouble());
                }
                response.setScores(scores);
            }

            return response;

        } catch (Exception e) {
            log.error("调用RAG服务搜索失败: {}", e.getMessage(), e);
            return new SearchResponse();
        }
    }
}
