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

import com.alibaba.cloud.ai.studio.admin.builder.generator.model.App;
import com.alibaba.cloud.ai.studio.admin.builder.generator.model.AppModeEnum;
import com.alibaba.cloud.ai.studio.admin.builder.generator.model.agent.Agent;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.dsl.DSLAdapter;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.generator.GraphProjectDescription;
import com.alibaba.cloud.ai.studio.admin.builder.generator.service.generator.ProjectGenerator;
import io.spring.initializr.generator.io.template.MustacheTemplateRenderer;
import io.spring.initializr.generator.io.template.TemplateRenderer;
import io.spring.initializr.generator.project.ProjectDescription;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.*;

/**
 * Agent 项目生成器：将 Agent Schema 转为最小可运行工程（编译 Agent 为 CompiledGraph）
 */
@Component
public class AgentProjectGenerator implements ProjectGenerator {

	private static final String AGENT_BUILDER_TEMPLATE_NAME = "AgentBuilder.java";

	private static final String GRAPH_RUN_TEMPLATE_NAME = "GraphRunController.java";

	private static final String PACKAGE_NAME = "packageName";

	private static final String IMPORT_SECTION = "importSection";

	private static final String AGENT_SECTION = "agentSection";

	private static final String HAS_RESOLVER = "hasResolver";

	private final DSLAdapter dslAdapter;

	private final TemplateRenderer templateRenderer;

	private final AgentTypeProviderRegistry providerRegistry;

	public AgentProjectGenerator(@Qualifier("agentDSLAdapter") DSLAdapter dslAdapter,
			ObjectProvider<MustacheTemplateRenderer> templateRenderer, AgentTypeProviderRegistry providerRegistry) {
		this.dslAdapter = dslAdapter;
		this.templateRenderer = templateRenderer
			.getIfAvailable(() -> new MustacheTemplateRenderer("classpath:/templates"));
		this.providerRegistry = providerRegistry;
	}

	@Override
	public Boolean supportAppMode(AppModeEnum appModeEnum) {
		return Objects.equals(appModeEnum, AppModeEnum.AGENT);
	}

	@Override
	public void generate(GraphProjectDescription projectDescription, Path projectRoot) {
		// 解析 DSL -> Agent 模型
		App app = dslAdapter.importDSL(projectDescription.getDsl());
		Agent root = (Agent) app.getSpec();

		// 渲染构造 Agent 的 Java 代码片段（支持递归/并行）
		RenderContext ctx = new RenderContext();
		CodeSections sections = collectSections(root, ctx);
		String agentSection = sections.getCode()
				+ String.format("%nreturn %s.getAndCompileGraph();%n", sections.getVarName());

		// 模板渲染并写入
		Map<String, Object> agentBuilderModel = new HashMap<>();
		agentBuilderModel.put(PACKAGE_NAME, projectDescription.getPackageName());
		agentBuilderModel.put(IMPORT_SECTION, String.join("\n", sections.getImports()));
		agentBuilderModel.put(AGENT_SECTION, agentSection);
		agentBuilderModel.put(HAS_RESOLVER, sections.isHasResolver());

		Map<String, Object> graphRunModel = Map.of(PACKAGE_NAME, projectDescription.getPackageName());

		renderAndWriteTemplates(List.of(AGENT_BUILDER_TEMPLATE_NAME, GRAPH_RUN_TEMPLATE_NAME),
				List.of(agentBuilderModel, graphRunModel), projectRoot, projectDescription);
	}

	private void renderAndWriteTemplates(List<String> templateNames, List<Map<String, Object>> models, Path projectRoot,
			ProjectDescription projectDescription) {
		Path fileRoot = createDirectory(projectRoot, projectDescription);

		for (int i = 0; i < templateNames.size(); i++) {
			String templateName = templateNames.get(i);
			Map<String, Object> model = models.get(i);
			Path filePath = fileRoot.resolve(templateName);

			try {
				String template = templateRenderer.render(templateName, model);

				// 覆盖写文件（自动创建/替换文件）
				Files.writeString(filePath, template, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
			}
			catch (IOException e) {
				throw new RuntimeException("Error processing template: " + templateName, e);
			}
		}
	}

	private Path createDirectory(Path projectRoot, ProjectDescription projectDescription) {
		StringBuilder pathBuilder = new StringBuilder("src/main/").append(projectDescription.getLanguage().id());
		String packagePath = projectDescription.getPackageName().replace('.', '/');
		pathBuilder.append("/").append(packagePath).append("/graph/");
		try {
			return Files.createDirectories(projectRoot.resolve(pathBuilder.toString()));
		}
		catch (Exception e) {
			throw new RuntimeException("Got error when creating files", e);
		}
	}

	private CodeSections collectSections(Agent agent, RenderContext ctx) {
		// 递归先处理子 agent，收集子 varNames
		List<String> childVars = new ArrayList<>();
		List<CodeSections> childSections = new ArrayList<>();
		if (agent.getSubAgents() != null) {
			for (Agent sub : agent.getSubAgents()) {
				CodeSections cs = collectSections(sub, ctx);
				childSections.add(cs);
				childVars.add(cs.getVarName());
			}
		}

		// 当前节点由 Provider 渲染
		AgentTypeProvider provider = providerRegistry.get(agent.getAgentClass());
		AgentShell shell = AgentShell.of(agent.getAgentClass(), agent.getName(), agent.getDescription(),
				agent.getInstruction(), agent.getInputKeys(), agent.getOutputKey());
		Map<String, Object> handle = agent.getHandle() == null ? java.util.Map.of() : agent.getHandle();
		CodeSections me = provider.render(shell, handle, ctx, childVars);

		// 合并 imports 与 hasResolver，拼接顺序：子在前、父在后
		for (CodeSections cs : childSections) {
			me.getImports().addAll(cs.getImports());
			if (cs.isHasResolver())
				me.resolver(true);
		}
		StringBuilder code = new StringBuilder();
		for (CodeSections cs : childSections)
			code.append(cs.getCode());
		code.append(me.getCode());
		me.code(code.toString());
		return me;
	}

}
