package com.lingyue.graph.service;

import com.lingyue.common.exception.ServiceException;
import com.lingyue.graph.entity.TextChunk;
import com.lingyue.graph.entity.VectorEmbedding;
import com.lingyue.graph.repository.VectorEmbeddingRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.reactive.function.client.WebClient;

import java.util.*;

/**
 * Ollama Embedding 服务
 * 调用 Ollama API 进行文本向量化
 *
 * @author lingyue
 * @since 2026-01-15
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class OllamaEmbeddingService {

    private final VectorEmbeddingRepository vectorEmbeddingRepository;
    private final WebClient ollamaWebClient;

    @Value("${ollama.embedding.model:nomic-embed-text}")
    private String embeddingModel;

    @Value("${rag.embedding.batch-size:10}")
    private int embeddingBatchSize;

    /**
     * 对单个文本进行向量化
     *
     * @param text 待向量化的文本
     * @return 向量（浮点数数组）
     */
    public float[] embed(String text) {
        if (text == null || text.trim().isEmpty()) {
            throw new ServiceException("文本不能为空");
        }

        try {
            long startTime = System.currentTimeMillis();
            log.info("调用 Ollama 向量化: model={}, textLength={}", embeddingModel, text.length());

            // 构建请求体
            Map<String, Object> request = new HashMap<>();
            request.put("model", embeddingModel);
            request.put("prompt", text);

            // 调用 Ollama API
            @SuppressWarnings("unchecked")
            Map<String, Object> response = ollamaWebClient
                    .post()
                    .uri("/api/embeddings")
                    .bodyValue(request)
                    .retrieve()
                    .bodyToMono(Map.class)
                    .block();

            if (response == null || !response.containsKey("embedding")) {
                throw new ServiceException("Ollama 返回结果无效");
            }

            // 解析向量
            @SuppressWarnings("unchecked")
            List<Double> embeddingList = (List<Double>) response.get("embedding");
            float[] vector = new float[embeddingList.size()];
            for (int i = 0; i < embeddingList.size(); i++) {
                vector[i] = embeddingList.get(i).floatValue();
            }

            long elapsedMs = System.currentTimeMillis() - startTime;
            log.info("Ollama 向量化完成: model={}, dim={}, elapsedMs={}",
                    embeddingModel, vector.length, elapsedMs);
            return vector;

        } catch (Exception e) {
            log.error("Ollama 向量化失败: {}", e.getMessage(), e);
            throw new ServiceException("向量化失败: " + e.getMessage(), e);
        }
    }

    /**
     * 将向量数组转换为 pgvector 格式字符串
     */
    public String vectorToString(float[] vector) {
        StringBuilder sb = new StringBuilder("[");
        for (int i = 0; i < vector.length; i++) {
            if (i > 0) {
                sb.append(",");
            }
            sb.append(vector[i]);
        }
        sb.append("]");
        return sb.toString();
    }

    /**
     * 对文本分块进行向量化并保存
     *
     * @param chunk 文本分块
     * @return 向量嵌入实体
     */
    @Transactional
    public VectorEmbedding embedAndSave(TextChunk chunk) {
        // 检查是否已存在
        if (vectorEmbeddingRepository.existsByChunkId(chunk.getId())) {
            log.debug("分块 {} 已有向量，跳过", chunk.getId());
            return vectorEmbeddingRepository.findByChunkId(chunk.getId());
        }

        // 向量化
        float[] vector = embed(chunk.getContent());
        String embeddingString = vectorToString(vector);

        // 保存
        VectorEmbedding embedding = new VectorEmbedding();
        embedding.setId(UUID.randomUUID().toString().replace("-", ""));
        embedding.setChunkId(chunk.getId());
        embedding.setEmbedding(embeddingString);
        embedding.setModelName(embeddingModel);
        embedding.setCreateTime(new Date());

        vectorEmbeddingRepository.insert(embedding);
        log.debug("分块 {} 向量化完成并保存", chunk.getId());

        return embedding;
    }

    /**
     * 批量向量化文本分块
     *
     * @param chunks 文本分块列表
     * @return 向量嵌入列表
     */
    @Transactional
    public List<VectorEmbedding> embedBatch(List<TextChunk> chunks) {
        List<VectorEmbedding> embeddings = new ArrayList<>();
        int processed = 0;
        int failed = 0;

        int batchSize = Math.max(1, embeddingBatchSize);
        for (int batchStart = 0; batchStart < chunks.size(); batchStart += batchSize) {
            int batchEnd = Math.min(batchStart + batchSize, chunks.size());
            List<TextChunk> batch = chunks.subList(batchStart, batchEnd);
            long batchStartTime = System.currentTimeMillis();
            log.info("向量化批次开始: batch={}-{}, batchSize={}, total={}",
                    batchStart + 1, batchEnd, batch.size(), chunks.size());

            for (TextChunk chunk : batch) {
                try {
                    VectorEmbedding embedding = embedAndSave(chunk);
                    embeddings.add(embedding);
                    processed++;
                } catch (Exception e) {
                    log.error("分块 {} 向量化失败: {}", chunk.getId(), e.getMessage());
                    // 继续处理其他分块
                    failed++;
                }
            }

            long batchElapsedMs = System.currentTimeMillis() - batchStartTime;
            log.info("向量化批次完成: processed={}/{} (成功: {}, 失败: {}, elapsedMs={})",
                    processed + failed, chunks.size(), processed, failed, batchElapsedMs);
        }

        log.info("批量向量化完成，成功 {}/{} 块", embeddings.size(), chunks.size());
        return embeddings;
    }

    /**
     * 根据分块ID删除向量
     */
    public void deleteByChunkId(String chunkId) {
        vectorEmbeddingRepository.deleteByChunkId(chunkId);
    }
}
