/*
 * 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.LLMNodeData;
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 org.springframework.stereotype.Component;

import java.util.List;

// TODO：支持异常分支、支持DashScope平台以外其他模型、Dify的结构化输出
@Component
public class LLMNodeSection implements NodeSection<LLMNodeData> {

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

	@Override
	public String render(Node node, String varName) {
		LLMNodeData nodeData = ((LLMNodeData) node.getData());
		return String.format("""
				// —— LLMNode [%s] ——
				stateGraph.addNode("%s", AsyncNodeAction.node_async(
				    createLLMNodeAction(chatModel, %s, %s, %s, %s, %s, %s, %s, %s, %s)
				));

				""", node.getId(), varName, ObjectToCodeUtil.toCode(nodeData.getChatModeName()),
				ObjectToCodeUtil.toCode(nodeData.getModeParams()),
				ObjectToCodeUtil.toCode(nodeData.getMessageTemplates()),
				ObjectToCodeUtil.toCode(nodeData.getMemoryKey()), ObjectToCodeUtil.toCode(nodeData.getMaxRetryCount()),
				ObjectToCodeUtil.toCode(nodeData.getRetryIntervalMs()),
				ObjectToCodeUtil.toCode(nodeData.getDefaultOutput()),
				ObjectToCodeUtil.toCode(nodeData.getErrorNextNode()),
				ObjectToCodeUtil.toCode(nodeData.getOutputKeyPrefix()));
	}

	@Override
	public String assistMethodCode(DSLDialectType dialectType) {

		return String.format(
				"""
						private record MessageTemplate(String template, List<String> keys, MessageType type) {
						    public Message render(OverAllState state) {
						        Map<String, Object> params = keys.stream()
						            .collect(Collectors.toMap(key -> key, key -> state.value(key, ""), (o1, o2) -> o2));
						        String text = new PromptTemplate(template).render(params);
						        return switch (type) {
						            case USER -> new UserMessage(text);
						            case SYSTEM -> new SystemMessage(text);
						            case TOOL -> throw new UnsupportedOperationException("Tool message not supported");
						            case ASSISTANT -> new AssistantMessage(text);
						        };
						    }
						}

						private NodeAction createLLMNodeAction(ChatModel chatModel,
						        String chatModelName, Map<String, Number> modeParams,
						        List<MessageTemplate> messageTemplates, String memoryKey, Integer maxRetryCount, Integer retryIntervalMs,
						        String defaultOutput, String errorNextNode, String outputKeyPrefix) {
						    // build chatClient with params
						    var chatOptionsBuilder = DashScopeChatOptions.builder().withModel(chatModelName);
						    Optional.ofNullable(modeParams.get("temperature")).ifPresent(val -> chatOptionsBuilder.withTemperature(val.doubleValue()));
						    Optional.ofNullable(modeParams.get("seed")).ifPresent(val -> chatOptionsBuilder.withSeed(val.intValue()));
						    Optional.ofNullable(modeParams.get("top_p")).ifPresent(val -> chatOptionsBuilder.withTopP(val.doubleValue()));
						    Optional.ofNullable(modeParams.get("top_k")).ifPresent(val -> chatOptionsBuilder.withTopK(val.intValue()));
						    Optional.ofNullable(modeParams.get("max_tokens")).ifPresent(val -> chatOptionsBuilder.withMaxToken(val.intValue()));
						    Optional.ofNullable(modeParams.get("repetition_penalty")).ifPresent(val -> chatOptionsBuilder.withRepetitionPenalty(val.doubleValue()));
						    final ChatClient chatClient = ChatClient.builder(chatModel).defaultOptions(chatOptionsBuilder.build()).build();

						    String nextNodeKey = "next_node";

						    return state -> {
						        String memories;
						        if (memoryKey == null || state.value(memoryKey).isEmpty()) {
						            memories = "This is the user's first request without context";
						        }
						        else {
						            memories = String.format("This is the history of previous requests:\\n %%s",
						                    state.value(memoryKey, List.of()).toString());
						        }

						        // call chatClient
						        int retryCount = Optional.ofNullable(maxRetryCount).orElse(1);
						        int retryInterval = Optional.ofNullable(retryIntervalMs).orElse(1000);
						        while (retryCount-- > 0) {
						            try {
						                // build messages
						                List<Message> messages = messageTemplates.stream()
						                    .map(messageTemplate -> messageTemplate.render(state))
						                    .toList();
						                String content = chatClient.prompt().system(memories).messages(messages).call().content();
						                if (content == null) {
						                    throw new RuntimeException("ChatClient error");
						                }
						                Map<String, Object> map = new HashMap<>(%s);
						                if (memoryKey != null) {
						                    map.put(memoryKey, content);
						                }
						                return map;
						            }
						            catch (Exception e) {
						                try {
						                    Thread.sleep(retryInterval);
						                } catch (InterruptedException ie) {
						                    Thread.currentThread().interrupt();
						                    break;
						                }
						            }
						        }

						        // error handling
						        if (defaultOutput != null) {
						            return %s;
						        }
						        else if (errorNextNode != null) {
						            return Map.of(nextNodeKey, errorNextNode);
						        }
						        else {
						            throw new IllegalStateException("No default output or error next node provided");
						        }
						    };
						}
						""",
				dialectType.equals(DSLDialectType.DIFY) ? "Map.of(outputKeyPrefix + \"text\", content)"
						: "Map.of(outputKeyPrefix + \"output\", content, outputKeyPrefix + \"reasoning_content\", content)",
				dialectType.equals(DSLDialectType.DIFY) ? "Map.of(outputKeyPrefix + \"text\", defaultOutput)"
						: "Map.of(outputKeyPrefix + \"output\", defaultOutput, outputKeyPrefix + \"reasoning_content\", defaultOutput)");
	}

	@Override
	public List<String> getImports() {
		return List.of("org.springframework.ai.chat.messages.Message",
				"org.springframework.ai.chat.messages.AssistantMessage",
				"org.springframework.ai.chat.messages.MessageType",
				"org.springframework.ai.chat.messages.SystemMessage",
				"org.springframework.ai.chat.messages.UserMessage",
				"com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions",
				"org.springframework.beans.factory.annotation.Autowired", "java.util.Optional");
	}

}
