/*
 * 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.graph.node;

import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class QuestionClassifierNode implements NodeAction {

	private static final String CLASSIFIER_PROMPT_TEMPLATE = """
				### Job Description',
				You are a text classification engine that analyzes text data and assigns categories based on user input or automatically determined categories.
				### Task
				Your task is to assign one category ONLY to the input text and only one category can be  returned in the output. Additionally, you need to extract the key words from the text that are related to the classification.
				### Format
				The input text is: {inputText}. Categories are specified as a category list: {categories}. Classification instructions may be included to improve the classification accuracy: {classificationInstructions}.
				### Constraint
				DO NOT include anything other than the JSON array in your response.
			""";

	private static final String QUESTION_CLASSIFIER_USER_PROMPT_1 = """
				{ "input_text": ["I recently had a great experience with your company. The service was prompt and the staff was very friendly."],
				"categories": ["Customer Service", "Satisfaction", "Sales", "Product"],
				"classification_instructions": ["classify the text based on the feedback provided by customer"]}
			""";

	private static final String QUESTION_CLASSIFIER_ASSISTANT_PROMPT_1 = """
				```json
					{"keywords": ["recently", "great experience", "company", "service", "prompt", "staff", "friendly"]
					"category_name": "Customer Service"}
				```
			""";

	private static final String QUESTION_CLASSIFIER_USER_PROMPT_2 = """
				{"input_text": ["bad service, slow to bring the food"],
				"categories": ["Food Quality", "Experience", "Price"],
				"classification_instructions": []}
			""";

	private static final String QUESTION_CLASSIFIER_ASSISTANT_PROMPT_2 = """
				```json
					{"keywords": ["bad service", "slow", "food", "tip", "terrible", "waitresses"],
					"category_name": "Experience"}
				```
			""";

	private static final Pattern VAR_TEMPLATE_PATTERN = Pattern.compile("\\{(\\w+)}");

	private final SystemPromptTemplate systemPromptTemplate;

	private final ChatClient chatClient;

	private String inputText;

	// Categories may contain {} placeholder variables that need to be processed
	private final Map<String, String> categories;

	// Classification instructions may contain {} placeholder variables that need to be processed
	private final List<String> classificationInstructions;

	private final String inputTextKey;

	private final String outputKey;

	public QuestionClassifierNode(ChatClient chatClient, String inputTextKey, Map<String, String> categories,
			List<String> classificationInstructions, String outputKey) {
		this.chatClient = chatClient;
		this.inputTextKey = inputTextKey;
		this.categories = categories;
		this.classificationInstructions = classificationInstructions;
		this.systemPromptTemplate = new SystemPromptTemplate(CLASSIFIER_PROMPT_TEMPLATE);
		this.outputKey = outputKey;
	}

	private String renderTemplate(OverAllState state, String template) {
		if (!StringUtils.hasText(template)) {
			return template;
		}
		Map<String, Object> params = Stream.of(template)
			.map(VAR_TEMPLATE_PATTERN::matcher)
			.map(Matcher::results)
			.map(results -> results.collect(Collectors.toUnmodifiableMap(r -> r.group(1),
					r -> state.value(r.group(1)).orElse(""), (a, b) -> b)))
			.findFirst()
			.orElseThrow();
		return new PromptTemplate(template).render(params);
	}

	private List<String> renderTemplates(OverAllState state, List<String> templates) {
		return templates.stream().map(template -> renderTemplate(state, template)).toList();
	}

	@Override
	public Map<String, Object> apply(OverAllState state) throws Exception {
		if (StringUtils.hasLength(inputTextKey)) {
			this.inputText = (String) state.value(inputTextKey).orElse(this.inputText);
		}

		List<Message> messages = new ArrayList<>();
		UserMessage userMessage1 = new UserMessage(QUESTION_CLASSIFIER_USER_PROMPT_1);
		AssistantMessage assistantMessage1 = new AssistantMessage(QUESTION_CLASSIFIER_ASSISTANT_PROMPT_1);
		UserMessage userMessage2 = new UserMessage(QUESTION_CLASSIFIER_USER_PROMPT_2);
		AssistantMessage assistantMessage2 = new AssistantMessage(QUESTION_CLASSIFIER_ASSISTANT_PROMPT_2);
		messages.add(userMessage1);
		messages.add(assistantMessage1);
		messages.add(userMessage2);
		messages.add(assistantMessage2);

		Map<String, String> renderedCategories = categories.entrySet()
			.stream()
			.map(e -> Map.entry(e.getKey(), renderTemplate(state, e.getValue())))
			.collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, Map.Entry::getValue));

		List<String> categoriesList = renderedCategories.values().stream().toList();

		ChatResponse response = chatClient.prompt()
			.system(systemPromptTemplate.render(Map.of("inputText", inputText, "categories", categoriesList,
					"classificationInstructions", renderTemplates(state, classificationInstructions))))
			.user(inputText)
			.messages(messages)
			.call()
			.chatResponse();

		Map<String, Object> updatedState = new HashMap<>();
		String output = Optional
			.ofNullable(Optional.ofNullable(response)
				.orElseThrow(() -> new RuntimeException("chat response is null"))
				.getResult()
				.getOutput()
				.getText())
			.orElseThrow(() -> new RuntimeException("chat response text is null"));
		String result = renderedCategories.entrySet()
			.stream()
			.filter(entry -> output.contains(entry.getValue()))
			.map(Map.Entry::getKey)
			.findFirst()
			.orElseThrow(() -> new RuntimeException(
					"chatClient returns [" + output + "], but it does not belong to the given category."));

		updatedState.put(outputKey, result);
		if (state.value("messages").isPresent()) {
			updatedState.put("messages", response.getResult().getOutput());
		}

		return updatedState;
	}

	public static Builder builder() {
		return new Builder();
	}

	public static class Builder {

		private String inputTextKey;

		private ChatClient chatClient;

		private Map<String, String> categories;

		private List<String> classificationInstructions;

		private String outputKey;

		public Builder inputTextKey(String input) {
			this.inputTextKey = input;
			return this;
		}

		public Builder chatClient(ChatClient chatClient) {
			this.chatClient = chatClient;
			return this;
		}

		/**
		 * Requires a Map object where key is category ID and value is the specific category name.
		 * Category names can contain variable placeholders like {@code {var_name}}, which will be replaced with values from {@code OverAllState}.
		 */
		public Builder categories(Map<String, String> categories) {
			this.categories = categories;
			return this;
		}

		// Compatible with old version
		@Deprecated
		public Builder categories(List<String> categories) {
			this.categories = categories.stream().collect(Collectors.toMap(k -> k, k -> k, (a, b) -> b));
			return this;
		}

		/**
		 * List of classification instructions. Instructions can contain variable placeholders like {@code {var_name}}, which will be replaced with values from {@code OverAllState}.
		 */
		public Builder classificationInstructions(List<String> classificationInstructions) {
			this.classificationInstructions = classificationInstructions;
			return this;
		}

		public Builder outputKey(String outputKey) {
			this.outputKey = outputKey;
			return this;
		}

		public QuestionClassifierNode build() {
			return new QuestionClassifierNode(chatClient, inputTextKey, categories, classificationInstructions,
					outputKey);
		}

	}

}
