#include "lupi.h"
#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/ioctl.h>
#include <linux/fs.h>
#include <linux/hdreg.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <errno.h>
#include <string.h>

struct blkdev {
	int fd;
	int secsize;
	int platters;
	size_t size;
	char * buffer;
};

int blk_open(lua_State * L) {
	char * path = luaL_checkstring(L, 1);
	char rpath[255];
	sprintf(rpath, "/dev/%s", path);
	int fd = open(rpath, O_RDWR | O_DSYNC);
	if (fd == -1) {
		luaL_error(L, "Failed to open %s: %s", rpath, strerror(errno));
	}
	struct blkdev * dev = lua_newuserdata(L, sizeof(struct blkdev));
	dev->fd = fd;
	struct hd_driveid driveid;
	if (ioctl(fd, HDIO_GET_IDENTITY, &driveid) < 0) {
		//luaL_error("Failed to get identity for %s", rpath);
		dev->platters = 3;
		// Use another ioctl call to get sector size.
		if (!ioctl(fd, BLKSSZGET, &dev->secsize)) {
			// Fuck it, i guess it's 512 bytes.
			dev->secsize = 512;
			// Maybe a warning here?
		}
		dev->size = lseek(fd, 0, SEEK_END);
		lseek(fd, 0, SEEK_SET);
	} else {
		dev->size = driveid.lba_capacity_2*driveid.sector_bytes;
		dev->platters = driveid.cur_heads/2;
		dev->secsize = driveid.sector_bytes;
		if (dev->size == 0) {
			dev->size = lseek(fd, 0, SEEK_END);
			lseek(fd, 0, SEEK_SET);
		}
		if (dev->platters == 0) {
			dev->platters = 3;
		}
		if (dev->secsize == 0) {
			if (!ioctl(fd, BLKSSZGET, &dev->secsize)) {
				// Fuck it, i guess it's 512 bytes.
				dev->secsize = 512;
				// Maybe a warning here too?
			}
		}
	}
	dev->buffer = malloc(dev->secsize);
	return 1;
}

int blk_readbyte(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	off_t offset = luaL_checkinteger(L, 2);
	lseek(dev->fd, offset-1, SEEK_SET);
	uint8_t c = 0;
	size_t ioc = 0;
	if ((ioc = read(dev->fd, &c, 1)) != 1) {
		luaL_error(L, "io error! %llu != %llu", ioc, 1);
	}
	lua_pushinteger(L, c);
	return 1;
}

int blk_writebyte(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	off_t offset = luaL_checkinteger(L, 2);
	long byte = luaL_checkinteger(L, 3);
	if (byte > 255 || byte < 0)
		luaL_error(L, "%lld is out of range (should be between 0 and 255)", byte);
	uint8_t i = byte & 0xFF;
	lseek(dev->fd, offset-1, SEEK_SET);
	size_t ioc = 0;
	if ((ioc = write(dev->fd, &i, 1)) != 1) {
		luaL_error(L, "io error! %llu != %llu", ioc, 1);
	}
	return 0;
}

int blk_readsector(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	off_t offset = luaL_checkinteger(L, 2);
	lseek(dev->fd, (offset-1)*dev->secsize, SEEK_SET);
	size_t ioc = 0;
	if ((ioc = read(dev->fd, dev->buffer, dev->secsize)) != dev->secsize) {
		luaL_error(L, "io error! %llu != %llu", ioc, dev->secsize);
	}
	lua_pushlstring(L, dev->buffer, dev->secsize);
	return 1;
}

int blk_writesector(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	off_t offset = luaL_checkinteger(L, 2);
	size_t len = 0;
	char * sec = luaL_checklstring(L, 3, &len);
	lseek(dev->fd, (offset-1)*dev->secsize, SEEK_SET);
	if (len > dev->secsize) {
		len = dev->secsize;
	}
	size_t ioc = 0;
	if ((ioc = write(dev->fd, sec, len)) != len) {
		luaL_error(L, "io error! %llu != %llu", ioc, dev->secsize);
	}
	return 0;
}

int blk_getsecsize(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	lua_pushinteger(L, dev->secsize);
	return 1;
}

int blk_getsize(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	lua_pushinteger(L, dev->size);
	return 1;
}

int blk_getplatcount(lua_State * L) {
	struct blkdev * dev = lua_touserdata(L, 1);
	lua_pushinteger(L, dev->platters);
	return 1;
}

void blk_start(lua_State * L) {
	struct luaL_Reg blklib[] = {
		{"open", blk_open},
		{"readbyte", blk_readbyte},
		{"writebyte", blk_writebyte},
		{"readsector", blk_readsector},
		{"writesector", blk_writesector},
		{"secsize", blk_getsecsize},
		{"size", blk_getsize},
		{"platcount", blk_getplatcount},
		{NULL, NULL}
	};
	luaL_openlib(L, "blk", blklib, 0);
}