/* ###
 * 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.app.plugin.core.debug.service.modules;

import java.io.File;
import java.util.*;
import java.util.function.BiPredicate;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import ghidra.app.plugin.core.debug.service.modules.ProgramModuleIndexer.IndexEntry;
import ghidra.debug.api.modules.*;
import ghidra.framework.model.DomainFile;
import ghidra.graph.*;
import ghidra.graph.jung.JungDirectedGraph;
import ghidra.program.model.listing.Program;
import ghidra.trace.model.memory.TraceMemoryRegion;
import ghidra.trace.model.modules.TraceModule;
import ghidra.trace.model.target.path.KeyPath;
import ghidra.util.Msg;

public enum DebuggerStaticMappingProposals {
	;

	protected static String getLastLower(String path) {
		return new File(path).getName().toLowerCase();
	}

	/**
	 * Check if either the program's name, its executable path, or its domain file name contains the
	 * given module name
	 * 
	 * @param program the program whose names to check
	 * @param moduleLowerName the module name to check for in lower case
	 * @return true if matched, false if not
	 */
	protected static boolean namesContain(Program program, String moduleLowerName) {
		DomainFile df = program.getDomainFile();
		if (df == null || df.getProjectLocator() == null) {
			return false;
		}
		String programName = getLastLower(program.getName());
		if (programName.contains(moduleLowerName)) {
			return true;
		}
		String exePath = program.getExecutablePath();
		if (exePath != null) {
			String execName = getLastLower(exePath);
			if (execName.contains(moduleLowerName)) {
				return true;
			}
		}
		String fileName = df.getName().toLowerCase();
		if (fileName.contains(moduleLowerName)) {
			return true;
		}
		return false;
	}

	protected interface ProposalGenerator<F, T, MP extends MapProposal<?, ?, ?>> {
		MP proposeMap(F from, long snap, T to);

		MP proposeBestMap(F from, long snap, Collection<? extends T> tos);

		Map<F, MP> proposeBestMaps(Collection<? extends F> froms, long snap,
				Collection<? extends T> tos);
	}

	protected abstract static class AbstractProposalGenerator //
	<F, T, J, MP extends MapProposal<?, ?, ?>> {
		protected final long snap;

		public AbstractProposalGenerator(long snap) {
			this.snap = snap;
		}

		protected abstract MP proposeMap(F from, T to);

		protected abstract J computeFromJoinKey(F from);

		protected abstract boolean isJoined(J key, T to);

		protected Collection<T> filterJoined(J key, Collection<? extends T> tos) {
			return tos.stream()
					.filter(t -> isJoined(key, t))
					// Need to preserve order here
					.collect(Collectors.toCollection(LinkedHashSet::new));
		}

		protected MP proposeBestMap(F from, Collection<? extends T> tos) {
			double bestScore = -1;
			MP bestMap = null;
			for (T t : tos) {
				MP map = proposeMap(from, t);
				double score = map.computeScore();
				// NOTE: Ties prefer first in candidate collection
				if (score > bestScore) {
					bestScore = score;
					bestMap = map;
				}
			}
			return bestMap;
		}

		protected Map<F, MP> proposeBestMaps(Collection<? extends F> froms,
				Collection<? extends T> tos) {
			Map<F, MP> result = new LinkedHashMap<>();
			for (F f : froms) {
				J joinKey = computeFromJoinKey(f);
				Collection<T> joined = filterJoined(joinKey, tos);
				MP map = proposeBestMap(f, joined);
				if (map != null) {
					result.put(f, map);
				}
			}
			return result;
		}
	}

	protected static class ModuleMapProposalGenerator
			implements ProposalGenerator<TraceModule, Program, ModuleMapProposal> {
		private final ProgramModuleIndexer indexer;

		public ModuleMapProposalGenerator(ProgramModuleIndexer indexer) {
			this.indexer = indexer;
		}

		@Override
		public ModuleMapProposal proposeMap(TraceModule from, long snap, Program to) {
			return new DefaultModuleMapProposal(from, snap, to);
		}

		@Override
		public ModuleMapProposal proposeBestMap(TraceModule from, long snap,
				Collection<? extends Program> tos) {
			Collection<IndexEntry> entries =
				indexer.filter(indexer.getBestEntries(from, snap), tos);
			DomainFile df = indexer.getBestMatch(from, snap, null, entries);
			if (df == null) {
				return null;
			}
			try (PeekOpenedDomainObject peek = new PeekOpenedDomainObject(df)) {
				return proposeMap(from, snap, (Program) peek.object);
			}
		}

		@Override
		public Map<TraceModule, ModuleMapProposal> proposeBestMaps(
				Collection<? extends TraceModule> froms, long snap,
				Collection<? extends Program> tos) {
			Map<TraceModule, ModuleMapProposal> result = new LinkedHashMap<>();
			for (TraceModule f : froms) {
				ModuleMapProposal map = proposeBestMap(f, snap, tos);
				if (map != null) {
					result.put(f, map);
				}
			}
			return result;
		}
	}

	public static class SectionMapProposalGenerator
			extends AbstractProposalGenerator<TraceModule, Program, String, SectionMapProposal> {
		public SectionMapProposalGenerator(long snap) {
			super(snap);
		}

		@Override
		protected SectionMapProposal proposeMap(TraceModule from, Program to) {
			return new DefaultSectionMapProposal(from, snap, to);
		}

		@Override
		protected String computeFromJoinKey(TraceModule from) {
			return getLastLower(from.getName(snap));
		}

		@Override
		protected boolean isJoined(String key, Program to) {
			return namesContain(to, key);
		}
	}

	public static class RegionMapProposalGenerator extends
			AbstractProposalGenerator<Collection<TraceMemoryRegion>, Program, Set<String>, //
					RegionMapProposal> {

		public RegionMapProposalGenerator(long snap) {
			super(snap);
		}

		@Override
		protected RegionMapProposal proposeMap(Collection<TraceMemoryRegion> from,
				Program to) {
			return new DefaultRegionMapProposal(from, snap, to);
		}

		@Override
		protected Set<String> computeFromJoinKey(Collection<TraceMemoryRegion> from) {
			return from.stream()
					.flatMap(r -> getLikelyModulesFromName(r).stream())
					.map(n -> getLastLower(n))
					.collect(Collectors.toSet());
		}

		@Override
		protected boolean isJoined(Set<String> key, Program to) {
			return key.stream().anyMatch(n -> namesContain(to, n));
		}
	}

	public static RegionMapProposal proposeRegionMap(
			Collection<? extends TraceMemoryRegion> regions, long snap,
			Collection<? extends Program> programs) {
		return new RegionMapProposalGenerator(snap)
				.proposeBestMap(Collections.unmodifiableCollection(regions), programs);
	}

	public static <V, J> Set<Set<V>> groupByComponents(Collection<? extends V> vertices,
			Function<V, J> precompute, BiPredicate<J, J> areConnected) {
		List<V> vs = List.copyOf(vertices);
		List<J> pres = vs.stream().map(precompute).collect(Collectors.toList());
		GDirectedGraph<V, GEdge<V>> graph = new JungDirectedGraph<>();
		for (V v : vs) {
			graph.addVertex(v); // Lone regions should still be considered
		}
		int size = vs.size();
		for (int i = 0; i < size; i++) {
			V v1 = vs.get(i);
			J j1 = pres.get(i);
			for (int j = i + 1; j < size; j++) {
				V v2 = vs.get(j);
				J j2 = pres.get(j);
				if (areConnected.test(j1, j2)) {
					graph.addEdge(new DefaultGEdge<>(v1, v2));
					graph.addEdge(new DefaultGEdge<>(v2, v1));
				}
			}
		}
		return GraphAlgorithms.getStronglyConnectedComponents(graph);
	}

	protected static Set<String> getLikelyModulesFromName(TraceMemoryRegion region) {
		String key;
		try {
			KeyPath path = KeyPath.parse(region.getPath());
			key = KeyPath.parseIfIndex(path.key());
		}
		catch (IllegalArgumentException e) { // Parse error
			Msg.error(DebuggerStaticMappingProposals.class,
				"Encountered unparsable path: " + region.getPath());
			// Path should always parse in object mode. In legacy mode, snap doesn't matter
			key = region.getName(0); // Not a great fallback, but it'll have to do
		}
		return Stream.of(key.split("\\s+"))
				.filter(n -> n.replaceAll("[0-9A-Fa-f]+", "").length() >= 5)
				.collect(Collectors.toSet());
	}

	public static Set<Set<TraceMemoryRegion>> groupRegionsByLikelyModule(
			Collection<? extends TraceMemoryRegion> regions) {
		return groupByComponents(regions, r -> getLikelyModulesFromName(r), (m1, m2) -> {
			return m1.stream().anyMatch(m2::contains);
		});
	}
}
