/* ###
 * 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.util.database.spatial;

import java.lang.reflect.Array;
import java.util.*;
import java.util.Map.Entry;
import java.util.function.Consumer;

import org.apache.commons.collections4.IteratorUtils;

import ghidra.util.LockHold;
import ghidra.util.database.DBSynchronizedIterator;
import ghidra.util.database.spatial.DBTreeDataRecord.RecordEntry;

public abstract class AbstractConstraintsTreeSpatialMap<
	DS extends BoundedShape<NS>,
	DR extends DBTreeDataRecord<DS, NS, T>,
	NS extends BoundingShape<NS>,
	T,
	Q extends Query<DS, NS>>
		implements SpatialMap<DS, T, Q> {

	protected final AbstractConstraintsTree<DS, DR, NS, ?, T, Q> tree;
	protected final Q query;

	public AbstractConstraintsTreeSpatialMap(AbstractConstraintsTree<DS, DR, NS, ?, T, Q> tree,
			Q query) {
		this.tree = tree;
		this.query = query;
	}

	@Override
	public T put(DS shape, T value) {
		try (LockHold hold = LockHold.lock(tree.dataStore.writeLock())) {
			DBTreeDataRecord<DS, NS, T> record = tree.doInsertData(shape, value);
			return record.getRecordValue();
		}
	}

	@Override
	public boolean remove(DS shape, T value) {
		try (LockHold hold = LockHold.lock(tree.dataStore.writeLock())) {
			return tree.doRemoveData(shape, value, query);
		}
	}

	@Override
	public boolean remove(Entry<DS, T> entry) {
		try (LockHold hold = LockHold.lock(tree.dataStore.writeLock())) {
			if (!(entry instanceof RecordEntry)) {
				return tree.doRemoveData(entry.getKey(), entry.getValue(), query);
			}
			@SuppressWarnings("rawtypes")
			DBTreeDataRecord rec = ((RecordEntry) entry).asRecord();
			// TODO: Should probably just if rec's store is tree.dataStore
			if (!tree.dataStore.asMap().containsValue(rec)) {
				return tree.doRemoveData(entry.getKey(), entry.getValue(), query);
			}
			@SuppressWarnings("unchecked")
			DR dr = (DR) rec;
			tree.doDeleteEntry(dr);
			return true;
		}
	}

	@Override
	public int size() {
		try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
			return tree.count(query);
		}
	}

	@Override
	public boolean isEmpty() {
		try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
			return tree.isEmpty(query);
		}
	}

	protected static int checkInt(long l) {
		if (l > Integer.MAX_VALUE) {
			throw new IllegalArgumentException();
		}
		return (int) l;
	}

	protected abstract static class ToArrayConsumer<A, T, U extends A> implements Consumer<T> {
		protected final A[] arr;
		protected int i = 0;

		public ToArrayConsumer(A[] arr) {
			this.arr = arr;
		}

		@Override
		public void accept(T t) {
			arr[i++] = transformed(t);
		}

		protected abstract U transformed(T t);

		protected void putNull() {
			if (i < arr.length) {
				arr[i] = null;
			}
		}
	}

	protected abstract static class ToListConsumer<A, T, U extends A> implements Consumer<T> {
		protected final List<A> list;

		public ToListConsumer(List<A> list) {
			this.list = list;
		}

		@Override
		public void accept(T t) {
			list.add(transformed(t));
		}

		protected abstract U transformed(T t);
	}

	@Override
	public Collection<Entry<DS, T>> entries() {
		return new AbstractCollection<>() {
			@Override
			public Iterator<Entry<DS, T>> iterator() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					return new DBSynchronizedIterator<>(
						IteratorUtils.transformedIterator(tree.iterator(query), r -> r.asEntry()),
						tree.dataStore.getLock());
				}
			}

			public Object[] toArray() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					/**
					 * One one hand, traversing twice (one to compute the size, and again to
					 * retrieve the elements) is not great. On the other, resizing the array list
					 * many times is also not great. Which is worse? IDK. LATER: Take some
					 * measurements with large traces.
					 */
					int size = AbstractConstraintsTreeSpatialMap.this.size();
					Object[] a = new Object[size];
					ToArrayConsumer<Object, DR, Object> consumer = new ToArrayConsumer<>(a) {
						protected Object transformed(DR t) {
							return t.asEntry();
						};
					};
					tree.visitAllData(query, consumer, false);
					return a;
				}
			}

			@SuppressWarnings("unchecked")
			public <U> U[] toArray(U[] a) {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					int size = AbstractConstraintsTreeSpatialMap.this.size();
					if (a.length < size) {
						a = (U[]) Array.newInstance(a.getClass().getComponentType(), size);
					}
					ToArrayConsumer<U, DR, U> consumer = new ToArrayConsumer<>(a) {
						@Override
						protected U transformed(DR t) {
							return (U) t.asEntry();
						}
					};
					tree.visitAllData(query, consumer, false);
					consumer.putNull();
					return a;
				}
			}

			@Override
			public int size() {
				return checkInt(AbstractConstraintsTreeSpatialMap.this.size());
			}

			@Override
			public boolean isEmpty() {
				return AbstractConstraintsTreeSpatialMap.this.isEmpty();
			}

			@Override
			public Spliterator<Entry<DS, T>> spliterator() {
				// Size estimate is more expensive than benefits of knowing size
				return Spliterators.spliteratorUnknownSize(iterator(), 0);
			}
		};
	}

	@Override
	public Collection<Entry<DS, T>> orderedEntries() {
		return new AbstractCollection<>() {
			@Override
			public Iterator<Entry<DS, T>> iterator() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					return new DBSynchronizedIterator<>(
						IteratorUtils.transformedIterator(tree.orderedIterator(query),
							r -> r.asEntry()),
						tree.dataStore.getLock());
				}
			}

			@Override
			public int size() {
				return checkInt(AbstractConstraintsTreeSpatialMap.this.size());
			}

			@Override
			public boolean isEmpty() {
				return AbstractConstraintsTreeSpatialMap.this.isEmpty();
			}

			@Override
			public Spliterator<Entry<DS, T>> spliterator() {
				// Size estimate is more expensive than benefits of knowing size
				return Spliterators.spliteratorUnknownSize(iterator(), 0);
			}
		};
	}

	@Override
	public Collection<DS> keys() {
		return new AbstractCollection<>() {
			@Override
			public Iterator<DS> iterator() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					return new DBSynchronizedIterator<>(
						IteratorUtils.transformedIterator(tree.iterator(query), r -> r.getShape()),
						tree.dataStore.getLock());
				}
			}

			public Object[] toArray() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					List<DS> result = new ArrayList<>();
					tree.visitAllData(query, new ToListConsumer<>(result) {
						@Override
						protected DS transformed(DR t) {
							return t.getShape();
						}
					}, false);
					return result.toArray();
				}
			}

			@SuppressWarnings("unchecked")
			public <U> U[] toArray(U[] a) {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					int size = AbstractConstraintsTreeSpatialMap.this.size();
					if (a.length < size) {
						a = (U[]) Array.newInstance(a.getClass().getComponentType(), size);
					}
					ToArrayConsumer<U, DR, U> consumer = new ToArrayConsumer<>(a) {
						@Override
						protected U transformed(DR t) {
							return (U) t.getShape();
						}
					};
					tree.visitAllData(query, consumer, false);
					consumer.putNull();
					return a;
				}
			}

			@Override
			public int size() {
				return checkInt(AbstractConstraintsTreeSpatialMap.this.size());
			}

			@Override
			public boolean isEmpty() {
				return AbstractConstraintsTreeSpatialMap.this.isEmpty();
			}

			@Override
			public Spliterator<DS> spliterator() {
				// Size estimate is more expensive than benefits of knowing size
				return Spliterators.spliteratorUnknownSize(iterator(), 0);
			}
		};
	}

	@Override
	public Collection<DS> orderedKeys() {
		return new AbstractCollection<>() {
			@Override
			public Iterator<DS> iterator() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					return new DBSynchronizedIterator<>(
						IteratorUtils.transformedIterator(tree.orderedIterator(query),
							r -> r.getShape()),
						tree.dataStore.getLock());
				}
			}

			@Override
			public int size() {
				return checkInt(AbstractConstraintsTreeSpatialMap.this.size());
			}

			@Override
			public boolean isEmpty() {
				return AbstractConstraintsTreeSpatialMap.this.isEmpty();
			}

			@Override
			public Spliterator<DS> spliterator() {
				// Size estimate is more expensive than benefits of knowing size
				return Spliterators.spliteratorUnknownSize(iterator(), 0);
			}
		};
	}

	@Override
	public Collection<T> values() {
		return new AbstractCollection<>() {
			@Override
			public Iterator<T> iterator() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					return new DBSynchronizedIterator<>(
						IteratorUtils.transformedIterator(tree.iterator(query),
							r -> r.getRecordValue()),
						tree.dataStore.getLock());
				}
			}

			public Object[] toArray() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					List<T> result = new ArrayList<>();
					tree.visitAllData(query, new ToListConsumer<>(result) {
						@Override
						protected T transformed(DR t) {
							return t.getRecordValue();
						}
					}, false);
					return result.toArray();
				}
			}

			@SuppressWarnings("unchecked")
			public <U> U[] toArray(U[] a) {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					int size = AbstractConstraintsTreeSpatialMap.this.size();
					if (a.length < size) {
						a = (U[]) Array.newInstance(a.getClass().getComponentType(), size);
					}
					ToArrayConsumer<U, DR, U> consumer = new ToArrayConsumer<>(a) {
						@Override
						protected U transformed(DR t) {
							return (U) t.getRecordValue();
						}
					};
					tree.visitAllData(query, consumer, false);
					consumer.putNull();
					return a;
				}
			}

			@Override
			public int size() {
				return checkInt(AbstractConstraintsTreeSpatialMap.this.size());
			}

			@Override
			public boolean isEmpty() {
				return AbstractConstraintsTreeSpatialMap.this.isEmpty();
			}

			@Override
			public Spliterator<T> spliterator() {
				// Size estimate is more expensive than benefits of knowing size
				return Spliterators.spliteratorUnknownSize(iterator(), 0);
			}
		};
	}

	@Override
	public Collection<T> orderedValues() {
		return new AbstractCollection<>() {
			@Override
			public Iterator<T> iterator() {
				try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
					return new DBSynchronizedIterator<>(
						IteratorUtils.transformedIterator(tree.orderedIterator(query),
							r -> r.getRecordValue()),
						tree.dataStore.getLock());
				}
			}

			@Override
			public int size() {
				return checkInt(AbstractConstraintsTreeSpatialMap.this.size());
			}

			@Override
			public boolean isEmpty() {
				return AbstractConstraintsTreeSpatialMap.this.isEmpty();
			}

			@Override
			public Spliterator<T> spliterator() {
				// Size estimate is more expensive than benefits of knowing size
				return Spliterators.spliteratorUnknownSize(iterator(), 0);
			}
		};
	}

	@Override
	public Entry<DS, T> firstEntry() {
		try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
			DBTreeDataRecord<DS, NS, T> first = tree.first(query);
			return first == null ? null : first.asEntry();
		}
	}

	@Override
	public DS firstKey() {
		try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
			DBTreeDataRecord<DS, NS, T> first = tree.first(query);
			return first == null ? null : first.getShape();
		}
	}

	@Override
	public T firstValue() {
		try (LockHold hold = LockHold.lock(tree.dataStore.readLock())) {
			DBTreeDataRecord<DS, NS, T> first = tree.first(query);
			return first == null ? null : first.getRecordValue();
		}
	}

	@Override
	public void clear() {
		try (LockHold hold = LockHold.lock(tree.dataStore.writeLock())) {
			tree.clear(query);
		}
	}
}
