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

import com.alibaba.cloud.ai.graph.agent.extension.interceptor.AssistantMessageSanitizerInterceptor;
import com.alibaba.cloud.ai.graph.agent.hook.Hook;
import com.alibaba.cloud.ai.graph.agent.interceptor.Interceptor;
import com.alibaba.cloud.ai.graph.agent.interceptor.ModelInterceptor;
import com.alibaba.cloud.ai.graph.agent.interceptor.ToolInterceptor;
import com.alibaba.cloud.ai.graph.agent.node.AgentLlmNode;
import com.alibaba.cloud.ai.graph.agent.node.AgentToolNode;
import com.alibaba.cloud.ai.graph.agent.tool.ToolCallbackUtils;
import io.micrometer.observation.ObservationRegistry;
import org.apache.commons.collections4.CollectionUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.converter.FormatProvider;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.execution.DefaultToolExecutionExceptionProcessor;
import org.springframework.util.StringUtils;

import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class DefaultBuilder extends Builder {

	private static final Logger logger = LoggerFactory.getLogger(DefaultBuilder.class);

	public static final String POSSIBLE_LLM_TOOL_NAME_CHANGE_WARNING
				= "LLM may have adapted the tool name '{}', especially if the name was truncated due to length limits. If this is the case, you can customize the prefixing and processing logic using McpToolNamePrefixGenerator";

	@Override
	public ReactAgent build() {

		// Validate name is not empty
		if (!StringUtils.hasText(this.name)) {
			throw new IllegalArgumentException("Agent name must not be empty");
		}

		// Validate either chatClient or model is provided
		if (chatClient == null && model == null) {
			throw new IllegalArgumentException("Either chatClient or model must be provided");
		}

		// Get source options from ChatClient (when provided) or from ChatModel (when building from model).
		ChatOptions sourceOptions = (chatClient != null)
				? getChatClientDefaultOptions(chatClient)
				: getChatModelDefaultOptions(model);
		ChatOptions effectiveOptions = mergeSourceOptionsWithAgentOptions(sourceOptions, this.chatOptions);

		if (chatClient == null) {
			ChatClient.Builder clientBuilder = ChatClient.builder(model,
					this.observationRegistry == null ? ObservationRegistry.NOOP : this.observationRegistry,
					this.customObservationConvention, this.advisorObservationConvention);
			if (effectiveOptions != null) {
				clientBuilder.defaultOptions(effectiveOptions);
			}
			chatClient = clientBuilder.build();
		} else {
			chatClient = chatClient.mutate().defaultOptions(effectiveOptions).build();
		}

		AgentLlmNode.Builder llmNodeBuilder = AgentLlmNode.builder()
				.agentName(this.name)
				.chatOptions(effectiveOptions)
				.chatClient(chatClient);

		if (outputKey != null && !outputKey.isEmpty()) {
			llmNodeBuilder.outputKey(outputKey);
		}

		if (systemPrompt != null) {
			llmNodeBuilder.systemPrompt(systemPrompt);
		}

		if (templateRenderer != null) {
			llmNodeBuilder.templateRenderer(templateRenderer);
    }
    
		if (instruction != null) {
			llmNodeBuilder.instruction(instruction);
		}

		String outputSchema = null;
		if (StringUtils.hasLength(this.outputSchema) ) {
			outputSchema = this.outputSchema;
		} else if (this.outputType != null) {
			FormatProvider formatProvider = new BeanOutputConverter<>(this.outputType);
			outputSchema = formatProvider.getFormat();
		}

		if (StringUtils.hasLength(outputSchema)) {
			llmNodeBuilder.outputSchema(outputSchema);
		}
		
		separateInterceptorsByType();

		registerDefaultModelInterceptors();

		List<ToolCallback> allTools = gatherLocalTools();
		
		// Set combined tools to LLM node
		if (CollectionUtils.isNotEmpty(allTools)) {
			llmNodeBuilder.toolCallbacks(Collections.unmodifiableList(allTools));
		}

		if (enableLogging) {
			llmNodeBuilder.enableReasoningLog(true);
		}

		AgentLlmNode llmNode = llmNodeBuilder.build();

		// Setup tool node with all available tools
		AgentToolNode toolNode;
		AgentToolNode.Builder toolBuilder = AgentToolNode.builder()
				.agentName(this.name)
				.parallelToolExecution(this.parallelToolExecution)
				.maxParallelTools(this.maxParallelTools)
				.toolExecutionTimeout(this.toolExecutionTimeout)
				.wrapSyncToolsAsAsync(this.wrapSyncToolsAsAsync);

		if (resolver != null) {
			toolBuilder.toolCallbackResolver(resolver);
		}
		if (CollectionUtils.isNotEmpty(allTools)) {
			toolBuilder.toolCallbacks(allTools);
		}

		if (enableLogging) {
			toolBuilder.enableActingLog(true);
		}
		if (toolExecutionExceptionProcessor == null) {
			toolBuilder.toolExecutionExceptionProcessor(DefaultToolExecutionExceptionProcessor.builder()
					.alwaysThrow(false)
					.build());
		} else {
			toolBuilder.toolExecutionExceptionProcessor(toolExecutionExceptionProcessor);
		}

		if (toolContext != null && !toolContext.isEmpty()) {
			toolBuilder.toolContext(toolContext);
		}

		toolNode = toolBuilder.build();

		return new ReactAgent(llmNode, toolNode, buildConfig(), this);
	}
	
	/**
	 * Register framework-provided model interceptors that are enabled by default.
	 * <p>
	 * The {@link AssistantMessageSanitizerInterceptor} is appended last so it sits innermost
	 * in the interceptor chain (closest to the model call). This way it normalizes the message
	 * list only after any user-configured interceptors have run, right before the request is
	 * sent to the model. Registration is skipped when disabled via
	 * {@link Builder#assistantMessageSanitizerEnabled(boolean)} or when an interceptor with the
	 * same name is already present.
	 */
	protected void registerDefaultModelInterceptors() {
		if (!this.assistantMessageSanitizerEnabled) {
			return;
		}
		if (modelInterceptors == null) {
			modelInterceptors = new ArrayList<>();
		}
		boolean alreadyPresent = modelInterceptors.stream()
				.anyMatch(interceptor -> AssistantMessageSanitizerInterceptor.class.equals(interceptor.getClass()));
		if (!alreadyPresent) {
			modelInterceptors.add(AssistantMessageSanitizerInterceptor.builder().build());
		}
	}

	/**
	 * Separate unified interceptors by type into modelInterceptors and toolInterceptors.
	 */
	protected void separateInterceptorsByType() {
		if (CollectionUtils.isNotEmpty(interceptors)) {
			modelInterceptors = new ArrayList<>();
			toolInterceptors = new ArrayList<>();
			
			for (Interceptor interceptor : interceptors) {
				if (interceptor instanceof ModelInterceptor) {
					modelInterceptors.add((ModelInterceptor) interceptor);
				}
				if (interceptor instanceof ToolInterceptor) {
					toolInterceptors.add((ToolInterceptor) interceptor);
				}
			}
		}
	}

	/**
	 * Collect all tools from various sources.
	 * <p>
	 * This method gathers tools from multiple sources in the following order:
	 * <ol>
	 *   <li>Direct tools provided via {@code tools}</li>
	 *   <li>Tools from {@code toolCallbackProviders}</li>
	 *   <li>Tools resolved by name via {@code toolNames}</li>
	 *   <li>Tools extracted from {@code resolver} (if no regular tools collected yet)</li>
	 *   <li>Tools from interceptors ({@code modelInterceptors})</li>
	 *   <li>Tools from hooks ({@code hooks})</li>
	 * </ol>
	 * <p>
	 * The final combined list prioritizes hook tools first, then interceptor tools, followed by regular tools.
	 *
	 * @return a combined list of all tools from hooks, interceptors and user-provided sources
	 */
	protected List<ToolCallback> gatherLocalTools() {
		// Collect tools from interceptors
		// - regularTools: user-provided tools
		// - interceptorTools: tools from interceptors
		List<ToolCallback> regularTools = new ArrayList<>();
		
		// Extract regular tools from user-provided tools
		if (CollectionUtils.isNotEmpty(tools)) {
			regularTools.addAll(tools);
		}
		
		if (CollectionUtils.isNotEmpty(toolCallbackProviders)) {
			for (var provider : toolCallbackProviders) {
				regularTools.addAll(List.of(provider.getToolCallbacks()));
			}
		}
		
		if (CollectionUtils.isNotEmpty(toolNames)) {
			for (String toolName : toolNames) {
				// Skip the tool if it is already present in the request toolCallbacks.
				// That might happen if a tool is defined in the options
				// both as a ToolCallback and as a tool name.
				if (regularTools.stream().anyMatch(tool -> tool.getToolDefinition().name().equals(toolName))) {
					continue;
				}
				
				if (this.resolver == null) {
					throw new IllegalStateException(
							"ToolCallbackResolver is null; cannot resolve tool name: " + toolName);
				}
				ToolCallback toolCallback = this.resolver.resolve(toolName);
				if (toolCallback == null) {
					logger.warn(POSSIBLE_LLM_TOOL_NAME_CHANGE_WARNING, toolName);
					throw new IllegalStateException("No ToolCallback found for tool name: " + toolName);
				}
				regularTools.add(toolCallback);
			}
		}
		
		// If regularTools is empty and resolver is provided, try to extract tools from resolver
		if (regularTools.isEmpty() && this.resolver != null) {
			// Check if resolver also implements ToolCallbackProvider
			if (this.resolver instanceof ToolCallbackProvider provider) {
				ToolCallback[] resolverTools = provider.getToolCallbacks();
				if (resolverTools != null && resolverTools.length > 0) {
					regularTools.addAll(List.of(resolverTools));
					if (logger.isDebugEnabled()) {
						logger.debug("Extracted {} tools from ToolCallbackResolver (ToolCallbackProvider)",
								resolverTools.length);
					}
				}
			} else {
				// This is a fallback for resolvers that don't implement ToolCallbackProvider
				try {
					Field toolsField = this.resolver.getClass().getDeclaredField("tools");
					toolsField.setAccessible(true);
					Object toolsObj = toolsField.get(this.resolver);
					if (toolsObj instanceof java.util.Map) {
						@SuppressWarnings("unchecked") java.util.Map<String, ToolCallback> toolsMap =
								(java.util.Map<String, ToolCallback>) toolsObj;
						if (!toolsMap.isEmpty()) {
							regularTools.addAll(toolsMap.values());
							if (logger.isDebugEnabled()) {
								logger.debug("Extracted {} tools from ToolCallbackResolver via reflection",
										toolsMap.size());
							}
						}
					}
				} catch (NoSuchFieldException | IllegalAccessException | ClassCastException e) {
					// Reflection failed, resolver doesn't have accessible tools field
					// This is expected for some resolver implementations
					if (logger.isTraceEnabled()) {
						logger.trace("Could not extract tools from resolver via reflection: {}", e.getMessage());
					}
				}
			}
		}
		
		// Extract interceptor tools
		List<ToolCallback> interceptorTools = new ArrayList<>();
		if (CollectionUtils.isNotEmpty(modelInterceptors)) {
			interceptorTools = modelInterceptors.stream().flatMap(interceptor -> interceptor.getTools().stream())
					.toList();
		}

		// Extract tools from hooks
		List<ToolCallback> hookTools = new ArrayList<>();
		if (CollectionUtils.isNotEmpty(hooks)) {
			for (Hook hook : hooks) {
				List<ToolCallback> toolsFromHook = hook.getTools();
				if (CollectionUtils.isNotEmpty(toolsFromHook)) {
					hookTools.addAll(toolsFromHook);
					if (logger.isDebugEnabled()) {
						logger.debug("Collected {} tools from hook '{}'", toolsFromHook.size(), hook.getName());
					}
				}
			}
		}

		// Combine all tools: hookTools + interceptorTools + regularTools
		List<ToolCallback> allTools = new ArrayList<>();
		allTools.addAll(hookTools);
		allTools.addAll(interceptorTools);
		allTools.addAll(regularTools);

		return ToolCallbackUtils.deduplicateByName(allTools);
	}

	/**
	 * Merge source options (from ChatModel or ChatClient) with agent-level chatOptions.
	 * Agent-level options take precedence. Uses ModelOptionsUtils.merge with the source
	 * options type as the target.
	 */
	@SuppressWarnings("unchecked")
	private static ChatOptions mergeSourceOptionsWithAgentOptions(ChatOptions sourceOptions,
			ChatOptions agentOptions) {
		if (sourceOptions == null) {
			return agentOptions;
		}
		if (agentOptions == null) {
			return sourceOptions;
		}
		Class<?> sourceOptionsClass = sourceOptions.getClass();
		if (sourceOptionsClass == agentOptions.getClass()) {
			logger.warn(
					"chatOptions type ({}) should be consistent with the default options type ({}) from ChatModel/ChatClient for proper merging.",
					agentOptions.getClass().getName(), sourceOptionsClass.getName());
		}
//		Class<ChatOptions> targetClass = (Class<ChatOptions>) sourceOptionsClass;
//		ChatOptions agentAsSourceType = ModelOptionsUtils.copyToTarget(agentOptions, ChatOptions.class, targetClass);
		return (ChatOptions) ModelOptionsUtils.merge(agentOptions, sourceOptions, sourceOptionsClass);
	}

	/**
	 * Get the default ChatOptions from a ChatModel via reflection (e.g. {@code getDefaultOptions()}).
	 * @param model the ChatModel instance (may be null)
	 * @return the model's default options, or null if not set or reflection fails
	 */
	private static ChatOptions getChatModelDefaultOptions(ChatModel model) {
		if (model == null) {
			return null;
		}
		try {
			Method m = model.getClass().getMethod("getDefaultOptions");
			Object o = m.invoke(model);
			return o instanceof ChatOptions ? (ChatOptions) o : null;
		}
		catch (ReflectiveOperationException e) {
			if (logger.isDebugEnabled()) {
				logger.debug("Could not read default options from ChatModel via reflection: {}", e.getMessage());
			}
			return null;
		}
	}

	/**
	 * Get the default ChatOptions from a ChatClient via reflection (ChatClient does not expose
	 * them publicly). Reads {@code defaultChatClientRequest.chatOptions} from the client.
	 * The options obtained here are the same as (consistent with) the default options in the
	 * ChatModel backing this ChatClient.
	 * @param chatClient the ChatClient instance
	 * @return the client's default options, or null if not set or reflection fails
	 */
	@SuppressWarnings("unchecked")
	private static ChatOptions getChatClientDefaultOptions(ChatClient chatClient) {
		try {
			Field defaultRequestField = chatClient.getClass().getDeclaredField("defaultChatClientRequest");
			defaultRequestField.setAccessible(true);
			Object defaultChatClientRequest = defaultRequestField.get(chatClient);
			if (defaultChatClientRequest == null) {
				return null;
			}
			Field chatOptionsField = defaultChatClientRequest.getClass().getDeclaredField("chatOptions");
			chatOptionsField.setAccessible(true);
			Object options = chatOptionsField.get(defaultChatClientRequest);
			return options instanceof ChatOptions ? (ChatOptions) options : null;
		}
		catch (NoSuchFieldException | IllegalAccessException e) {
			if (logger.isDebugEnabled()) {
				logger.debug("Could not read default options from ChatClient via reflection: {}", e.getMessage());
			}
			return null;
		}
	}

}
