content: /*
   * Copyright (C) 2020 ActiveJ LLC.
   *
   * 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
   *
   * http://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 io.activej.dataflow.graph;
  
  import com.fasterxml.jackson.core.JsonProcessingException;
  import com.fasterxml.jackson.core.type.TypeReference;
  import com.fasterxml.jackson.core.util.DefaultPrettyPrinter;
  import com.fasterxml.jackson.core.util.Separators;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.fasterxml.jackson.databind.ObjectWriter;
  import com.fasterxml.jackson.databind.SerializationFeature;
  import io.activej.async.process.AsyncCloseable;
  import io.activej.common.collection.Try;
  import io.activej.common.ref.RefInt;
  import io.activej.dataflow.DataflowClient;
  import io.activej.dataflow.DataflowClient.Session;
  import io.activej.dataflow.exception.DataflowException;
  import io.activej.dataflow.node.Node;
  import io.activej.dataflow.node.impl.Download;
  import io.activej.dataflow.node.impl.Upload;
  import io.activej.promise.Promise;
  import io.activej.promise.Promises;
  import io.activej.reactor.AbstractReactive;
  import io.activej.reactor.ImplicitlyReactive;
  import io.activej.reactor.Reactor;
  
  import java.util.*;
  import java.util.concurrent.ThreadLocalRandom;
  
  import static io.activej.reactor.Reactive.checkInReactorThread;
  import static java.util.stream.Collectors.*;
  
  /**
   * Represents a graph of partitions, nodeStats and streams in datagraph system.
   */
  public final class DataflowGraph extends AbstractReactive {
  	private static final ObjectWriter OBJECT_WRITER = new ObjectMapper()
  		.configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false)
  		.writerFor(new TypeReference<List<Node>>() {})
  		.with(new DefaultPrettyPrinter().withSeparators(Separators.createDefaultInstance().withRootSeparator("%n")));
  
  	private final Map<Node, Partition> nodePartitions = new LinkedHashMap<>();
  	private final Map<StreamId, Node> streams = new LinkedHashMap<>();
  
  	private final DataflowClient client;
  	private final List<Partition> availablePartitions;
  
  	public DataflowGraph(Reactor reactor, DataflowClient client, List<Partition> availablePartitions) {
  		super(reactor);
  		this.client = client;
  		this.availablePartitions = availablePartitions;
  	}
  
  	public List<Partition> getAvailablePartitions() {
  		return availablePartitions;
  	}
  
  	public Partition getPartition(Node node) {
  		return nodePartitions.get(node);
  	}
  
  	public Partition getPartition(StreamId streamId) {
  		return getPartition(streams.get(streamId));
  	}
  
  	private Map<Partition, List<Node>> getNodesByPartition() {
  		return nodePartitions.entrySet().stream()
  			.collect(groupingBy(Map.Entry::getValue, mapping(Map.Entry::getKey, toList())));
  	}
  
  	public static class PartitionSession extends ImplicitlyReactive implements AsyncCloseable {
  		private final Partition partition;
  		private final Session session;
  
  		private PartitionSession(Partition partition, Session session) {
  			this.partition = partition;
  			this.session = session;
  		}
  
  		public Promise<Void> execute(long taskId, List<Node> nodes) {
  			return session.execute(taskId, nodes);
  		}
  
  		@Override
  		public void closeEx(Exception e) {
  			session.closeEx(e);
  		}
  	}
  
  	/**
  	 * Executes the defined operations on all partitions.
  	 */
  	public Promise<Void> execute() {
  		checkInReactorThread(this);
  		Map<Partition, List<Node>> nodesByPartition = getNodesByPartition();
  		long taskId = ThreadLocalRandom.current().nextInt() & (Integer.MAX_VALUE >>> 1);
  		return connect(nodesByPartition.keySet())
  			.then(sessions ->
  				Promises.all(
  						sessions.stream()
  							.map(session -> session.execute(taskId, nodesByPartition.get(session.partition))))
  					.whenException(() -> sessions.forEach(PartitionSession::close)));
  	}
  
  	private Promise<List<PartitionSession>> connect(Set<Partition> partitions) {
  		return Promises.toList(partitions.stream()
  				.map(partition -> client.connect(partition.address()).map(session -> new PartitionSession(partition, session)).toTry()))
  			.map(tries -> {
  				List<PartitionSession> sessions = tries.stream()
  					.filter(Try::isSuccess)
  					.map(Try::get)
  					.collect(toList());
  
  				if (sessions.size() != partitions.size()) {
  					sessions.forEach(PartitionSession::close);
  					throw new DataflowException("Cannot connect to all partitions");
  				}
  				return sessions;
  			});
  	}
  
  	public void addNode(Partition partition, Node node) {
  		nodePartitions.put(node, partition);
  		for (StreamId streamId : node.getOutputs()) {
  			streams.put(streamId, node);
  		}
  	}
  
  	public void addNodeStream(Node node, StreamId streamId) {
  		streams.put(streamId, node);
  	}
  
  	public List<Partition> getPartitions(List<? extends StreamId> channels) {
  		List<Partition> partitions = new ArrayList<>();
  		for (StreamId streamId : channels) {
  			Partition partition = getPartition(streamId);
  			partitions.add(partition);
  		}
  		return partitions;
  	}
  
  	public String toGraphViz() {
  		return toGraphViz(false, -1);
  	}
  
  	public String toGraphViz(boolean streamLabels) {
  		return toGraphViz(streamLabels, -1);
  	}
  
  	public String toGraphViz(int maxPartitions) {
  		return toGraphViz(false, maxPartitions);
  	}
  
  	@SuppressWarnings("StringConcatenationInsideStringBufferAppend")
  	public String toGraphViz(boolean streamLabels, int maxPartitions) {
  		StringBuilder sb = new StringBuilder("digraph {\n\n");
  
  		RefInt nodeCounter = new RefInt(0);
  		RefInt clusterCounter = new RefInt(0);
  
  		Map<StreamId, Node> nodesByInput = new HashMap<>();
  		Map<StreamId, StreamId> network = new HashMap<>();
  
  		List<Upload<?>> uploads = new ArrayList<>();
  
  		// collect network streams and populate the nodesByInput lookup map
  		for (Node node : nodePartitions.keySet()) {
  			if (node instanceof Download<?> download) {
  				network.put(download.streamId, download.output);
  			} else if (node instanceof Upload) {
  				uploads.add((Upload<?>) node);
  			} else {
  				node.getInputs().forEach(input -> nodesByInput.put(input, node));
  			}
  		}
  		// check for upload nodeStats not connected to download ones, add them
  		for (Upload<?> upload : uploads) {
  			StreamId streamId = upload.streamId;
  			if (!network.containsKey(streamId)) {
  				nodesByInput.put(streamId, upload);
  			}
  		}
  
  		Map<Node, String> nodeIds = new HashMap<>();
  
  		// define nodeStats and group them by partitions using graphviz clusters
  		getNodesByPartition()
  			.entrySet()
  			.stream()
  			.limit(maxPartitions == -1 ? availablePartitions.size() : maxPartitions)
  			.forEach(e -> {
  				sb.append("  subgraph cluster_")
  					.append(++clusterCounter.value)
  					.append(" {\n")
  					.append("    label=\"")
  					.append(e.getKey().address())
  					.append("\";\n    style=rounded;\n\n");
  				for (Node node : e.getValue()) {
  					// upload and download nodeStats have no common connections
  					// download nodeStats are never drawn, and upload only has an input
  					if (node instanceof Download ||
  						node instanceof Upload &&
  						network.containsKey(((Upload<?>) node).streamId)
  					) {
  						continue;
  					}
  					String nodeId = "n" + ++nodeCounter.value;
  					String name = node.getClass().getSimpleName();
  					sb.append("    ")
  						.append(nodeId)
  						.append(" [label=\"")
  						.append(name)
  						.append("\"];\n");
  					nodeIds.put(node, nodeId);
  				}
  				sb.append("  }\n\n");
  			});
  
  		Set<String> notFound = new HashSet<>();
  
  		// walk over each node outputs and build the connections
  		nodeIds.forEach((node, id) -> {
  			for (StreamId output : node.getOutputs()) {
  				Node outputNode = nodesByInput.get(output);
  				boolean net = false;
  				boolean forceLabel = false;
  				// check if unbound stream is a network one
  				StreamId prev = null;
  				if (outputNode == null) {
  					StreamId through = network.get(output);
  					if (through != null) {
  						prev = output;
  						output = through;
  						// connect to network stream output and
  						// set the 'net' flag to true for dashed arrow
  						outputNode = nodesByInput.get(output);
  						net = outputNode != null;
  					}
  				}
  				String nodeId = nodeIds.get(outputNode);
  				// if still unbound and not net (because of partition limit),
  				// draw as point and force the stream label
  				if (nodeId == null && !net) {
  					nodeId = "s" + output.getId();
  					notFound.add(nodeId);
  					forceLabel = true;
  				}
  				if (nodeId != null) { // nodeId might be null only for net nodeStats here, see previous 'if'
  					sb.append("  ")
  						.append(id)
  						.append(" -> ")
  						.append(nodeId)
  						.append(" [");
  					if (streamLabels || forceLabel) {
  						if (prev != null) {
  							sb.append("taillabel=\"")
  								.append(prev)
  								.append("\", headlabel=\"")
  								.append(output)
  								.append("\"");
  						} else {
  							sb.append("xlabel=\"")
  								.append(output)
  								.append("\"");
  						}
  						if (net) {
  							sb.append(", ");
  						}
  					}
  					if (net) {
  						sb.append("style=dashed");
  					}
  					sb.append("];\n");
  				}
  			}
  		});
  
  		// draw the nodeStats that were never defined as points that still have connections
  		if (!notFound.isEmpty()) {
  			sb.append('\n');
  			notFound.forEach(id -> sb.append("  " + id + " [shape=point];\n"));
  		}
  
  		sb.append("}");
  
  		return sb.toString();
  	}
  
  	@SuppressWarnings("StringConcatenationInsideStringBufferAppend")
  	@Override
  	public String toString() {
  		StringBuilder sb = new StringBuilder();
  		Map<Partition, List<Node>> map = getNodesByPartition();
  		for (Map.Entry<Partition, List<Node>> entry : map.entrySet()) {
  			sb.append("--- " + entry.getKey() + "\n\n");
  			String nodes;
  			try {
  				nodes = OBJECT_WRITER.writeValueAsString(entry.getValue());
  			} catch (JsonProcessingException e) {
  				nodes = "<UNKNOWN>";
  			}
  			sb.append(nodes);
  			sb.append("\n\n");
  		}
  		return sb.toString();
  	}
  }
  