#include <socks/status.h>
#include <socks/object.h>
#include <socks/device.h>
#include <socks/printk.h>
#include <socks/bitmap.h>
#include <socks/libc/stdio.h>

#define DEVICE_CAST(p) OBJECT_C_CAST(struct device, dev_base, &device_type, p)

static struct object *dev_folder = NULL;
static struct device *__root_device = NULL;
static struct device *__misc_device = NULL;
static kern_status_t device_object_destroy(struct object *);
static kern_status_t device_object_read(struct object *obj, void *, size_t *, socks_flags_t);
static kern_status_t device_object_write(struct object *obj, const void *, size_t *, socks_flags_t);
static kern_status_t device_object_query_name(struct object *, char out[OBJECT_NAME_MAX]);
static kern_status_t device_object_get_child_at(struct object *, size_t, struct object **);
static kern_status_t device_object_get_child_named(struct object *, const char *, struct object **);

extern kern_status_t init_driver_tree(void);

extern struct device_type_ops input_type_ops;
extern struct device_type_ops framebuffer_type_ops;
extern struct device_type_ops bus_type_ops;

static struct device_type_ops *type_ops[] = {
	[DEV_TYPE_UNKNOWN] = NULL,
	[DEV_TYPE_BLOCK] = NULL,
	[DEV_TYPE_CHAR] = NULL,
	[DEV_TYPE_NET] = NULL,
	[DEV_TYPE_INPUT] = &input_type_ops,
	[DEV_TYPE_BUS] = &bus_type_ops,
	[DEV_TYPE_FRAMEBUFFER] = &framebuffer_type_ops,
};

static struct object_type device_type = {
	.ob_name = "device",
	.ob_size = sizeof(struct device),
	.ob_ops = {
		.read = device_object_read,
		.write = device_object_write,
		.destroy = device_object_destroy,
		.query_name = device_object_query_name,
		.get_at = device_object_get_child_at,
		.get_named = device_object_get_child_named,
	}
};

static kern_status_t set_root_device(struct device *dev)
{
	if (__root_device) {
		set_remove_object(dev_folder, &__root_device->dev_base);
		object_deref(&__root_device->dev_base);
	}

	object_ref(&dev->dev_base);
	set_add_object(dev_folder, &dev->dev_base);
	__root_device = dev;

	return KERN_OK;
}

kern_status_t device_init(void)
{
	object_type_register(&device_type);
	dev_folder = set_create("dev");
	object_publish(global_namespace(), "/", dev_folder);

	kern_status_t status = init_driver_tree();
	if (status != KERN_OK) {
		return status;
	}

	struct bus_device *system_dev = bus_device_create();
	struct device *system_dev_base = bus_device_base(system_dev);
	snprintf(system_dev_base->dev_name, sizeof system_dev_base->dev_name, "system");
	set_root_device(bus_device_base(system_dev));

	struct bus_device *misc_dev = bus_device_create();
	struct device *misc_dev_base = bus_device_base(misc_dev);
	snprintf(misc_dev_base->dev_name, sizeof misc_dev_base->dev_name, "misc");
	__misc_device = misc_dev_base;

	struct driver *system = system_driver();

	device_register(__root_device, system, NULL);
	device_register(__misc_device, system, __root_device);

	return KERN_OK;
}

struct device *root_device(void)
{
	return __root_device;
}

struct device *misc_device(void)
{
	return __misc_device;
}

struct device *device_alloc(void)
{
	struct object *dev_object = object_create(&device_type);
	if (!dev_object) {
		return NULL;
	}

	return DEVICE_CAST(dev_object);
}

struct device *generic_device_create(void)
{
	struct device *dev = device_alloc();
	if (!dev) {
		return NULL;
	}

	dev->dev_type = DEV_TYPE_UNKNOWN;

	return dev;
}

kern_status_t device_read(struct device *dev, void *buf, size_t size, size_t *bytes_read, socks_flags_t flags)
{
	switch (dev->dev_type) {
	case DEV_TYPE_INPUT:
		return input_device_read(dev, buf, size, bytes_read, flags);
	default:
		return KERN_UNSUPPORTED;
	}
}

kern_status_t device_write(struct device *dev, const void *buf, size_t size, size_t *bytes_written, socks_flags_t flags)
{
	return KERN_UNSUPPORTED;
}

struct device *cast_to_device(struct object *obj)
{
	return DEVICE_CAST(obj);
}

static kern_status_t device_object_read(struct object *obj, void *p, size_t *count, socks_flags_t flags)
{
	struct device *dev = DEVICE_CAST(obj);
	return device_read(dev, p, *count, count, flags);
}

static kern_status_t device_object_write(struct object *obj, const void *p, size_t *count, socks_flags_t flags)
{
	struct device *dev = DEVICE_CAST(obj);
	return device_write(dev, p, *count, count, flags);
}

static kern_status_t device_object_destroy(struct object *obj)
{
	return KERN_OK;
}

static kern_status_t device_object_query_name(struct object *obj, char out[OBJECT_NAME_MAX])
{
	struct device *dev = DEVICE_CAST(obj);
	if (!dev) {
		return KERN_INVALID_ARGUMENT;
	}

	strncpy(out, dev->dev_name, OBJECT_NAME_MAX - 1);
	out[OBJECT_NAME_MAX - 1] = 0;

	return KERN_OK;
}

static kern_status_t device_object_get_child_at(struct object *obj, size_t at, struct object **out)
{
	struct device *dev = DEVICE_CAST(obj);
	size_t i = 0;
	queue_foreach(struct device, child, &dev->dev_children, dev_childent) {
		if (i == at) {
			*out = object_ref(&child->dev_base);
			return KERN_OK;
		}

		i++;
	}

	return KERN_NO_ENTRY;
}

static kern_status_t device_object_get_child_named(struct object *obj, const char *name, struct object **out)
{
	struct device *dev = DEVICE_CAST(obj);
	if (!dev) {
		return KERN_INVALID_ARGUMENT;
	}

	queue_foreach(struct device, child, &dev->dev_children, dev_childent) {
		if (!strcmp(child->dev_name, name)) {
			*out = object_ref(&child->dev_base);
			return KERN_OK;
		}
	}

	return KERN_NO_ENTRY;
}

static kern_status_t add_device_to_parent(struct device *dev, struct device *parent)
{
	kern_status_t status = KERN_OK;

	queue_foreach (struct device, child, &parent->dev_children, dev_childent) {
		if (!strcmp(dev->dev_name, child->dev_name)) {
			status = KERN_NAME_EXISTS;
			break;
		}
	}

	if (status != KERN_OK) {
		return status;
	}

	queue_push_back(&parent->dev_children, &dev->dev_childent);

	return KERN_OK;
}

kern_status_t device_register(struct device *dev, struct driver *owner, struct device *parent)
{
	unsigned long flags;
	device_lock_irqsave(dev, &flags);

	if (dev->dev_owner) {
		struct driver *prev_owner = dev->dev_owner;

		/* migrate device to new driver */
		driver_lock(prev_owner);
		driver_remove_device(prev_owner, dev);
		driver_free_minor(prev_owner, dev->dev_minor);
		dev->dev_minor = DEV_MINOR_INVALID;
		dev->dev_owner = NULL;
		driver_unlock(prev_owner);
	}

	driver_lock(owner);

	if (owner->drv_major == DEV_MAJOR_INVALID) {
		driver_unlock(owner);
		device_unlock_irqrestore(dev, flags);
		/* TODO better error message for lack of resources? */
		return KERN_INVALID_ARGUMENT;
	}

	unsigned int minor = driver_alloc_minor(owner);

	if (minor == DEV_MINOR_INVALID) {
		driver_unlock(owner);
		device_unlock_irqrestore(dev, flags);
		/* TODO better error message for lack of resources? */
		return KERN_BUSY;
	}

	kern_status_t status = KERN_OK;
	if (parent) {
		status = add_device_to_parent(dev, parent);
	}

	if (status != KERN_OK) {
		driver_unlock(owner);
		device_unlock_irqrestore(dev, flags);
		return status;
	}

	dev->dev_minor = minor;
	dev->dev_owner = owner;

	driver_add_device(owner, dev);

	if (type_ops[dev->dev_type] && type_ops[dev->dev_type]->register_device) {
		status = type_ops[dev->dev_type]->register_device(dev);
	}

	/* TODO remove device if registration failed */

	driver_unlock(owner);
	device_unlock_irqrestore(dev, flags);

	return status;
}