package io.github.qifan777.knowledge.ai.agent.documentSearch;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.annotation.JsonPropertyDescription;
import io.github.qifan777.knowledge.domain.po.SearchDocument;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Description;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * <p>
 *
 * </p>
 *
 * @author taohongrun
 * @since 2025/3/17
 */
@Component
@Description("返回用户想要查询的文档连接")
@Slf4j
@RequiredArgsConstructor
public class SearchDocumentFunction implements Function<SearchDocumentFunction.Request, String> {
    private final ElasticsearchClient client;

    @Override
    public String apply(SearchDocumentFunction.Request request) {
        //获取的参数
        log.info("用户想要查询的文档名称: {}", request.fileName);
        log.info("用户的角色类型: {}", request.roleName);

        try {
            List<String> urlList = new ArrayList<>();

            urlList = client.search(s -> s.index("documents_index")
                                    .query(q -> q
                                            .bool(b -> b
                                                    .must(m -> m
                                                            .match(t -> t
                                                                    .field("title")
                                                                    .query(request.fileName)))
                                                    //"permissions"为数组，判断是否包含roleName

                                                    .filter(m -> m
                                                            .term(t -> t
                                                                    .field("permissions")
                                                                    .value(request.roleName)))
                                            ))
                            , SearchDocument.class)
                    .hits()
                    .hits()
                    .stream()
                    .map(h -> {
                        assert h.source() != null;
                        return h.source().getUrl();
                    })
                    .collect(Collectors.toList());
            log.info("查询结果: {}", urlList);
            return String.join("\n", urlList);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public record Request(@JsonProperty(required = true)
                          @JsonPropertyDescription("用户想要查询的文档名称")
                          String fileName
            ,
                          @JsonProperty(required = true)
                          @JsonPropertyDescription("用户的角色类型,如老师，学生")
                          String roleName) {
    }
}
