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

import io.activej.async.function.AsyncRunnable;
import io.activej.async.function.AsyncRunnables;
import io.activej.async.function.AsyncSupplier;
import io.activej.async.service.ReactiveService;
import io.activej.common.builder.AbstractBuilder;
import io.activej.common.function.ConsumerEx;
import io.activej.fs.IFileSystem;
import io.activej.fs.exception.FileSystemException;
import io.activej.fs.exception.FileSystemIOException;
import io.activej.jmx.api.attribute.JmxAttribute;
import io.activej.promise.Promise;
import io.activej.promise.Promises;
import io.activej.reactor.AbstractReactive;
import io.activej.reactor.Reactor;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static io.activej.async.util.LogUtils.toLogger;
import static io.activej.fs.cluster.ServerSelector.RENDEZVOUS_HASH_SHARDER;
import static io.activej.reactor.Reactive.checkInReactorThread;
import static java.util.stream.Collectors.toList;

public final class FileSystemPartitions extends AbstractReactive
	implements ReactiveService {

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

	static final FileSystemException LOCAL_EXCEPTION = new FileSystemException("Local exception");

	private final IDiscoveryService discoveryService;

	private final Map<Object, IFileSystem> alivePartitions = new HashMap<>();
	private final Map<Object, IFileSystem> alivePartitionsView = Collections.unmodifiableMap(alivePartitions);

	private final Map<Object, IFileSystem> deadPartitions = new HashMap<>();
	private final Map<Object, IFileSystem> deadPartitionsView = Collections.unmodifiableMap(deadPartitions);

	private final AsyncRunnable checkAllPartitions = AsyncRunnables.reuse(this::doCheckAllPartitions);
	private final AsyncRunnable checkDeadPartitions = AsyncRunnables.reuse(this::doCheckDeadPartitions);

	private final Map<Object, IFileSystem> partitions = new HashMap<>();
	private final Map<Object, IFileSystem> partitionsView = Collections.unmodifiableMap(partitions);

	private ServerSelector serverSelector = RENDEZVOUS_HASH_SHARDER;

	private FileSystemPartitions(Reactor reactor, IDiscoveryService discoveryService) {
		super(reactor);
		this.discoveryService = discoveryService;
	}

	public static FileSystemPartitions create(Reactor reactor, IDiscoveryService discoveryService) {
		return builder(reactor, discoveryService).build();
	}

	public static Builder builder(Reactor reactor, IDiscoveryService discoveryService) {
		return new FileSystemPartitions(reactor, discoveryService).new Builder();
	}

	public final class Builder extends AbstractBuilder<Builder, FileSystemPartitions> {
		private Builder() {}

		/**
		 * Sets the server selection strategy based on file name and alive partitions
		 */
		public Builder withServerSelector(ServerSelector serverSelector) {
			checkNotBuilt(this);
			FileSystemPartitions.this.serverSelector = serverSelector;
			return this;
		}

		@Override
		protected FileSystemPartitions doBuild() {
			return FileSystemPartitions.this;
		}
	}

	/**
	 * Returns an unmodifiable view of all partitions
	 */
	public Map<Object, IFileSystem> getPartitions() {
		return partitionsView;
	}

	/**
	 * Returns an unmodifiable view of alive partitions
	 */
	public Map<Object, IFileSystem> getAlivePartitions() {
		return alivePartitionsView;
	}

	/**
	 * Returns an unmodifiable view of dead partitions
	 */
	public Map<Object, IFileSystem> getDeadPartitions() {
		return deadPartitionsView;
	}

	/**
	 * Returns alive {@link IFileSystem} by given id
	 *
	 * @param partitionId id of {@link IFileSystem}
	 * @return alive {@link IFileSystem}
	 */
	public @Nullable IFileSystem get(Object partitionId) {
		return alivePartitions.get(partitionId);
	}

	/**
	 * Starts a check process, which pings all partitions and marks them as dead or alive accordingly
	 *
	 * @return promise of the check
	 */
	public Promise<Void> checkAllPartitions() {
		checkInReactorThread(this);
		return checkAllPartitions.run()
			.whenComplete(toLogger(logger, "checkAllPartitions"));
	}

	/**
	 * Starts a check process, which pings all dead partitions to possibly mark them as alive.
	 * This is the preferred method as it does nothing when no partitions are marked as dead,
	 * and RemoteF operations themselves do mark nodes as dead on connection failures.
	 *
	 * @return promise of the check
	 */
	public Promise<Void> checkDeadPartitions() {
		checkInReactorThread(this);
		return checkDeadPartitions.run()
			.whenComplete(toLogger(logger, "checkDeadPartitions"));
	}

	/**
	 * Mark a partition as dead. It means that no operations will use it, and it would not be given to the server selector.
	 * Next call to {@link #checkDeadPartitions()} or {@link #checkAllPartitions()} will ping this partition and possibly
	 * mark it as alive again.
	 *
	 * @param partitionId id of the partition to be marked
	 * @param e           optional exception for logging
	 * @return <code>true</code> if partition was alive and <code>false</code> otherwise
	 */
	@SuppressWarnings("UnusedReturnValue")
	public boolean markDead(Object partitionId, @Nullable Exception e) {
		checkInReactorThread(this);
		IFileSystem partition = alivePartitions.remove(partitionId);
		if (partition != null) {
			logger.warn("marking {} as dead ", partitionId, e);
			deadPartitions.put(partitionId, partition);
			return true;
		}
		return false;
	}

	public void markAlive(Object partitionId) {
		checkInReactorThread(this);
		IFileSystem partition = deadPartitions.remove(partitionId);
		if (partition != null) {
			logger.info("Partition {} is alive again!", partitionId);
			alivePartitions.put(partitionId, partition);
		}
	}

	/**
	 * If partition has returned exception other than {@link FileSystemException} that indicates that there were connection problems
	 * or that there were no response at all
	 */
	public void markIfDead(Object partitionId, Exception e) {
		checkInReactorThread(this);
		if (!(e instanceof FileSystemException) || e instanceof FileSystemIOException) {
			markDead(partitionId, e);
		}
	}

	public ConsumerEx<Exception> wrapDeathFn(Object partitionId) {
		return e -> {
			checkInReactorThread(this);
			markIfDead(partitionId, e);
			if (e instanceof FileSystemException) {
				throw e;
			}
			logger.warn("Node failed", e);
			throw new FileSystemIOException("Node failed");
		};
	}

	public List<Object> select(String filename) {
		checkInReactorThread(this);
		return serverSelector.selectFrom(filename, alivePartitions.keySet());
	}

	public ServerSelector getServerSelector() {
		return serverSelector;
	}

	@Override
	public Promise<?> start() {
		checkInReactorThread(this);
		AsyncSupplier<Map<Object, IFileSystem>> discoverySupplier = discoveryService.discover();
		return discoverySupplier.get()
			.whenResult(result -> {
				this.partitions.putAll(result);
				this.alivePartitions.putAll(result);
				checkAllPartitions()
					.whenComplete(() -> rediscover(discoverySupplier));
			});
	}

	@Override
	public Promise<?> stop() {
		checkInReactorThread(this);
		return Promise.complete();
	}

	@Override
	public String toString() {
		return "FileSystemPartitions{partitions=" + partitions + ", deadPartitions=" + deadPartitions + '}';
	}

	private void rediscover(AsyncSupplier<Map<Object, IFileSystem>> discoverySupplier) {
		discoverySupplier.get()
			.whenResult(result -> {
				updatePartitions(result);
				checkAllPartitions()
					.whenComplete(() -> rediscover(discoverySupplier));
			})
			.whenException(e -> {
				logger.warn("Could not discover partitions", e);
				reactor.delayBackground(Duration.ofSeconds(1), () -> rediscover(discoverySupplier));
			});
	}

	private void updatePartitions(Map<Object, IFileSystem> newPartitions) {
		this.partitions.clear();
		this.partitions.putAll(newPartitions);

		alivePartitions.keySet().retainAll(this.partitions.keySet());
		deadPartitions.keySet().retainAll(this.partitions.keySet());

		for (Map.Entry<Object, IFileSystem> entry : this.partitions.entrySet()) {
			Object partitionId = entry.getKey();
			IFileSystem fs = entry.getValue();

			IFileSystem deadFS = deadPartitions.get(partitionId);
			if (deadFS != null) {
				if (deadFS == fs) continue;

				deadPartitions.remove(partitionId);
			}
			alivePartitions.put(partitionId, fs);
		}

		alivePartitions.clear();
		deadPartitions.clear();
	}

	private Promise<Void> doCheckAllPartitions() {
		return Promises.all(
			partitions.entrySet().stream()
				.map(entry -> {
					Object id = entry.getKey();
					return entry.getValue()
						.ping()
						.map(($, e) -> {
							if (e == null) {
								markAlive(id);
							} else {
								markDead(id, e);
							}
							return null;
						});
				}));
	}

	private Promise<Void> doCheckDeadPartitions() {
		return Promises.all(
			deadPartitions.entrySet().stream()
				.map(entry -> entry.getValue()
					.ping()
					.map(($, e) -> {
						if (e == null) {
							markAlive(entry.getKey());
						}
						return null;
					})
				));
	}

	@JmxAttribute
	public List<String> getAllPartitions() {
		return partitions.keySet().stream()
			.map(Object::toString)
			.collect(toList());
	}
}
