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

import io.activej.common.Checks;
import io.activej.common.StringFormatUtils;
import io.activej.common.builder.AbstractBuilder;
import io.activej.common.time.CurrentTimeProvider;
import io.activej.dns.protocol.DnsProtocol.ResponseErrorCode;
import io.activej.dns.protocol.DnsQuery;
import io.activej.dns.protocol.DnsQueryException;
import io.activej.dns.protocol.DnsResponse;
import io.activej.jmx.api.attribute.JmxAttribute;
import io.activej.jmx.api.attribute.JmxOperation;
import io.activej.jmx.api.attribute.JmxReducers.JmxReducerSum;
import io.activej.promise.Promise;
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.net.InetAddress;
import java.time.Duration;
import java.util.*;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Predicate;
import java.util.stream.Collectors;

import static io.activej.reactor.Reactive.checkInReactorThread;

/**
 * Represents a cache for storing resolved domains during it's time to live.
 */
public final class DnsCache extends AbstractReactive {
	private static final Logger logger = LoggerFactory.getLogger(DnsCache.class);
	private static final boolean CHECKS = Checks.isEnabled(DnsCache.class);

	public static final Duration DEFAULT_ERROR_CACHE_EXPIRATION = Duration.ofMinutes(1);
	public static final Duration DEFAULT_TIMED_OUT_EXPIRATION = Duration.ofSeconds(1);
	public static final Duration DEFAULT_HARD_EXPIRATION_DELTA = Duration.ofMinutes(1);
	public static final Duration DEFAULT_MAX_TTL = null;

	private final Map<DnsQuery, CachedDnsQueryResult> cache = new ConcurrentHashMap<>();

	private long errorCacheExpiration = DEFAULT_ERROR_CACHE_EXPIRATION.toMillis();
	private long timedOutExpiration = DEFAULT_TIMED_OUT_EXPIRATION.toMillis();
	private long hardExpirationDelta = DEFAULT_HARD_EXPIRATION_DELTA.toMillis();
	private long maxTtl = Long.MAX_VALUE;

	private final AtomicBoolean cleaningUpNow = new AtomicBoolean(false);
	private final PriorityQueue<CachedDnsQueryResult> expirations = new PriorityQueue<>();

	final CurrentTimeProvider now;

	/**
	 * Creates a new DNS cache.
	 *
	 * @param reactor reactor
	 */
	private DnsCache(Reactor reactor) {
		super(reactor);
		this.now = reactor;
	}

	public static DnsCache create(Reactor reactor) {
		return builder(reactor).build();
	}

	public static Builder builder(Reactor reactor) {
		return new DnsCache(reactor).new Builder();
	}

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

		/**
		 * @param errorCacheExpiration expiration time for errors without time to live
		 */
		public Builder withErrorCacheExpiration(Duration errorCacheExpiration) {
			checkNotBuilt(this);
			DnsCache.this.errorCacheExpiration = errorCacheExpiration.toMillis();
			return this;
		}

		/**
		 * @param timedOutExpiration expiration time for timed out exception
		 */
		public Builder withTimedOutExpiration(Duration timedOutExpiration) {
			checkNotBuilt(this);
			DnsCache.this.timedOutExpiration = timedOutExpiration.toMillis();
			return this;
		}

		/**
		 * @param hardExpirationDelta delta between time at which entry is considered resolved, but needs
		 */
		public Builder withHardExpirationDelta(Duration hardExpirationDelta) {
			checkNotBuilt(this);
			DnsCache.this.hardExpirationDelta = hardExpirationDelta.toMillis();
			return this;
		}

		public Builder withMaxTtl(@Nullable Duration maxTtl) {
			checkNotBuilt(this);
			DnsCache.this.maxTtl = maxTtl == null ? Long.MAX_VALUE : maxTtl.toMillis();
			return this;
		}

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

	/**
	 * Tries to get status of the entry for some query from the cache.
	 *
	 * @param query DNS query
	 * @return DnsQueryCacheResult for this query
	 */
	public @Nullable DnsCache.DnsQueryCacheResult tryToResolve(DnsQuery query) {
		CachedDnsQueryResult cachedResult = cache.get(query);

		if (cachedResult == null) {
			logger.trace("{} cache miss", query);
			return null;
		}

		DnsResponse result = cachedResult.response;
		assert result != null; // results with null responses should never be in cache map
		if (result.isSuccessful()) {
			logger.trace("{} cache hit", query);
		} else {
			logger.trace("{} error cache hit", query);
		}

		if (isExpired(cachedResult)) {
			logger.trace("{} hard TTL expired", query);
			return null;
		} else if (isSoftExpired(cachedResult)) {
			logger.trace("{} soft TTL expired", query);
			return new DnsQueryCacheResult(result, true);
		}
		return new DnsQueryCacheResult(result, false);
	}

	private boolean isExpired(CachedDnsQueryResult cachedResult) {
		return now.currentTimeMillis() >= cachedResult.expirationTime + hardExpirationDelta;
	}

	private boolean isSoftExpired(CachedDnsQueryResult cachedResult) {
		return now.currentTimeMillis() >= cachedResult.expirationTime;
	}

	/**
	 * Adds DnsResponse to this cache
	 *
	 * @param response response to add
	 */
	public void add(DnsQuery query, DnsResponse response) {
		if (CHECKS) checkInReactorThread(this);
		long expirationTime = now.currentTimeMillis();
		if (response.isSuccessful()) {
			assert response.getRecord() != null; // where are my advanced contracts so that the IDE would know it's true here without an assertion?
			long minTtl = response.getRecord().getMinTtl() * 1000L;
			if (minTtl == 0) {
				return;
			}
			expirationTime += Math.min(minTtl, maxTtl);
		} else {
			expirationTime += response.getErrorCode() == ResponseErrorCode.TIMED_OUT ?
				timedOutExpiration :
				errorCacheExpiration;
		}
		CachedDnsQueryResult cachedResult = new CachedDnsQueryResult(response, expirationTime);
		CachedDnsQueryResult old = cache.put(query, cachedResult);
		expirations.add(cachedResult);

		if (old != null) {
			old.response = null; // mark old cache response as refreshed (see performCleanup)
			logger.trace("Refreshed cache entry for {}", query);
		} else {
			logger.trace("Added cache entry for {}", query);
		}
	}

	public void performCleanup() {
		if (!cleaningUpNow.compareAndSet(false, true)) {
			return;
		}
		long currentTime = now.currentTimeMillis();

		CachedDnsQueryResult peeked;
		while ((peeked = expirations.peek()) != null && peeked.expirationTime <= currentTime) {
			DnsResponse response = peeked.response;
			if (response != null) { // if it was not refreshed(so there is a newer response in the queue)
				DnsQuery query = response.getTransaction().getQuery();
				cache.remove(query); // we drop it from cache
				logger.trace("Cache entry expired for {}", query);
			}
			expirations.poll();
		}
		cleaningUpNow.set(false);
	}

	@JmxAttribute
	public Duration getErrorCacheExpiration() {
		return Duration.ofMillis(errorCacheExpiration);
	}

	@JmxAttribute
	public void setErrorCacheExpiration(Duration errorCacheExpiration) {
		this.errorCacheExpiration = errorCacheExpiration.toMillis();
	}

	@JmxAttribute
	public Duration getTimedOutExpiration() {
		return Duration.ofMillis(timedOutExpiration);
	}

	@JmxAttribute
	public void setTimedOutExpiration(Duration timedOutExpiration) {
		this.timedOutExpiration = timedOutExpiration.toMillis();
	}

	@JmxAttribute
	public Duration getHardExpirationDelta() {
		return Duration.ofMillis(hardExpirationDelta);
	}

	@JmxAttribute
	public void setHardExpirationDelta(Duration hardExpirationDelta) {
		this.hardExpirationDelta = hardExpirationDelta.toMillis();
	}

	@JmxAttribute
	public String getMaxTtl() {
		return maxTtl == Long.MAX_VALUE ? "" : StringFormatUtils.formatDuration(Duration.ofMillis(maxTtl));
	}

	@JmxAttribute
	public void setMaxTtl(String s) {
		if (s == null || s.isEmpty()) {
			maxTtl = Long.MAX_VALUE;
		} else {
			maxTtl = StringFormatUtils.parseDuration(s).toMillis();
		}
	}

	@JmxAttribute(reducer = JmxReducerSum.class)
	public int getDomainsCount() {
		return cache.size();
	}

	@JmxAttribute(reducer = JmxReducerSum.class)
	public int getFailedDomainsCount() {
		return (int) cache.values().stream()
			.filter(cachedResult -> {
				assert cachedResult.response != null;
				return !cachedResult.response.isSuccessful();
			})
			.count();
	}

	@JmxOperation
	public List<String> getResolvedDomains() {
		return getDomainNames(DnsResponse::isSuccessful);
	}

	@JmxOperation
	public List<String> getFailedDomains() {
		return getDomainNames(response -> !response.isSuccessful());
	}

	private List<String> getDomainNames(Predicate<DnsResponse> predicate) {
		return cache.entrySet()
			.stream()
			.filter(entry -> predicate.test(entry.getValue().response))
			.map(Entry::getKey)
			.map(DnsQuery::getDomainName)
			.collect(Collectors.toList());
	}

	public class RecordFormatter {
		final String domain;
		final CachedDnsQueryResult result;

		RecordFormatter(String domain, CachedDnsQueryResult result) {
			this.domain = domain;
			this.result = result;
		}

		ResponseErrorCode getStatus() {
			if (result.response == null)
				return ResponseErrorCode.UNKNOWN;
			return result.response.getErrorCode();
		}

		Collection<InetAddress> getIps() {
			if (result.response == null || result.response.getRecord() == null)
				return List.of();
			return List.of(result.response.getRecord().getIps());
		}

		int getMinTtlSeconds() {
			if (result.response == null || result.response.getRecord() == null)
				return 0;
			return result.response.getRecord().getMinTtl();
		}

		String getSecondsToSoftExpiration() {
			long secs = (result.expirationTime - now.currentTimeMillis()) / 1000;
			return formatExpired(secs);
		}

		String getSecondsToHardExpiration() {
			long secs = (result.expirationTime + hardExpirationDelta - now.currentTimeMillis()) / 1000;
			return formatExpired(secs);
		}

		private String formatExpired(long secs) {
			return secs < 0 ? "expired" : Long.toString(secs);
		}

		public List<String> formatMultiline() {
			List<String> lines = new ArrayList<>();
			lines.add("DomainName:\t" + domain);
			lines.add("Status:\t" + getStatus());
			lines.add("IP:\t" + getIps());
			lines.add("MinTtlSeconds:\t" + getMinTtlSeconds());
			lines.add("SecondsToSoftExpiration:\t" + getSecondsToSoftExpiration());
			lines.add("SecondsToHardExpiration:\t" + getSecondsToHardExpiration());
			return lines;
		}
	}

	@JmxOperation
	public List<String> getDomainRecord(String domain) {
		return cache.entrySet().stream()
			.filter(e -> e.getKey().getDomainName().equalsIgnoreCase(domain))
			.findFirst()
			.map(e -> new RecordFormatter(e.getKey().getDomainName(), e.getValue()).formatMultiline())
			.orElse(List.of());
	}

	@JmxOperation
	public List<String> getDomainRecords() {
		if (cache.isEmpty())
			return List.of();

		List<String> lines = new ArrayList<>(cache.size());
		lines.add("DomainName;Status;IP;MinTtlSeconds;SecondsToSoftExpiration;SecondsToHardExpiration");
		StringBuilder sb = new StringBuilder();
		cache.forEach((domainName, cachedResult) -> {
			RecordFormatter formatter = new RecordFormatter(domainName.getDomainName(), cachedResult);
			lines.add(sb
				.append(formatter.domain)
				.append(";")
				.append(formatter.getStatus())
				.append(";")
				.append(formatter.getIps())
				.append(";")
				.append(formatter.getMinTtlSeconds())
				.append(";")
				.append(formatter.getSecondsToSoftExpiration())
				.append(";")
				.append(formatter.getSecondsToHardExpiration())
				.toString());
			sb.setLength(0);
		});
		return lines;
	}

	@JmxOperation
	public void clear() {
		cache.clear();
		reactor.submit(expirations::clear);
	}

	public static final class DnsQueryCacheResult {
		private final DnsResponse response;
		private final boolean needsRefreshing;

		private @Nullable DnsQueryException exception;

		public DnsQueryCacheResult(DnsResponse response, boolean needsRefreshing) {
			this.response = response;
			this.needsRefreshing = needsRefreshing;
		}

		public Promise<DnsResponse> getResponseAsPromise() {
			if (response.getErrorCode() == ResponseErrorCode.NO_ERROR) {
				return Promise.of(response);
			}
			if (exception == null) {
				exception = new DnsQueryException(response);
			}
			return Promise.ofException(exception);
		}

		public boolean doesNeedRefreshing() {
			return needsRefreshing;
		}
	}

	public static final class CachedDnsQueryResult implements Comparable<CachedDnsQueryResult> {
		@Nullable DnsResponse response;
		final long expirationTime;

		CachedDnsQueryResult(@Nullable DnsResponse response, long expirationTime) {
			this.response = response;
			this.expirationTime = expirationTime;
		}

		@Override
		public int compareTo(CachedDnsQueryResult o) {
			return Long.compare(expirationTime, o.expirationTime);
		}
	}
}
