/*******************************************************************************
 * Copyright (c) 2017 Association for Decentralized Information Management
 * in Industry THTH ry.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Eclipse Public License v1.0
 * which accompanies this distribution, and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 *
 * Contributors:
 *     Semantum Oy - #7330 - initial API and implementation
 *******************************************************************************/
package org.simantics.utils.datastructures.file;

import static java.nio.file.StandardWatchEventKinds.ENTRY_CREATE;
import static java.nio.file.StandardWatchEventKinds.ENTRY_DELETE;
import static java.nio.file.StandardWatchEventKinds.ENTRY_MODIFY;
import static java.nio.file.StandardWatchEventKinds.OVERFLOW;

import java.io.Closeable;
import java.io.IOException;
import java.nio.file.FileSystems;
import java.nio.file.FileVisitResult;
import java.nio.file.FileVisitor;
import java.nio.file.Files;
import java.nio.file.NoSuchFileException;
import java.nio.file.Path;
import java.nio.file.SimpleFileVisitor;
import java.nio.file.WatchEvent;
import java.nio.file.WatchEvent.Kind;
import java.nio.file.WatchKey;
import java.nio.file.WatchService;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.LongConsumer;

import org.simantics.databoard.util.BijectionMap;
import org.simantics.utils.FileUtils;
import org.simantics.utils.datastructures.MapSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import gnu.trove.map.TObjectLongMap;
import gnu.trove.map.hash.TObjectLongHashMap;

/**
 * @author Tuukka Lehtonen
 * @since 1.30.0
 */
public class DirectorySizeTracker implements Runnable, Closeable {

	private static final Logger LOGGER = LoggerFactory.getLogger(DirectorySizeTracker.class);

	private static final int TRACE_NONE = 0;
	@SuppressWarnings("unused")
	private static final int TRACE_DIRS = 1;
	@SuppressWarnings("unused")
	private static final int TRACE_EVENTS = 2;
	@SuppressWarnings("unused")
	private static final int TRACE_FILES = 3;

	private static final Kind<?>[] ALL_EVENTS = { ENTRY_CREATE, ENTRY_DELETE, ENTRY_MODIFY };

	private static AtomicInteger threadCounter = new AtomicInteger();

	private LongConsumer consumer;
	private Kind<?>[] events;

	private WatchService watcher;

	private Object lock = new Object();

	// For watched directory tracking
	private BijectionMap<WatchKey, Path> trackedDirs = new BijectionMap<>();
	private MapSet<Path, Path> subdirs = new MapSet.Hash<Path, Path>();
	private MapSet<Path, Path> files = new MapSet.Hash<Path, Path>();

	// For directory and file size tracking
	private TObjectLongMap<Path> entrySizes = new TObjectLongHashMap<>(1024, 0.5f, -1L);
	private TObjectLongMap<Path> dirSizes = new TObjectLongHashMap<>(512, 0.5f, -1L);
	private long totalSize = 0L;

	private int traceLevel = TRACE_NONE;
	private boolean running = true;
	private Thread thread;

	public static DirectorySizeTracker startTracker(LongConsumer sizeChangeListener) throws IOException {
		DirectorySizeTracker watcher = new DirectorySizeTracker(sizeChangeListener, ALL_EVENTS);
		watcher.thread = new Thread(watcher, DirectorySizeTracker.class.getSimpleName() + threadCounter.get());
		watcher.thread.start();
		return watcher;
	}

	private DirectorySizeTracker(LongConsumer sizeChangeListener, Kind<?>[] events) throws IOException {
		this.consumer = sizeChangeListener;
		this.events = events;
		this.watcher = FileSystems.getDefault().newWatchService();
	}

	public void close(boolean joinThread) throws InterruptedException {
		running = false;
		thread.interrupt();
		if (joinThread)
			thread.join();
	}

	@Override
	public void close() throws IOException {
		try {
			close(true);
		} catch (InterruptedException e) {
			throw new IOException(e);
		}
	}

	/**
	 * Register the given directory with the WatchService to listen for the
	 * default set of events provided to the constructor.
	 */
	public void track(Path dir) throws IOException {
		synchronized (lock) {
			if (trackedDirs.containsRight(dir))
				return;
			if (traceLevel > 0)
				LOGGER.info("Starting to track entire directory " + dir + " with a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
			Files.walkFileTree(dir, REGISTER);
			if (traceLevel > 0)
				LOGGER.info("Now tracking entire directory " + dir + " with a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
		}
	}

	private FileVisitor<Path> REGISTER = new SimpleFileVisitor<Path>() {
		private Path currentDir;

		public java.nio.file.FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) throws IOException {
			currentDir = dir;
			if (traceLevel > 2)
				LOGGER.info("Set current dir to " + currentDir);
			return FileVisitResult.CONTINUE;
		}

		public java.nio.file.FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
			registerFile(currentDir, file, attrs);
			return FileVisitResult.CONTINUE;
		}

		public java.nio.file.FileVisitResult visitFileFailed(Path file, IOException exc) throws IOException {
			LOGGER.warn("Failed to visit file " + file, exc);
			return FileVisitResult.CONTINUE;
		}

		public java.nio.file.FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
			// Start tracking directories only after registering all entries
			// within them.
			Path parent = dir.getParent();
			if (trackDir(dir))
				subdirs.add(parent, dir);
			currentDir = parent;
			return FileVisitResult.CONTINUE;
		}
	};

	/**
	 * Register the given directory with the WatchService to listen for the
	 * specified set of events.
	 */
	private boolean trackDir(Path dir) throws IOException {
		if (trackedDirs.containsRight(dir))
			return false;

		WatchKey key = dir.register(watcher, events);

		if (traceLevel > 1) {
			Path prev = trackedDirs.getRight(key);
			if (prev == null) {
				LOGGER.info("Tracking new directory {}\n", dir);
			} else {
				if (!dir.equals(prev)) {
					LOGGER.info("Tracked directory update: {} -> {}\n", prev, dir);
				}
			}
		}

		trackedDirs.map(key, dir);
		return true;
	}

	private boolean registerFile(Path dir, Path file, BasicFileAttributes attrs) {
		if (files.add(dir, file)) {
			long entrySize = attrs.size();
			entrySizes.put(file, entrySize);
			dirSizes.adjustOrPutValue(dir, entrySize, entrySize);
			totalSize += entrySize;
			if (traceLevel > 2)
				LOGGER.info("Registered file " + file + " size " + entrySize + " for total size " + totalSize);
			return true;
		} else {
			long size = attrs.size();
			long oldSize = entrySizes.put(file, size);
			long sizeDelta = oldSize >= 0 ? size - oldSize : size;
			totalSize += sizeDelta;
			dirSizes.adjustOrPutValue(dir, sizeDelta, sizeDelta);
			if (traceLevel > 2)
				LOGGER.info("Modified " + file + " size from " + oldSize + " to " + size
						+ " with delta " + sizeDelta + ", dir size = " + dirSizes.get(dir)
						+ ", total size = " + totalSize);
		}
		return false;
	}

	private boolean unregisterFile(Path dir, Path file) {
		long fileSize = entrySizes.remove(file);
		if (fileSize >= 0) {
			totalSize -= fileSize;
			if (files.remove(dir, file))
				dirSizes.adjustValue(dir, -fileSize);
			if (traceLevel > 2)
				LOGGER.info("Unregistered file " + file + " of size " + fileSize + ", dirSize = " + dirSizes.get(dir) + ", totalSize = " + totalSize);
			return true;
		}
		return false;
	}

	public void untrack(Path dir) {
		synchronized (lock) {
			if (!trackedDirs.containsRight(dir))
				return;
			if (traceLevel > 0)
				LOGGER.info("Starting to untrack entire directory " + dir + " with total tracked size " + totalSize);
			untrackTree(dir);
			subdirs.remove(dir.getParent(), dir);
			if (traceLevel > 0)
				LOGGER.info("Done untracking entire directory " + dir + " with total tracked size " + totalSize);
		}
	}

	private void untrackTree(Path dir) {
		Set<Path> subdirs = this.subdirs.removeValues(dir);
		for (Path subdir : subdirs)
			untrackTree(subdir);
		untrackDir(dir);
	}

	private void untrackDir(Path dir) {
		if (traceLevel > 1)
			LOGGER.info("Untrack directory " + dir + " with total tracked size " + totalSize);
		WatchKey key = trackedDirs.removeWithRight(dir);
		if (key != null)
			key.cancel();

		Set<Path> registeredFiles = files.removeValues(dir);
		for (Path file : registeredFiles)
			unregisterFile(dir, file);
		dirSizes.remove(dir);
	}

	@SuppressWarnings("unchecked")
	static <T> WatchEvent<T> cast(WatchEvent<?> event) {
		return (WatchEvent<T>) event;
	}

	/**
	 * Process all events for keys queued to the watcher
	 */
	void processEvents() {
		while (running) {
			// wait for key to be signaled
			WatchKey key;
			try {
				key = watcher.take();
			} catch (InterruptedException x) {
				return;
			}

			Path dir = trackedDirs.getRight(key);
			if (dir == null) {
				LOGGER.error("WatchKey not registered: " + key);
				continue;
			}

			synchronized (lock) {
				for (WatchEvent<?> event : key.pollEvents()) {
					WatchEvent.Kind<?> kind = event.kind();

					// TBD - provide example of how OVERFLOW event is handled
					if (kind == OVERFLOW)
						continue;

					// Context for directory entry event is the file name of entry
					WatchEvent<Path> evt = cast(event);
					Path name = evt.context();
					Path child = dir.resolve(name);

					if (traceLevel > 1)
						LOGGER.info(String.format("%s: %s", event.kind().name(), child));

					if (kind == ENTRY_CREATE) {
						try {
							BasicFileAttributes attrs = Files.readAttributes(child, BasicFileAttributes.class);
							if (attrs.isDirectory())
								track(child);
							else if (attrs.isRegularFile()) {
								registerFile(dir, child, attrs);
							}
						} catch (IOException ioe) {
							LOGGER.error("Failed to read attribute for path " + child, ioe);
						}
					} else if (kind == ENTRY_MODIFY) {
						try {
							BasicFileAttributes attrs = Files.readAttributes(child, BasicFileAttributes.class);
							if (attrs.isRegularFile()) {
								registerFile(dir, child, attrs);
							}
						} catch (NoSuchFileException ioe) {
							// It is possible that child is a directory that has been removed.
							// In this case, just untrack the whole tree of child.
							if (!entrySizes.containsKey(child)) {
								untrack(child);
							}
						} catch (IOException ioe) {
							LOGGER.error("Failed to read attribute for path " + child, ioe);
						}
					} else if (kind == ENTRY_DELETE) {
						if (!unregisterFile(dir, child)) {
							// This must have been a directory since it isn't registered as a file.
							untrack(child);
						}
					}
				}

				// reset key and remove from set if directory no longer accessible
				boolean valid = key.reset();
				if (!valid) {
					if (traceLevel > 0)
						LOGGER.info("WatchKey for dir " + dir + " is no longer valid. Untracking it.");
					untrack(dir);
					if (trackedDirs.isEmpty()) {
						// all directories are inaccessible
						break;
					}
				}
			}

			if (traceLevel > 1)
				LOGGER.info("STATUS: Tracking a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
			if (consumer != null)
				consumer.accept(totalSize);
		}
	}

	@Override
	public void run() {
		try {
			while (running)
				processEvents();
		} finally {
			FileUtils.uncheckedClose(watcher);
		}
	}

	/**
	 * @return total size of the tracked directories in bytes
	 */
	public long getTotalSize() {
		return totalSize;
	}

//	public static void main(String[] args) {
//		try {
//			DirectorySizeTracker tracker = DirectorySizeTracker.startTracker(null);
//			tracker.track(Paths.get("d:/track-test"));
//			LOGGER.info("AFTER TRACK: Total size from " + tracker.entrySizes.size() + " files and " + tracker.trackedDirs.size() + " directories is " + tracker.getTotalSize());
//			Thread.sleep(2000);
//			//tracker.untrack(Paths.get("d:/track-test"));
//			//LOGGER.info("AFTER UNTRACK: Total size from " + tracker.entrySizes.size() + " files and " + tracker.trackedDirs.size() + " directories is " + tracker.getTotalSize());
//		} catch (IOException | InterruptedException e) {
//			LOGGER.error("test failed", e);
//		}
//	}

}