#include <socks/device.h>
#include <socks/printk.h>
#include <socks/bitmap.h>
#include <socks/btree.h>
#include <socks/kext.h>
#include <socks/vm.h>

static struct vm_cache driver_cache = {
	.c_name = "driver",
	.c_obj_size = sizeof(struct driver),
};

static DECLARE_BITMAP(major_numbers, DEV_MAJOR_MAX);
static struct btree drivers;
static spin_lock_t drivers_lock = SPIN_LOCK_INIT;

static struct driver *__system_driver;

BTREE_DEFINE_SIMPLE_INSERT(struct driver, drv_ent, drv_major, put_driver)
BTREE_DEFINE_SIMPLE_GET(struct driver, unsigned int, drv_ent, drv_major, get_driver_by_major)

BTREE_DEFINE_SIMPLE_INSERT(struct device, dev_driverent, dev_minor, put_device)
BTREE_DEFINE_SIMPLE_GET(struct device, unsigned int, dev_driverent, dev_minor, get_device_by_minor)

kern_status_t init_driver_tree(void)
{
	vm_cache_init(&driver_cache);
	bitmap_set(major_numbers, 0);

	__system_driver = driver_create(kernel_kext(), "system");
	driver_register(__system_driver);

	return KERN_OK;
}

struct driver *driver_create(struct kext *self, const char *name)
{
	if (!self || !name) {
		return NULL;
	}

	struct driver *driver = vm_cache_alloc(&driver_cache, 0);
	if (!driver) {
		return NULL;
	}

	memset(driver, 0x0, sizeof *driver);

	driver->drv_owner = self;
	strncpy(driver->drv_name, name, sizeof driver->drv_name - 1);
	driver->drv_name[sizeof driver->drv_name - 1] = 0;

	bitmap_set(driver->drv_minors, DEV_MINOR_INVALID);

	return driver;
}

kern_status_t driver_destroy(struct driver *driver)
{
	if (btree_first(&driver->drv_children)) {
		return KERN_BUSY;
	}

	return KERN_UNIMPLEMENTED;
}

kern_status_t driver_register(struct driver *driver)
{
	if (driver->drv_major != DEV_MAJOR_INVALID) {
		return KERN_INVALID_ARGUMENT;
	}

	unsigned long flags;
	spin_lock_irqsave(&drivers_lock, &flags);

	struct btree_node *cur = btree_first(&drivers);
	while (cur) {
		struct driver *cur_driver = BTREE_CONTAINER(struct driver, drv_ent, cur);
		if (driver->drv_owner == cur_driver->drv_owner && !strcmp(driver->drv_name, cur_driver->drv_name)) {
			spin_unlock_irqrestore(&drivers_lock, flags);
			return KERN_NAME_EXISTS;
		}

		cur = btree_next(cur);
	}

	unsigned long major = bitmap_lowest_clear(major_numbers, DEV_MAJOR_MAX);
	bitmap_set(major_numbers, major);

	driver->drv_major = major;
	put_driver(&drivers, driver);

	spin_unlock_irqrestore(&drivers_lock, flags);
	return KERN_UNIMPLEMENTED;
}

kern_status_t driver_unregister(struct driver *driver)
{
	if (driver->drv_major == DEV_MAJOR_INVALID) {
		return KERN_INVALID_ARGUMENT;
	}

	unsigned long flags;
	spin_lock_irqsave(&drivers_lock, &flags);

	btree_delete(&drivers, &driver->drv_ent);
	bitmap_clear(major_numbers, driver->drv_major);
	driver->drv_major = DEV_MAJOR_INVALID;

	spin_unlock_irqrestore(&drivers_lock, flags);
	return KERN_OK;
}

unsigned int driver_alloc_minor(struct driver *driver)
{
	unsigned int minor = bitmap_lowest_clear(driver->drv_minors, DEV_MINOR_MAX);
	if (minor == BITMAP_NPOS) {
		return DEV_MINOR_INVALID;
	}

	bitmap_set(driver->drv_minors, minor);
	return minor;
}

void driver_free_minor(struct driver *driver, unsigned int minor)
{
	bitmap_clear(driver->drv_minors, minor);
}

struct driver *system_driver(void)
{
	return __system_driver;
}

struct device *driver_get_device(struct driver *driver, unsigned int minor)
{
	return get_device_by_minor(&driver->drv_children, minor);
}

kern_status_t driver_add_device(struct driver *driver, struct device *dev)
{
	if (get_device_by_minor(&driver->drv_children, dev->dev_minor)) {
		return KERN_NAME_EXISTS;
	}

	put_device(&driver->drv_children, dev);

	return KERN_OK;
}