/*
 * Copyright 2024-2026 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.alibaba.cloud.ai.studio.admin.builder.generator.service.generator.workflow.sections;

import com.alibaba.cloud.ai.studio.admin.builder.generator.model.workflow.Node;
import com.alibaba.cloud.ai.studio.admin.builder.generator.model.workflow.NodeType;
import com.alibaba.cloud.ai.studio.admin.builder.generator.model.workflow.nodedata.KnowledgeRetrievalNodeData;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.dsl.DSLDialectType;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.generator.workflow.NodeSection;

import com.alibaba.cloud.ai.studio.admin.builder.generator.utils.ObjectToCodeUtil;
import com.alibaba.cloud.ai.studio.core.config.StudioProperties;
import com.alibaba.cloud.ai.studio.core.rag.DocumentService;
import com.alibaba.cloud.ai.studio.runtime.domain.PagingList;
import com.alibaba.cloud.ai.studio.runtime.domain.knowledgebase.Document;
import com.alibaba.cloud.ai.studio.runtime.domain.knowledgebase.DocumentQuery;
import com.alibaba.cloud.ai.studio.runtime.enums.DocumentType;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;
import java.util.Optional;

// TODO: 支持其他格式的文档，如PDF、ZIP等
// TODO: 解析并应用RerankModel、EmbeddingModel配置
// TODO: 支持从OSS获取资源文件，或者在生成项目中从OSS获取资源文件
@Component
public class KnowledgeRetrievalNodeSection implements NodeSection<KnowledgeRetrievalNodeData> {

	@Override
	public boolean support(NodeType nodeType) {
		return NodeType.RETRIEVER.equals(nodeType);
	}

	// 用于获取Studio存储的文档
	private final DocumentService studioDocumentService;

	private final String studioStoragePath;

	public KnowledgeRetrievalNodeSection(@Autowired(required = false) DocumentService studioDocumentService,
			@Autowired(required = false) StudioProperties properties) {
		this.studioDocumentService = studioDocumentService;
		this.studioStoragePath = properties != null ? properties.getStoragePath() : null;
	}

	@Override
	public String render(Node node, String varName) {
		KnowledgeRetrievalNodeData nodeData = (KnowledgeRetrievalNodeData) node.getData();

		if (DSLDialectType.STUDIO.equals(nodeData.getDialectType())) {
			if (this.studioDocumentService == null || this.studioStoragePath == null) {
				throw new IllegalArgumentException(
						"The current mode does not support Studio's knowledge retrieval node code generation. Please start the complete StudioApplication class");
			}
			// 根据knowledgeBaseIds获取对应的资源文件
			List<ResourceFile> resourceFiles = Optional.ofNullable(nodeData.getKnowledgeBaseIds())
				.orElse(List.of())
				.stream()
				.map(kbId -> {
					PagingList<Document> getSize = this.studioDocumentService.listDocuments(kbId, new DocumentQuery());
					Long total = getSize.getTotal();
					DocumentQuery query = new DocumentQuery();
					query.setSize(total.intValue());
					PagingList<Document> pagingList = this.studioDocumentService.listDocuments(kbId, query);
					return pagingList.getRecords();
				})
				.flatMap(List::stream)
				.filter(Document::getEnabled)
				.filter(d -> StringUtils.hasText(d.getPath()))
				.map(document -> {
					// 文件类型
					String contentType = document.getMetadata().getContentType();
					// 存储形式
					DocumentType documentType = document.getType();
					// 存储路径
					String path = switch (documentType) {
						case FILE -> {
							{
								Path p = Path.of(studioStoragePath);
								Path resolvedPath = p.resolve(document.getPath()).normalize();

								// 安全检查：确保解析后的路径仍在允许的目录范围内
								if (!resolvedPath.startsWith(p.normalize())) {
									throw new SecurityException("非法路径访问尝试: " + document.getPath());
								}

								yield resolvedPath.toAbsolutePath().toString();
							}
						}
						case URL -> {
							// 对URL路径进行基本验证
							String urlPath = document.getPath();
							if (urlPath == null || urlPath.trim().isEmpty()) {
								throw new IllegalArgumentException("URL路径不能为空");
							}
							yield urlPath;
						}
						default ->
							throw new UnsupportedOperationException("unsupported document type: " + documentType);
					};
					String fileName = document.getName();
					// 构造文件记录
					return new ResourceFile(fileName, switch (documentType) {
						case FILE -> ResourceFile.Type.CLASS_PATH;
						case URL -> ResourceFile.Type.URL;
						default ->
							throw new UnsupportedOperationException("unsupported document type: " + documentType);
					}, () -> {
						try {
							return Files.newInputStream(Path.of(path));
						}
						catch (IOException e) {
							throw new RuntimeException(e);
						}
					});
				})
				.toList();
			nodeData.setResourceFiles(resourceFiles);
		}

		return String.format("""
				// —— KnowledgeRetrievalNode [%s] ——
				KnowledgeRetrievalNode %s = KnowledgeRetrievalNode.builder()
				    .topK(%s)
				    .similarityThreshold(%s)
				    .inputKey(%s)
				    .outputKey(%s)
				    .vectorStore(createVectorStore(%s))
				    .build();
				stateGraph.addNode("%s", AsyncNodeAction.node_async(wrapperRetrievalNodeAction(%s, "%s")));

				""", node.getId(), varName, ObjectToCodeUtil.toCode(nodeData.getTopK()),
				ObjectToCodeUtil.toCode(nodeData.getThreshold()), ObjectToCodeUtil.toCode(nodeData.getInputKey()),
				ObjectToCodeUtil.toCode(nodeData.getOutputKey()),
				ObjectToCodeUtil.toCode(DSLDialectType.STUDIO.equals(nodeData.getDialectType())
						? nodeData.getResourceFiles() : List.of("please_config_your_own_resource_files")),
				varName, varName, nodeData.getOutputKey());
	}

	@Override
	public String assistMethodCode(DSLDialectType dialectType) {
		StringBuilder sb = new StringBuilder();
		sb.append("""
				@Autowired
				private ResourceLoader resourceLoader;

				@Autowired
				private EmbeddingModel embeddingModel;

				""");
		if (!DSLDialectType.STUDIO.equals(dialectType)) {
			sb.append(
					"// todo: Please manually modify the parameter values passed to this method to point to the correct resource paths");
			sb.append(String.format("%n"));
		}
		sb.append("""
				private VectorStore createVectorStore(List<String> paths) {
				    List<Resource> resources = Optional.ofNullable(paths).orElse(List.of())
				            .stream().map(resourceLoader::getResource).toList();
				    List<Document> documents = resources.stream().map(TextReader::new).map(TextReader::read)
				            .flatMap(List::stream).toList();
				    List<Document> chunks = new TokenTextSplitter().transform(documents);
				    SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
				    vectorStore.write(chunks);
				    return vectorStore;
				}
				""");
		sb.append(switch (dialectType) {
			case DIFY ->
				"""
						 private NodeAction wrapperRetrievalNodeAction(NodeAction nodeAction, String key) {
						     return state -> {
						         // Convert the result to the variable format required by the Dify workflow
						         Map<String, Object> result = nodeAction.apply(state);
						         Object object = result.get(key);
						         if(object instanceof List<?> list && !list.isEmpty() && list.get(0) instanceof Document) {
						             // Return value is Array[Object] (using List<Map>)
						             List<Document> documentList = (List<Document>) list;
						             List<Map<String, Object>> mapList = documentList.stream().map(document ->
						                             Map.of("content", document.getFormattedContent(), "title", document.getId(), "url", "", "icon", "", "metadata", document.getMetadata()))
						                     .toList();
						             return Map.of(key, mapList);
						         } else {
						             return Map.of(key, List.of(Map.of("content", object.toString(), "title", "unknown", "url", "unknown", "icon", "unknown", "metadata", object)));
						         }
						     };
						 }
						""";
			case STUDIO ->
				"""
						   private NodeAction wrapperRetrievalNodeAction(NodeAction nodeAction, String key) {
						       return state -> {
						           // Convert the result to the variable format required by the workflow
						           Map<String, Object> result = nodeAction.apply(state);
						           Object object = result.get(key);
						           if (object instanceof List<?> list && !list.isEmpty() && list.get(0) instanceof Document) {
						               // Return value is Array[Object] (using List<Map>)
						               List<Document> documentList = (List<Document>) list;
						               List<Map<String, Object>> mapList = documentList.stream()
						                   .map(document -> Map.<String, Object>of("doc_id", document.getId(), "doc_name",
						                           Optional.ofNullable(document.getText()).orElse("unknown"), "title", document.getId(),
						                           "text", document.getFormattedContent(), "score",
						                           Optional.ofNullable(document.getScore()).orElse(0.0), "page_number", 0, "chunk_id",
						                           document.getId()))
						                   .toList();
						               return Map.of(key, mapList);
						           }
						           else {
						               return Map.of(key, List.of());
						           }
						       };
						   }
						""";
			default -> "";
		});
		return sb.toString();
	}

	@Override
	public List<ResourceFile> resourceFiles(DSLDialectType dialectType, KnowledgeRetrievalNodeData nodeData) {
		return Optional.ofNullable(nodeData.getResourceFiles())
			.orElse(NodeSection.super.resourceFiles(dialectType, nodeData));
	}

	@Override
	public List<String> getImports() {
		return List.of("com.alibaba.cloud.ai.graph.node.KnowledgeRetrievalNode",
				"org.springframework.ai.embedding.EmbeddingModel", "org.springframework.ai.reader.TextReader",
				"org.springframework.ai.transformer.splitter.TokenTextSplitter",
				"org.springframework.ai.vectorstore.SimpleVectorStore",
				"org.springframework.ai.vectorstore.VectorStore", "org.springframework.beans.factory.annotation.Value",
				"org.springframework.core.io.Resource", "org.springframework.ai.document.Document",
				"org.springframework.beans.factory.annotation.Autowired", "org.springframework.core.io.ResourceLoader",
				"java.util.Optional");
	}

}
