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

import com.alibaba.cloud.ai.graph.DiagramGenerator;

import static com.alibaba.cloud.ai.graph.StateGraph.END;
import static com.alibaba.cloud.ai.graph.StateGraph.START;
import static java.lang.String.format;
import static java.util.Optional.ofNullable;

/**
 * This class represents a MermaidGenerator that extends DiagramGenerator. It generates a
 * flowchart using Mermaid syntax. The flowchart includes various nodes such as start,
 * stop, web_search, retrieve, grade_documents, generate, transform_query, and different
 * conditional states.
 */
public class MermaidGenerator extends DiagramGenerator {

	public static final char SUBGRAPH_PREFIX = '_';

	private String formatNode(String id, Context ctx) {
		if (!ctx.isSubGraph()) {
			return id;
		}

		if (ctx.anySubGraphWithId(id)) {
			return id;
		}

		if (isStart(id) || isEnd(id)) {
			return format("%s%s", id, ctx.title());
		}

		return format("%s_%s", id, ctx.title());
	}

	@Override
	protected void appendHeader(Context ctx) {
		if (ctx.isSubGraph()) {
			ctx.sb()
				.append(format("subgraph %s\n", ctx.title()))
				.append(format("\t%1$s((start)):::%1$s\n", formatNode(START, ctx)))
				.append(format("\t%1$s((stop)):::%1$s\n", formatNode(END, ctx)));
		}
		else {
			ofNullable(ctx.title()).map(title -> ctx.sb().append(format("---\ntitle: %s\n---\n", title)))
				.orElseGet(ctx::sb)
				.append("flowchart TD\n")
				.append(format("\t%s((start))\n", START))
				.append(format("\t%s((stop))\n", END));
		}
	}

	@Override
	protected void appendFooter(Context ctx) {
		if (ctx.isSubGraph()) {
			ctx.sb().append("end\n");
		}
		else {
			ctx.sb()
				.append('\n')
				.append(format("\tclassDef %s fill:black,stroke-width:1px,font-size:xx-small;\n",
						formatNode(START, ctx)))
				.append(format("\tclassDef %s fill:black,stroke-width:1px,font-size:xx-small;\n",
						formatNode(END, ctx)));
		}
	}

	@Override
	protected void declareConditionalStart(Context ctx, String name) {
		ctx.sb().append('\t');
		ctx.sb().append(format("%s{\"check state\"}\n", formatNode(name, ctx)));
	}

	@Override
	protected void declareNode(Context ctx, String name) {
		ctx.sb().append('\t');
		ctx.sb().append(format("%s(\"%s\")\n", formatNode(name, ctx), name));
	}

	@Override
	protected void declareConditionalEdge(Context ctx, int ordinal) {
		ctx.sb().append('\t');
		ctx.sb().append(format("%s{\"check state\"}\n", formatNode(format("condition%d", ordinal), ctx)));
	}

	@Override
	protected void commentLine(Context ctx, boolean yesOrNo) {
		if (yesOrNo)
			ctx.sb().append("\t%%");
	}

	@Override
	protected void call(Context ctx, String from, String to, CallStyle style) {
		ctx.sb().append('\t');

		from = formatNode(from, ctx);
		to = formatNode(to, ctx);
		ctx.sb().append(switch (style) {
			case CONDITIONAL -> format("%1$s:::%1$s -.-> %2$s:::%2$s\n", from, to);
			default -> format("%1$s:::%1$s --> %2$s:::%2$s\n", from, to);
		});
	}

	@Override
	protected void call(Context ctx, String from, String to, String description, CallStyle style) {
		ctx.sb().append('\t');
		from = formatNode(from, ctx);
		to = formatNode(to, ctx);

		ctx.sb().append(switch (style) {
			case CONDITIONAL -> format("%1$s:::%1$s -.->|%2$s| %3$s:::%3$s\n", from, description, to);
			default -> format("%1$s:::%1s -->|%2$s| %3$s:::%3$s\n", from, description, to);
		});
	}

}
