package com.lanyuanxiaoyao.service.ai.knowledge.service;

import club.kingon.sql.builder.SqlBuilder;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import com.lanyuanxiaoyao.service.ai.knowledge.entity.Knowledge;
import com.lanyuanxiaoyao.service.ai.knowledge.entity.vo.KnowledgeVO;
import com.lanyuanxiaoyao.service.common.Constants;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;
import org.eclipse.collections.api.factory.Lists;
import org.eclipse.collections.api.list.ImmutableList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

/**
 * @author lanyuanxiaoyao
 * @version 20250522
 */
@Service
public class KnowledgeService {
  public static final String KNOWLEDGE_TABLE_NAME = Constants.DATABASE_NAME + ".service_ai_knowledge";
  private static final Logger logger = LoggerFactory.getLogger(KnowledgeService.class);
  private static final RowMapper<Knowledge> knowledgeMapper = (rs, row) -> {
    Knowledge knowledge = new Knowledge();
    knowledge.setId(rs.getLong(1));
    knowledge.setVectorSourceId(rs.getLong(2));
    knowledge.setName(rs.getString(3));
    knowledge.setStrategy(rs.getString(4));
    knowledge.setCreatedTime(rs.getTimestamp(5).getTime());
    knowledge.setModifiedTime(rs.getTimestamp(6).getTime());
    return knowledge;
  };
  private final JdbcTemplate template;
  private final EmbeddingModel embeddingModel;
  private final QdrantClient client;
  private final GroupService groupService;

  public KnowledgeService(JdbcTemplate template, EmbeddingModel embeddingModel, VectorStore vectorStore, GroupService groupService) {
    this.template = template;
    this.embeddingModel = embeddingModel;
    this.client = (QdrantClient) vectorStore.getNativeClient().orElseThrow();
    this.groupService = groupService;
  }

  public Knowledge get(Long id) {
    return template.queryForObject(
        SqlBuilder.select("id", "vector_source_id", "name", "strategy", "created_time", "modified_time")
            .from(KNOWLEDGE_TABLE_NAME)
            .whereEq("id", "?")
            .precompileSql(),
        knowledgeMapper,
        id
    );
  }

  @Transactional(rollbackFor = Exception.class)
  public void add(String name, String strategy) throws ExecutionException, InterruptedException {
    Integer count = template.queryForObject(
        SqlBuilder.select("count(*)")
            .from(KNOWLEDGE_TABLE_NAME)
            .whereEq("name", "?")
            .precompileSql(),
        Integer.class,
        name
    );
    if (count > 0) {
      throw new RuntimeException("名称已存在");
    }

    long id = IdUtil.getSnowflakeNextId();
    long vectorSourceId = IdUtil.getSnowflakeNextId();
    template.update(
        SqlBuilder.insertInto(KNOWLEDGE_TABLE_NAME, "id", "vector_source_id", "name", "strategy")
            .values()
            .addValue("?", "?", "?", "?")
            .precompileSql(),
        id,
        vectorSourceId,
        name,
        strategy
    );
    client.createCollectionAsync(
        String.valueOf(vectorSourceId),
        Collections.VectorParams.newBuilder()
            .setDistance(Collections.Distance.valueOf(strategy))
            .setSize(embeddingModel.dimensions())
            .build()
    ).get();
  }

  public String getName(Long id) {
    return template.queryForObject(
        SqlBuilder.select("name")
            .from(KNOWLEDGE_TABLE_NAME)
            .whereEq("id", id)
            .orderByDesc("created_time")
            .build(),
        String.class
    );
  }

  public ImmutableList<KnowledgeVO> list() {
    return template.query(
            SqlBuilder.select("id", "vector_source_id", "name", "strategy", "created_time", "modified_time")
                .from(KNOWLEDGE_TABLE_NAME)
                .orderByDesc("created_time")
                .build(),
            knowledgeMapper
        )
        .stream()
        .map(knowledge -> {
          try {
            Collections.CollectionInfo info = client.getCollectionInfoAsync(String.valueOf(knowledge.getVectorSourceId())).get();
            KnowledgeVO vo = new KnowledgeVO();
            vo.setId(String.valueOf(knowledge.getId()));
            vo.setName(knowledge.getName());
            vo.setPoints(info.getPointsCount());
            vo.setSegments(info.getSegmentsCount());
            vo.setStatus(info.getStatus().name());
            Collections.VectorParams vectorParams = info.getConfig().getParams().getVectorsConfig().getParams();
            vo.setStrategy(vectorParams.getDistance().name());
            vo.setSize(vectorParams.getSize());
            vo.setCreatedTime(vo.getCreatedTime());
            vo.setModifiedTime(vo.getModifiedTime());
            return vo;
          } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
          }
        })
        .collect(Collectors.toCollection(Lists.mutable::empty))
        .toImmutable();
  }

  @Transactional(rollbackFor = Exception.class)
  public void remove(Long id) throws ExecutionException, InterruptedException {
    Knowledge knowledge = get(id);
    if (ObjectUtil.isNull(knowledge)) {
      throw new RuntimeException(StrUtil.format("{} 不存在"));
    }
    template.update(
        SqlBuilder.delete(KNOWLEDGE_TABLE_NAME)
            .whereEq("id", "?")
            .precompileSql(),
        knowledge.getId()
    );
    groupService.removeByKnowledgeId(knowledge.getId());
    client.deleteCollectionAsync(String.valueOf(knowledge.getVectorSourceId())).get();
  }
}