/*
Reference implementation of Object runtime code.

Feel free to rewrite this in other languages that can export C symbols.
*/

#include <cstdlib>
#include <cstdint>
#include <cstdio>
#include <vector>
#include <atomic>
#include <Object/Object.h>
#include "Schema.hpp"


#define LENGTHOF(array) (sizeof(array) / sizeof((array)[0]))


/** Size of ObjectClass is part of the ABI. */
static_assert(sizeof(ObjectClass) == 256, "ObjectClass size must be 256 bytes");


static std::atomic<uint64_t> aliveCount{0};


static SchemaNode* rootNode_get() {
	static SchemaNode* const rootNode = new SchemaNode;
	return rootNode;
}


struct alignas(64) Object {
	const SchemaNode* schemaNode = rootNode_get();
	std::atomic<const Schema*> schema{NULL};
	/** Packed reference counts.
	Low 32 bits = strong refs, high 32 bits = weak refs.
	*/
	std::atomic<uint64_t> refs{1};
	void* slotsInline[4] = {};
	void** slots = slotsInline;
};


static const Schema* Object_schema_get(const Object* self) {
	const Schema* schema = self->schema.load(std::memory_order_acquire);
	if (schema)
		return schema;
	schema = SchemaNode_schema_get(self->schemaNode);
	const_cast<Object*>(self)->schema.store(schema, std::memory_order_release);
	return schema;
}


Object* Object_create() {
	Object* self = new Object;
	// assert(self);
	aliveCount.fetch_add(1, std::memory_order_relaxed);
	return self;
}


void Object_ref(const Object* self) {
	if (!self)
		return;
	// This check isn't part of the thread-safety guarantee, but it protects against obtaining a reference within a free() function.
	uint64_t refs = self->refs.load();
	if ((refs & 0xFFFFFFFF) == 0)
		return;
	// Increment strong reference count
	const_cast<Object*>(self)->refs.fetch_add(1);
}


void Object_unref(const Object* self) {
	if (!self)
		return;
	uint64_t refs = self->refs.load();
	// Decrement the strong reference count.
	// If this removes the last strong reference, add a weak reference in the same atomic operation.
	// The weak reference keeps the Object shell alive while its free() callbacks run.
	while (true) {
		uint32_t refsStrong = refs & 0xFFFFFFFF;
		if (refsStrong == 0)
			return;
		uint64_t refsNew = refs - 1;
		if (refsStrong == 1)
			refsNew += uint64_t(1) << 32;
		// Retry if another thread changed either reference count.
		if (!const_cast<Object*>(self)->refs.compare_exchange_weak(refs, refsNew))
			continue;
		// Only the thread that removed the last strong reference runs the free() callbacks.
		if (refsStrong != 1)
			return;
		break;
	}
	// Remove all classes from top to bottom
	const ObjectClass* classBottom = NULL;
	for (const SchemaNode* node = self->schemaNode; node; node = node->parent) {
		if (node->delta.type == SchemaDelta::CLASS)
			classBottom = node->delta.class_;
	}
	if (classBottom)
		Object_classes_remove(const_cast<Object*>(self), classBottom);
	// Release the weak reference, allowing the Object to be deleted if no other weak references remain.
	Object_weakUnref(self);
}


uint32_t Object_refs_get(const Object* self) {
	if (!self)
		return 0;
	return self->refs.load() & 0xFFFFFFFF;
}


void Object_weakRef(const Object* self) {
	if (!self)
		return;
	const_cast<Object*>(self)->refs.fetch_add(uint64_t(1) << 32);
}


void Object_weakUnref(const Object* self) {
	if (!self)
		return;
	uint64_t refs = self->refs.load();
	if ((refs >> 32) == 0)
		return;
	// Decrement weak reference count
	refs = const_cast<Object*>(self)->refs.fetch_sub(uint64_t(1) << 32);
	uint32_t refsStrong = refs & 0xFFFFFFFF;
	uint32_t refsWeak = refs >> 32;
	// Free Object shell if this was the last weak ref and strong refs are already gone
	if (refsWeak == 1 && refsStrong == 0) {
		aliveCount.fetch_sub(1, std::memory_order_relaxed);
		if (self->slots != self->slotsInline)
			free(self->slots);
		delete self;
	}
}


uint32_t Object_weakRefs_get(const Object* self) {
	if (!self)
		return 0;
	return self->refs.load() >> 32;
}


bool Object_weakLock(const Object* self) {
	if (!self)
		return false;
	// Atomically increment strong refs only if currently > 0
	uint64_t refs = self->refs.load();
	while ((refs & 0xFFFFFFFF) > 0) {
		if (const_cast<Object*>(self)->refs.compare_exchange_weak(refs, refs + 1))
			return true;
	}
	return false;
}


void Object_classes_push(Object* self, const ObjectClass* class_, void* slot) {
	if (!self || !class_ || !slot)
		return;
	// Fail silently if class already existed
	const Schema* schema = Object_schema_get(self);
	if (schema->slotIndices.find(class_))
		return;
	uint32_t slotIndex = schema->slotIndices.entryCount;
	self->schemaNode = SchemaNode_child_findOrCreate(self->schemaNode, SchemaDelta_classPush(class_));
	self->schema.store(self->schemaNode->schema.load(std::memory_order_acquire), std::memory_order_relaxed);
	if (slotIndex == LENGTHOF(self->slotsInline) && self->slots == self->slotsInline) {
		self->slots = (void**) malloc(2 * slotIndex * sizeof(void*));
		for (uint32_t inlineIndex = 0; inlineIndex < slotIndex; inlineIndex++)
			self->slots[inlineIndex] = self->slotsInline[inlineIndex];
	}
	else if (slotIndex > LENGTHOF(self->slotsInline) && (slotIndex & (slotIndex - 1)) == 0)
		self->slots = (void**) realloc(self->slots, 2 * slotIndex * sizeof(void*));
	self->slots[slotIndex] = slot;
}


// Don't allow inlining into callers when link-time optimization (LTO) is enabled because it overflows the instruction cache.
__attribute__((noinline))
void* Object_slots_get(const Object* self, const ObjectClass* class_) {
	if (!self || !class_)
		return NULL;
	const Schema* schema = Object_schema_get(self);
	const uint32_t* slotIndex = schema->slotIndices.find(class_);
	if (!slotIndex)
		return NULL;
	return self->slots[*slotIndex];
}


void Object_classes_remove(Object* self, const ObjectClass* class_) {
	if (!self)
		return;

	// Fail silently if the object does not have the class
	bool classFound = false;
	for (const SchemaNode* node = self->schemaNode; node; node = node->parent) {
		if (node->delta.type == SchemaDelta::CLASS && node->delta.class_ == class_) {
			classFound = true;
			break;
		}
	}
	if (!classFound)
		return;

	// Remove classes from top down through the requested class.
	for (const SchemaNode* node = self->schemaNode; node; node = node->parent) {
		if (node->delta.type != SchemaDelta::CLASS)
			continue;
		const ObjectClass* currentClass = node->delta.class_;
		if (currentClass->free)
			currentClass->free(self);
		// Set parent class
		self->schemaNode = node->parent;
		self->schema.store(self->schemaNode->schema.load(std::memory_order_acquire), std::memory_order_relaxed);
		// Stop at the requested class.
		if (currentClass == class_)
			return;
	}
}


void Object_methods_push(Object* self, void* dispatcher, void* method) {
	if (!self || !dispatcher || !method)
		return;
	// Find and return existing SchemaNode with the exact delta
	SchemaDelta delta = SchemaDelta_methodPush(dispatcher, method);
	SchemaNode* child = SchemaNode_child_find(self->schemaNode, delta);
	if (child) {
		self->schemaNode = child;
		self->schema.store(child->schema.load(std::memory_order_acquire), std::memory_order_relaxed);
		return;
	}
	// Check if the dispatcher already has a method
	if (SchemaNode_method_find(self->schemaNode, dispatcher)) {
		// Don't allow overriding with a method that was already pushed
		if (SchemaNode_dispatcher_find(self->schemaNode, method))
			return;
	}
	self->schemaNode = SchemaNode_child_findOrCreate(self->schemaNode, delta);
	self->schema.store(self->schemaNode->schema.load(std::memory_order_acquire), std::memory_order_relaxed);
}


__attribute__((noinline))
void* Object_methods_get(const Object* self, void* dispatcher) {
	if (!self || !dispatcher)
		return NULL;
	const Schema* schema = Object_schema_get(self);
	void* const* method = schema->methods.find(dispatcher);
	if (!method)
		return NULL;
	return *method;
}


__attribute__((noinline))
void* Object_supermethods_get(const Object* self, void* method) {
	if (!self || !method)
		return NULL;
	const Schema* schema = Object_schema_get(self);
	void* const* supermethod = schema->supermethods.find(method);
	if (!supermethod)
		return NULL;
	return *supermethod;
}


char* Object_inspect(const Object* self) {
	if (!self)
		return NULL;

	uint64_t refs = self->refs.load();
	uint32_t strongCount = refs & 0xFFFFFFFF;
	uint32_t weakCount = refs >> 32;

	// Collect classes in push order
	std::vector<const ObjectClass*> classes;
	for (const SchemaNode* node = self->schemaNode; node; node = node->parent) {
		if (node->delta.type == SchemaDelta::CLASS)
			classes.push_back(node->delta.class_);
	}

	// Fills `text` with inspect text if non-null.
	// Returns text length, or -1 if failed.
	auto format = [&](char* text, size_t capacity) {
		int textSize = 0;
		auto append = [&](const char* format, auto... args) {
			size_t available = text ? capacity - textSize : 0;
			int size = snprintf(text ? text + textSize : NULL, available, format, args...);
			if (size < 0)
				return false;
			if (text && (size_t) size >= available)
				return false;
			textSize += size;
			return true;
		};
		if (!append("Object(%p)[%u,%u]:", (const void*) self, strongCount, weakCount))
			return -1;
		for (size_t classIndex = classes.size(); classIndex > 0; classIndex--) {
			const ObjectClass* class_ = classes[classIndex - 1];
			if (!append(" %s(%p)", class_->name, Object_slots_get(self, class_)))
				return -1;
		}
		return textSize;
	};

	// Calculate and allocate result string length
	int textSize = format(NULL, 0);
	if (textSize < 0)
		return NULL;
	size_t capacity = (size_t) textSize + 1;
	char* text = (char*) malloc(capacity);
	if (!text)
		return NULL;
	// Fill result string
	if (format(text, capacity) != textSize) {
		free(text);
		return NULL;
	}
	return text;
}


uint64_t Object_aliveCount_get() {
	return aliveCount.load(std::memory_order_relaxed);
}


uint64_t Object_schemaNodes_count_get() {
	return SchemaNode_count_get(rootNode_get());
}