/* ###
 * IP: GHIDRA
 *
 * 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 ghidra.features.base.memsearch.searcher;

import java.util.function.Predicate;

import ghidra.features.base.memsearch.bytesource.AddressableByteSource;
import ghidra.features.base.memsearch.matcher.ByteMatcher;
import ghidra.program.model.address.*;
import ghidra.util.bytesearch.*;
import ghidra.util.datastruct.Accumulator;
import ghidra.util.task.TaskMonitor;

/**
 * Class for searching bytes from a byteSource (memory) using a {@link ByteMatcher}. It handles
 * breaking the search down into a series of searches, handling gaps in the address set and
 * breaking large address ranges down into manageable sizes.
 * <P>
 * It is created with a specific byte source, matcher, address set, and search limit. Clients can
 * then either call the {@link #findAll(Accumulator, TaskMonitor)} method or use it to incrementally
 * search using {@link #findNext(Address, TaskMonitor)}, 
 * {@link #findPrevious(Address, TaskMonitor)}, or {@link #findOnce(Address, boolean, TaskMonitor)}.
 * @param <T> The client object type used to identify matching patterns
 */

public class MemorySearcher<T> {
	private static final int DEFAULT_CHUNK_SIZE = 16 * 1024;
	private static final int OVERLAP_SIZE = 100;
	private AddressableByteSequence preBytes;
	private AddressableByteSequence mainBytes;
	private AddressableByteSequence postBytes;
	private final ByteMatcher<T> matcher;
	private final int chunkSize;

	private Predicate<MemoryMatch<T>> filter = r -> true;
	private final int searchLimit;
	private final AddressSetView searchSet;

	/**
	 * Constructor
	 * @param byteSource the source of the bytes to be searched
	 * @param matcher the matcher that can find matches in a byte sequence
	 * @param addresses the address in the byte source to search
	 * @param searchLimit the max number of hits before stopping
	 */
	public MemorySearcher(AddressableByteSource byteSource, ByteMatcher<T> matcher,
			AddressSetView addresses, int searchLimit) {
		this(byteSource, matcher, addresses, searchLimit, DEFAULT_CHUNK_SIZE);
	}

	/**
	 * Constructor
	 * @param byteSource the source of the bytes to be searched
	 * @param matcher the matcher that can find matches in a byte sequence
	 * @param addresses the address in the byte source to search
	 * @param searchLimit the max number of hits before stopping
	 * @param chunkSize the maximum number of bytes to feed to the matcher at any one time. 
	 */
	public MemorySearcher(AddressableByteSource byteSource, ByteMatcher<T> matcher,
			AddressSetView addresses, int searchLimit, int chunkSize) {
		this.matcher = matcher;
		this.searchSet = addresses;
		this.searchLimit = searchLimit;
		this.chunkSize = chunkSize;

		preBytes = new AddressableByteSequence(byteSource, chunkSize);
		mainBytes = new AddressableByteSequence(byteSource, chunkSize);
		postBytes = new AddressableByteSequence(byteSource, chunkSize);
	}

	/**
	 * Sets any match filters. The filter can be used to exclude matches that don't meet some
	 * criteria that is not captured in the byte matcher such as alignment and code unit type.
	 * @param filter the predicate to use to filter search results
	 */
	public void setMatchFilter(Predicate<MemoryMatch<T>> filter) {
		this.filter = filter;
	}

	/**
	 * Searches all the addresses in this search's {@link AddressSetView} using the byte matcher to
	 * find matches. As each match is found (and passes any filters), the match is given to the 
	 * accumulator. The search continues until either the entire address set has been search or
	 * the search limit has been reached.
	 * @param accumulator the accumulator for found matches
	 * @param monitor the task monitor
	 * @return true if the search completed searching through the entire address set.
	 */
	public boolean findAll(Accumulator<MemoryMatch<T>> accumulator, TaskMonitor monitor) {
		monitor.initialize(searchSet.getNumAddresses(), "Searching...");

		for (AddressRange range : searchSet.getAddressRanges()) {
			if (!findAll(accumulator, range, monitor)) {
				return false;
			}
		}
		return true;
	}

	/**
	 * Searches forwards or backwards starting at the given address until a match is found or
	 * the start or end of the address set is reached. It does not currently wrap the search.
	 * @param start the address to start searching
	 * @param forward if true, search forward, otherwise, search backwards.
	 * @param monitor the task monitor
	 * @return the first match found or null if no match found.
	 */
	public MemoryMatch<T> findOnce(Address start, boolean forward, TaskMonitor monitor) {
		if (forward) {
			return findNext(start, monitor);
		}
		return findPrevious(start, monitor);
	}

	/**
	 * Searches forwards starting at the given address until a match is found or
	 * the end of the address set is reached. It does not currently wrap the search.
	 * @param start the address to start searching
	 * @param monitor the task monitor
	 * @return the first match found or null if no match found.
	 */
	public MemoryMatch<T> findNext(Address start, TaskMonitor monitor) {

		long numAddresses = searchSet.getNumAddresses() - searchSet.getAddressCountBefore(start);
		monitor.initialize(numAddresses, "Searching....");

		for (AddressRange range : searchSet.getAddressRanges(start, true)) {
			range = range.intersectRange(start, range.getMaxAddress());
			MemoryMatch<T> match = findFirst(range, monitor);
			if (match != null) {
				return match;
			}
			if (monitor.isCancelled()) {
				break;
			}
		}
		return null;
	}

	/**
	 * Searches backwards starting at the given address until a match is found or
	 * the beginning of the address set is reached. It does not currently wrap the search.
	 * @param start the address to start searching
	 * @param monitor the task monitor
	 * @return the first match found or null if no match found.
	 */
	public MemoryMatch<T> findPrevious(Address start, TaskMonitor monitor) {

		monitor.initialize(searchSet.getAddressCountBefore(start) + 1, "Searching....");

		for (AddressRange range : searchSet.getAddressRanges(start, false)) {
			MemoryMatch<T> match = findLast(range, start, monitor);
			if (match != null) {
				return match;
			}
			if (monitor.isCancelled()) {
				break;
			}
		}
		return null;
	}

	private MemoryMatch<T> findFirst(AddressRange range, TaskMonitor monitor) {
		preBytes.clear();
		mainBytes.clear();
		postBytes.clear();

		AddressRangeIterator it = new AddressRangeSplitter(range, chunkSize, true);
		AddressRange first = it.next();

		mainBytes.setRange(first);
		while (it.hasNext()) {
			AddressRange next = it.next();
			postBytes.setRange(next);

			MemoryMatch<T> match = findFirst(monitor);
			if (match != null) {
				return match;
			}
			if (monitor.isCancelled()) {
				break;
			}
			rotateBuffers();	// main becomes pre, post becomes main, and pre becomes post
		}
		// last segment, no extra bytes to overlap, so just search the primary buffer
		postBytes.clear();
		return findFirst(monitor);
	}

	private MemoryMatch<T> findLast(AddressRange range, Address start, TaskMonitor monitor) {
		preBytes.clear();
		mainBytes.clear();
		postBytes.clear();

		if (range.contains(start)) {
			Address min = range.getMinAddress();
			Address max = range.getMaxAddress();
			range = new AddressRangeImpl(min, start);
			AddressRange remaining = new AddressRangeImpl(start.next(), max);
			AddressRange extraRange = new AddressRangeSplitter(remaining, chunkSize, true).next();
			postBytes.setRange(extraRange);
		}

		AddressRangeIterator it = new AddressRangeSplitter(range, chunkSize, false);
		mainBytes.setRange(it.next());

		while (it.hasNext()) {
			AddressRange next = it.next();
			preBytes.setRange(next);
			MemoryMatch<T> match = findLast(monitor);
			if (match != null) {
				return match;
			}
			if (monitor.isCancelled()) {
				break;
			}
			rotateBuffersBackwards(); // main becomes post, pre becomes main, post becomes pre
		}
		preBytes.clear();
		return findLast(monitor);
	}

	private MemoryMatch<T> findFirst(TaskMonitor monitor) {

		ExtendedByteSequence sequence =
			new ExtendedByteSequence(mainBytes, preBytes, postBytes, OVERLAP_SIZE);

		for (Match<T> byteMatch : matcher.match(sequence)) {
			Address address = mainBytes.getAddress((int) byteMatch.getStart());
			byte[] bytes = sequence.getBytes((int) byteMatch.getStart(), byteMatch.getLength());
			MemoryMatch<T> match = new MemoryMatch<>(address, bytes, byteMatch.getPattern());
			if (filter.test(match)) {
				return match;
			}
			if (monitor.isCancelled()) {
				break;
			}
		}
		monitor.incrementProgress(mainBytes.getLength());
		return null;
	}

	private MemoryMatch<T> findLast(TaskMonitor monitor) {

		MemoryMatch<T> last = null;

		ExtendedByteSequence searchSequence =
			new ExtendedByteSequence(mainBytes, preBytes, postBytes, OVERLAP_SIZE);

		for (Match<T> byteMatch : matcher.match(searchSequence)) {
			Address address = mainBytes.getAddress((int) byteMatch.getStart());
			byte[] bytes =
				searchSequence.getBytes((int) byteMatch.getStart(), byteMatch.getLength());
			MemoryMatch<T> match = new MemoryMatch<>(address, bytes, byteMatch.getPattern());
			if (filter.test(match)) {
				last = match;
			}
			if (monitor.isCancelled()) {
				return null;
			}
		}
		monitor.incrementProgress(mainBytes.getLength());
		return last;
	}

	private boolean findAll(Accumulator<MemoryMatch<T>> accumulator, AddressRange range,
			TaskMonitor monitor) {
		preBytes.clear();
		mainBytes.clear();
		postBytes.clear();

		AddressRangeIterator it = new AddressRangeSplitter(range, chunkSize, true);
		AddressRange first = it.next();

		mainBytes.setRange(first);
		while (it.hasNext()) {
			AddressRange next = it.next();
			postBytes.setRange(next);
			if (!findAllInBuffers(accumulator, monitor)) {
				return false;
			}
			rotateBuffers();	// main becomes pre, post becomes main, and pre becomes post
		}
		postBytes.clear();
		return findAllInBuffers(accumulator, monitor);
	}

	private void rotateBuffers() {
		AddressableByteSequence tmp = preBytes;
		preBytes = mainBytes;
		mainBytes = postBytes;
		postBytes = tmp;
	}

	private void rotateBuffersBackwards() {
		// moving backwards through memory
		AddressableByteSequence tmp = postBytes;
		postBytes = mainBytes;
		mainBytes = preBytes;
		preBytes = tmp;
	}

	private boolean findAllInBuffers(Accumulator<MemoryMatch<T>> accumulator, TaskMonitor monitor) {

		if (monitor.isCancelled()) {
			return false;
		}

		ExtendedByteSequence searchSequence =
			new ExtendedByteSequence(mainBytes, preBytes, postBytes, OVERLAP_SIZE);

		for (Match<T> match : matcher.match(searchSequence)) {
			Address address = mainBytes.getAddress((int) match.getStart());
			byte[] bytes = searchSequence.getBytes((int) match.getStart(), match.getLength());
			MemoryMatch<T> memMatch = new MemoryMatch<>(address, bytes, match.getPattern());
			if (filter.test(memMatch)) {
				if (accumulator.size() >= searchLimit) {
					return false;
				}
				accumulator.add(memMatch);
			}
			if (monitor.isCancelled()) {
				return false;
			}

		}
		// Reset the monitor message, since clients may change the message (such as the 
		// incremental table loader)
		monitor.setMessage("Searching...");
		monitor.incrementProgress(mainBytes.getLength());
		return true;
	}
}
