use ui;
use core:lang;
use lang:bs:macro;

/**
 * Some data in a program.
 *
 * This can either be a primitive that is a part of some large data structure, a data structure
 * itself, or some combination of the two.
 */
class Data on Render {
	// Type of this data.
	Type? type() : abstract;

	// Update data inside this object. If this object consists of multiple objects that are showed
	// as a single allocation, it should also add them to the traversal.
	void update(World:Traversal t, unsafe:RawPtr object, Nat offset) {}

	// Traverse this data. Most implementations of 'data' will likely also update their
	// representation of their data here as well, as they do not consist of separate
	// allocations. Find and update any pointers to data outside of this data structure.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : abstract;

	// Visit all sub-objects. Note that the target of pointers are *not* visited.
	void visit(DataVisitor v) {
		v.visit(this);
	}

	// Find the sub-data allocated at a particular position. Also attempts to match the desired type
	// of the sub-data if possible.
	Data atOffset(Type? type, Nat offset) { this; }

	// Get the title of this data element. This will be used to create suitable labels when the data
	// is visualized.
	Str title(World world) : abstract;

	// Copy this data element.
	Data copy() : final {
		Data->Data copies;
		copy(copies);
	}

	// Copy this data element, uses 'copies' to keep track of cycles.
	Data copy(Data->Data copies) : final {
		if (copy = copies.at(this)) {
			return copy;
		} else {
			var result = copyObject();
			copies.put(this, result);
			result.copyData(copies);
			result;
		}
	}

	// Create a summary of this object. See the entry function in World for more details.
	void summary(StrBuf to, SummaryData data) : final {
		Nat count = data.id.count;
		Nat id = data.id.get(this, count);
		to << id;
		if (id == count) {
			data.id.put(this, id);

			to << ":{";
			putSummary(to, data);
			to << "}";
		} else {
			// Output a delimiter so that summaries can be concatenated back-to-back elsewhere
			// without issues.
			to << ".";
		}
	}

	// Internal helper for summary. This is called when we have determined that we need to output
	// the contents of this data object.
	protected void putSummary(StrBuf to, SummaryData data) : abstract;

	// Create a copy of this object. Don't copy any data yet.
	protected Data copyObject() : abstract;

	// Copy all data inside.
	protected void copyData(Data->Data copies) : abstract;
}

/**
 * Visitor for all data in an allocation.
 *
 * Visiting a node generally does not traverse pointers. Visiting the allocation can be thought of
 * collecting all data that resides within the box that represents the allocation itself.
 */
class DataVisitor on Render {
	void visit(Data d) : abstract;
}

/**
 * Data passed along to the "summary" functions.
 */
class SummaryData on Render {
	// Create. Provide data for the allocation for each data member.
	init(Data->Allocation alloc, Nat->Nat threadId) {
		init { alloc = alloc; threadId = threadId; }
	}

	// Map from data elements to their allocated ID:s.
	Data->Nat id;

	// Map from data elements to the allocation they are contained in.
	Data->Allocation alloc;

	// Map for thread ID:s. From real thread ID to sequential ID.
	Nat->Nat threadId;
}

/**
 * Unknown data.
 */
class Unknown extends Data {
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {}
	Str title(World world) : override { "Unknown"; }
	Type? type() : override { null; }
	protected void putSummary(StrBuf to, SummaryData data) : override { to << "?"; }
	protected Data copyObject() : override { Unknown(); }
	protected void copyData(Data->Data copies) : override {}
}

/**
 * Pad data.
 *
 * This is essentially a dummy element that can be used as a pointer target in certain
 * circumstances. One example is a pointer past the end in C/C++ arrays.
 */
class Pad extends Data {
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {}
	Str title(World world) : override { "Pad"; }
	Type? type() : override { null; }
	protected void putSummary(StrBuf to, SummaryData data) : override { /* no need, we are padding */ }
	protected Data copyObject() : override { Pad(); }
	protected void copyData(Data->Data copies) : override {}
}

/**
 * Atomic data of some sort.
 *
 * This means that it does not consist of other types. As such, it makes sense to keep track of
 * reads and writes for this data.
 */
class Atom extends Data {
	// Create, give the size of this atom-object.
	init(Nat size) {
		init() {
			size = size;
		}
	}

	// Copy.
	init(Atom original) {
		init() {
			size = original.size;
			access = original.access;
		}
	}

	// Recent memory access to this atom.
	MemAccess access;

	// Size of this object to detect reads/writes properly.
	private Nat size;

	// Update memory access. We do this here rather than in traverse since pointers override
	// traverse in a way that it is difficult to call the super class in some cases.
	void update(World:Traversal t, unsafe:RawPtr object, Nat offset) {
		access = t.memory.query(access, object, offset, size);
	}

	// Add the access to the summary.
	protected void putSummary(StrBuf to, SummaryData data) : override {
		Nat[] writes;
		Nat[] reads;
		Nat[] current = writes;
		for (thread in access.threads) {
			if (thread == 0) {
				current = reads;
			} else {
				current << data.threadId.get(thread);
			}
		}

		writes.sort();
		reads.sort();
		for (x in writes)
			to << "w" << x;
		for (x in reads)
			to << "r" << x;
	}
}

/**
 * Primitive type of some sort (e.g. integer, char, ...)
 */
class Primitive extends Atom {
	// Create.
	init(Type type) {
		init(Value(type).size.current) {
			pType = type;
		}
	}

	// Copy.
	init(Primitive original) {
		init(original) {
			pType = original.pType;
			value = original.value;
		}
	}

	// The type of data in this primitive.
	private Type pType;

	// Value of this primitive, as a string.
	Str value;

	// Traverse.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		if (v = convertPrimitive(pType, object, offset))
			value = v;
	}

	// Title.
	Str title(World world) : override {
		world.typeTitle(pType);
	}

	// Type.
	Type? type() : override { pType; }

	// Summary.
	protected void putSummary(StrBuf to, SummaryData data) : override {
		to << value;
		super:putSummary(to, data);
	}

	// Copy.
	protected Data copyObject() : override {
		Primitive(this);
	}
	protected void copyData(Data->Data copies) : override {
		// Nothing to do here.
	}

}

/**
 * Special case for String. It assumes that it is at offset zero as it is a reference type.
 */
class String extends Primitive {
	// Create.
	init() {
		init(named{core:Str}) {}
	}

	// Traverse.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		if (offset == 0) {
			if (s = object.asObject() as Str) {
				value = s;
				if (value.empty)
					value = " ";
			} else {
				value = "?";
			}
		} else {
			value = "<invalid>";
		}
	}
}


/**
 * Special case for thread ID:s. This visualizes the ThreadId class, which is itself essentially an
 * integer. We need to handle this specially since we need to translate the thread ID inside the
 * summary function. We could also chose to visualize thread ID:s with "pointers" to the thread.
 */
class ThreadId extends Primitive {
	// Create.
	init() {
		init(named{progvis:program:ThreadId}) {
			id = progvis:program:ThreadId:noInit();
		}
	}

	// Copy.
	init(ThreadId original) {
		init(original) {
			id = original.id;
		}
	}

	// Last ThreadID value we read.
	private progvis:program:ThreadId id;

	// Traverse.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		id.v = object.readNat(offset);
		if (id.noInit()) {
			value = "?";
		} else if (id.empty()) {
			value = "–";
		} else {
			value = id.v.toS;
		}
	}

	// Title.
	Str title(World world) : override {
		world.typeTitle(named{Int});
	}

	// Summary.
	protected void putSummary(StrBuf to, SummaryData data) : override {
		Str oldValue = value;
		if (id.any) {
			// Note: It is possible that the thread id refers to a thread that no longer exists.
			value = data.threadId.get(id.v, 0).toS;
		}
		super:putSummary(to, data);
		value = oldValue;
	}

	// Copy.
	protected Data copyObject() : override {
		ThreadId(this);
	}
	protected void copyData(Data->Data copies) : override {
		super:copyData(copies);
		// Nothing to do here.
	}
}

/**
 * Special case for MaybePositive. Shows the special -1 value as "nothing", otherwise works like a
 * regular integer type.
 */
class MaybePositive extends Primitive {
	// Create.
	init() {
		init(named{progvis:program:MaybePositive}) {}
	}

	// Copy.
	init(MaybePositive original) {
		init(original) {}
	}

	// Traverse.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		progvis:program:MaybePositive val(object.readNat(offset));
		if (val.empty()) {
			value = "?";
		} else {
			value = val.v.toS;
		}
	}

	// Title.
	Str title(World world) : override {
		world.typeTitle(named{Int});
	}

	// Copy.
	protected Data copyObject() : override {
		MaybePositive(this);
	}
	protected void copyData(Data->Data copies) : override {
		super:copyData(copies);
		// Nothing to do here.
	}
}

/**
 * Pointer or reference.
 */
class Pointer extends Atom {
	// Create. Provide the type of the pointer if it is known.
	init(Type? type) {
		var sz = core:asm:sPtr;
		if (type)
			sz = Value(type).size;

		init(sz.unaligned.current) {
			pType = type;
		}
	}

	// Copy.
	init(Pointer original) {
		init(original) {
			pType = original.pType;
			to = original.to;
		}
	}

	// Type of the pointer, if known.
	private Type? pType;

	// Object we're pointing to.
	Data? to;

	// Traverse the pointer.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		traverse(t, object.readPtr(offset));
	}

	// Traverse a pointer.
	protected void traverse(World:Traversal t, unsafe:RawPtr ptr) {
		if (ptr.any) {
			to = t.findOrCreate(ptr);
		} else {
			to = null;
		}
	}

	// Title of the pointer.
	Str title(World world) : override {
		if (pType)
			world.typeTitle(pType);
		else
			"";
	}

	// Type.
	Type? type() : override { pType; }

	// Summary.
	protected void putSummary(StrBuf to, SummaryData data) : override {
		if (z = this.to) {
			Data root = z;

			var found = data.alloc.find(z);
			if (found != data.alloc.end)
				root = found.v.data;

			// If the root is not the same object as the one we're pointing to, output the root first.
			if (root !is z) {
				// If the root is already present in the output, we don't need to output it
				// again. It is enough for the second part to output a reference to the exact
				// object. Note: this will not be ambiguous since we only ever omit the first part
				// of the representation.
				if (!data.id.has(root)) {
					root.summary(to, data);
				}
			}

			// Then put the element. Either puts the entire thing, or just a reference to a sub-object.
			z.summary(to, data);
		} else {
			to << "-";
		}

		super:putSummary(to, data);
	}

	// Copy.
	protected Data copyObject() : override {
		Pointer(this);
	}
	protected void copyData(Data->Data copies) : override {
		if (to)
			this.to = to.copy(copies);
	}
}

/**
 * Pointer on the stack. Examines the pointer passed directly to the object rather than
 * dereferencing it once first.
 */
class StackPointer extends Pointer {
	// Note: The type of stack pointers is generally not known exactly, and it does not matter for our usage.
	init() {
		init(null) {}
	}
	init(StackPointer original) {
		init(original) {}
	}

	// Traverse.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		traverse(t, object);
		// TODO: How to update access here? I don't think we are able to.
	}

	// Copy.
	protected Data copyObject() : override {
		StackPointer(this);
	}
}

/**
 * Data that represents an object (i.e. variables with values).
 */
class Composite extends Data {
	/**
	 * An individual variable in the composite type.
	 */
	value Part {
		Str name;
		Nat offset;
		Data value;

		// First element in a particular logical structure (used to show inheritance).
		Bool first;

		init(Str name, Nat offset, Data value, Bool first) {
			init {
				name = name;
				offset = offset;
				value = value;
				first = first;
			}
		}
	}

	// Parts in here.
	Part[] parts;

	// The type of this composite object.
	private Type? myType;

	// Create.
	init(Type? type) {
		init {
			myType = type;
		}
	}

	// Copy.
	init(Composite original) {
		init {
			myType = original.myType;
			parts = Part[](original.parts);
		}
	}

	// Add an element.
	void add(Str name, Nat offset, Data value, Bool first) {
		parts << Part(name, offset, value, first);
	}

	// Update.
	void update(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		for (v in parts)
			v.value.update(t, object, offset + v.offset);
	}

	// Traverse.
	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		for (v in parts)
			v.value.traverse(t, object, offset + v.offset);
	}

	// Title.
	Str title(World world) : override {
		if (myType)
			world.typeTitle(myType);
		else
			"?";
	}

	// Type.
	Type? type() : override {
		myType;
	}

	// Find offset.
	Data atOffset(Type? type, Nat offset) {
		if (parts.empty)
			return this;

		// If we are the type that is sought after, give a pointer to us.
		// Note: This assumes that we are not an array.
		// Note: This does not work for "one past end" pointers.
		if (type) {
			if (myType) {
				if (core:lang:isA(myType, type)) {
					return this;
				}
			}
		}

		Nat found = parts.count - 1;
		for (Nat i = 1; i < parts.count; i++) {
			if (parts[i].offset > offset) {
				found = i - 1;
				break;
			}
		}

		var elem = parts[found];
		return elem.value.atOffset(type, offset - elem.offset);
	}

	// Visit.
	void visit(DataVisitor v) : override {
		super:visit(v);
		for (p in parts)
			p.value.visit(v);
	}

	// Summary.
	protected void putSummary(StrBuf to, SummaryData data) : override {
		for (p in parts) {
			to << p.name << ":";
			p.value.summary(to, data);
		}
	}

	// Copy.
	protected Data copyObject() : override {
		Composite(this);
	}
	protected void copyData(Data->Data copies) : override {
		for (Nat i = 0; i < parts.count; i++)
			parts[i].value = parts[i].value.copy(copies);
	}
}

/**
 * Array data. Similar to objects (keys + values), but may be shown differently, and all elements
 * are of the same type.
 */
class Array extends Data {
	// Create.
	init(Type elemType) {
		init() {
			elemType = elemType;
			elemValue = Value(elemType).isValue;
		}
	}

	init(Type elemType, Bool isValue) {
		init() {
			elemType = elemType;
			elemValue = isValue;
		}
	}

	// Copy.
	init(Array original) {
		init() {
			elemType = original.elemType;
			elemValue = original.elemValue;
			contents = Data[](original.contents);
		}
	}

	// Type of elements in the array.
	Type elemType;

	// Is the type in the array a value type?
	private Bool elemValue;

	// Data elements.
	Data[] contents;

	// Helper function to update the array.
	protected void updateArray(World:Traversal t, unsafe:RawPtr array) {
		// Note: readFilled is safe to call even if array is null.
		updateArray(t, array, array.readFilled);
	}

	// With explicit size specified.
	protected void updateArray(World:Traversal t, unsafe:RawPtr array, Nat size) {
		if (array.empty) {
			contents.clear();
			return;
		}

		Nat stride = array.readSize;
		TypeInfo type(elemType, elemValue, false);

		// Note: We assume we don't need to change the type of elements here, as all arrays are of a
		// fixed type.
		for (Nat i = 0; i < size; i++) {
			if (i >= contents.count)
				contents << t.create(type);

			contents[i].update(t, array, i*stride);
		}
		while (contents.count > size)
			contents.pop();
	}

	// Helper function to traverse the array.
	protected void traverseArray(World:Traversal t, unsafe:RawPtr array) {
		if (array.empty)
			return;

		Nat stride = array.readSize;

		for (i, elem in contents)
			elem.traverse(t, array, i*stride);
	}

	// Visit.
	void visit(DataVisitor v) : override {
		super:visit(v);
		for (c in contents)
			c.visit(v);
	}

	// Summary.
	protected void putSummary(StrBuf to, SummaryData data) : override {
		for (c in contents)
			c.summary(to, data);
	}

	// Copy.
	protected void copyData(Data->Data copies) : override {
		for (Nat i = 0; i < contents.count; i++)
			contents[i] = contents[i].copy(copies);
	}
}

/**
 * Standard array, extracts data from the core:Array implementation.
 */
class StdArray extends Array {
	private Type arrayType;

	init(Type arrayType, Type elemType) {
		init(elemType) {
			arrayType = arrayType;
		}
	}

	init(StdArray original) {
		init(original) {
			arrayType = original.arrayType;
		}
	}

	void update(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		updateArray(t, arrayData(object, offset));
	}

	void traverse(World:Traversal t, unsafe:RawPtr object, Nat offset) : override {
		traverseArray(t, arrayData(object, offset));
	}

	private unsafe:RawPtr arrayData(unsafe:RawPtr object, Nat offset) {
		if (offset != 0) {
			print("WARNING: The supplied object does not seem to be a Storm array.");
			return unsafe:RawPtr();
		}

		// TODO: Extract the actual offset! We want to get the offset of the 'data' member, but as
		// it is of type GcType<void>, it has no type in the Storm type system, and is therefore not
		// visible. We could expose it as a UNKNOWN(PTR_GC) though.
		Nat dataOffset = object.readSize - core:asm:sPtr.current;
		object.readPtr(dataOffset);
	}

	Str title(World world) : override {
		world.typeTitle(elemType) + "[]";
	}

	Type? type() : override {
		arrayType;
	}

	// Copy.
	protected Data copyObject() : override {
		StdArray(this);
	}
	protected void copyData(Data->Data copies) : override {
		super:copyData(copies);
	}
}