#define MODULE_NAME "drm"
#include "../twm.h"
#include <stdlib.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <unistd.h>
#include <xf86drm.h>
#include <xf86drmMode.h>
#include <drm_fourcc.h>

#define max(a,b) (((a)>(b))?(a):(b))

typedef struct {
	uint32_t id; // Unknown use
	uint32_t width;
	uint32_t height;
	uint32_t pitch;
	uint32_t handle;
	uint64_t size;
	uint8_t *data;
} DRMFb;

typedef struct _DRMConnector {
	uint32_t id;
	bool enabled;

	drmModeCrtcPtr old_mode;

	uint32_t crtc_id;
	drmModeModeInfo mode;

	Vec2 size;
	DRMFb frame;

	struct _DRMConnector *next;
} DRMConnector;

typedef struct {
	int drm_fd;
	DRMFb frameBuffer;
	DRMConnector *connectors;
} DRMState;

static DRMState *data;

uint32_t find_crtc(int fd, drmModeResPtr resources, drmModeConnectorPtr connector, uint32_t *taken) {
	for (int i = 0; i < connector->count_encoders; i++) {
		drmModeEncoderPtr encoder = drmModeGetEncoder(fd, connector->encoders[i]);
		if (!encoder) continue;
		for (int j = 0; j < resources->count_crtcs; j++) {
			if ((encoder->possible_crtcs & BIT(j)) == 0) continue; // Not compatible
			if ((*taken & BIT(j)) != 0) continue; // Already taken
			drmModeFreeEncoder(encoder);
			*taken |= BIT(j);
			return resources->crtcs[j];
		}
		drmModeFreeEncoder(encoder);
	}
	LOG_WARNING("Failed to find CRTC for connector %d", connector->connector_id);
	return 0;
}

bool create_fb(int fd, DRMFb *out, Vec2 size) {
	if(drmModeCreateDumbBuffer(fd, size.x, size.y, 32, 0, &out->handle, &out->pitch, &out->size) < 0) return false;

	uint32_t handles[4] = { out->handle };
	uint32_t pitches[4] = { out->pitch };
	uint32_t offsets[4] = { 0 };

	if (drmModeAddFB2(fd, size.x, size.y, DRM_FORMAT_XRGB8888, handles, pitches, offsets, &out->id, 0) < 0) goto err_dumb;

	uint64_t offset;
	if(drmModeMapDumbBuffer(fd,out->handle, &offset) < 0) goto err_dumb;

	out->data = mmap(0, out->size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, offset);
	if (out->data == MAP_FAILED) goto err_fb;
	return true;

err_fb:
	drmModeRmFB(fd, out->id);
err_dumb:
	drmModeDestroyDumbBuffer(fd, out->handle);
	return false;
}

void destroy_fb(int fd, DRMFb *fb) {
	munmap(fb->data, fb->size);
	drmModeRmFB(fd, fb->id);
	drmModeDestroyDumbBuffer(fd, fb->handle);
}

void display_put_pixel(Vec2 position, Color col) {
	DRMConnector *conn = data->connectors;
	while (conn) {
		if (conn->enabled) {
			if (position.x < 0 || position.x >= conn->size.x || position.y < 0 || position.y >= conn->size.y) continue;
			uint8_t *row = conn->frame.data + position.y*conn->frame.pitch;
			row[position.x * 4 + 0] = col.blue;
			row[position.x * 4 + 1] = col.green;
			row[position.x * 4 + 2] = col.red;
			row[position.x * 4 + 3] = 0x00;
		}
		conn = conn->next;
	}
}

void display_blit(Vec2 pos, Image image) {
	for (int x = 0; x < image.size.x; x++) {
		for (int y = 0; y < image.size.y; y++) {
			display_put_pixel((Vec2){x+pos.x,y+pos.y}, image.data[y][x]);
		}
	}
}

void display_flip() {
	DRMConnector *conn = data->connectors;
	while (conn) {
		if (conn->enabled) {
			drmModeClip clip = (drmModeClip){.x1 = 0, .y1 = 0, .x2 = conn->frame.width, .y2 = conn->frame.height};
			drmModeDirtyFB(data->drm_fd, conn->frame.id, &clip, 1);
		}
		conn = conn->next;
	}
}

Vec2 display_get_size() {
	int maxX = 0;
	int maxY = 0;
	DRMConnector *conn = data->connectors;
	while (conn) {
		if (conn->enabled) {
			maxX = max(maxX, conn->size.x);
			maxY = max(maxY, conn->size.y);
		}
		conn = conn->next;
	}
	return (Vec2){maxX, maxY};
}

//--------------------------------------------------------

const uint64_t FEATURE_LISTS[] = {
	BIT(BASIC_RENDERING) | BIT(ADVANCED_RENDERING),
	0
};

const FeatureFunc FUNCTIONS[] = {
	(FeatureFunc){ DISPLAY_PUT_PIXEL,    display_put_pixel },
	(FeatureFunc){ DISPLAY_BLIT,         display_blit },
	(FeatureFunc){ DISPLAY_GET_SIZE,     display_get_size },
	(FeatureFunc){ DISPLAY_FLIP,         display_flip },
	(FeatureFunc){ FUNC_END,             NULL }
};

FeatureFunc* module_init(size_t features_selected) {
	data = malloc(sizeof(DRMState));
	if (data == NULL) return NULL;

	data->drm_fd = open("/dev/dri/card0", O_RDWR);

	drmModeResPtr resources = drmModeGetResources(data->drm_fd);
	twm_assert_custom(resources, "Failed to find drm resources", return NULL);

	data->connectors = NULL;
	uint32_t taken_crtcs = 0;

	for (int i = 0; i < resources->count_connectors; i++) {
		drmModeConnectorPtr drm_connector = drmModeGetConnector(data->drm_fd, resources->connectors[i]);
		twm_assert_custom(drm_connector != NULL, "Failed to open connector", continue);
		
		DRMConnector *connector = malloc(sizeof(DRMConnector));
		twm_assert_custom(connector != NULL, "Failed to allocate connector data", goto cleanup);
		connector->next = data->connectors;
		data->connectors = connector;

		connector->id = drm_connector->connector_id;
		connector->enabled = drm_connector->connection == DRM_MODE_CONNECTED;

		LOG_INFO("Found connector %d (%s)", connector->id, connector->enabled?"Enabled":"Disabled");

		if (!connector->enabled) goto cleanup;
		twm_assert_custom(drm_connector->modes != NULL, "Connector has no modes", connector->enabled = false; goto cleanup);

		connector->crtc_id = find_crtc(data->drm_fd, resources, drm_connector, &taken_crtcs);
		twm_assert_custom(connector->crtc_id != 0, "Failed to find crtc", connector->enabled = false; goto cleanup);

		LOG_INFO("\tFound CRTC %d", connector->crtc_id);

		connector->mode = drm_connector->modes[0];
		connector->size = (Vec2){connector->mode.hdisplay, connector->mode.vdisplay};

		LOG_INFO("\tUsing Mode %dx%d", connector->size.x, connector->size.y);

		twm_assert_custom(create_fb(data->drm_fd, &connector->frame, connector->size), "Failed to make framebuffer", connector->enabled = false; goto cleanup);
		LOG_INFO("\tMade framebuffer with id %d", connector->frame.id);

		connector->old_mode = drmModeGetCrtc(data->drm_fd, connector->crtc_id);

		twm_assert_custom(drmModeSetCrtc(data->drm_fd, connector->crtc_id, connector->frame.id, 0, 0, &connector->id, 1, &connector->mode) >= 0, "Failed to set CRTC mode", connector->enabled = false; goto cleanup);

	cleanup:
		drmModeFreeConnector(drm_connector);
	}

	drmModeFreeResources(resources);

	return (FeatureFunc*)FUNCTIONS;
}


void module_close() {
	while (data->connectors) {
		if (data->connectors->enabled) {
			destroy_fb(data->drm_fd, &data->connectors->frame);
			drmModeCrtcPtr crtc = data->connectors->old_mode;
			if (crtc) {
				drmModeSetCrtc(data->drm_fd, crtc->crtc_id, crtc->buffer_id,
					crtc->x, crtc->y, &data->connectors->id, 1, &crtc->mode);
				drmModeFreeCrtc(crtc);
			}
		}
		DRMConnector *tmp = data->connectors->next;
		free(data->connectors);
		data->connectors = tmp;
	}
	close(data->drm_fd);
	free(data);
}