#include <stdio.h>
#include <stdlib.h>

#include <lua.h>
#include <lualib.h>
#include <luaT/luaT.h>

#include "../threads/lib/luaTHRD.h"
#include "Device.h"


const char *fastnn_device_tname = "fastnn.CDevice";


static int device_new(lua_State *L)
{
	Device *device = NULL;
	if(lua_gettop(L) == 0)
        {
                device = Device_new();
        }
        else if(lua_gettop(L) == 1)
        {
                long id = luaL_checkinteger(L, 1);
                device = Device_newWithId(id);
        }
        else
                luaL_error(L, "device: device new invalid arguments");
        if (!device)
                luaL_error(L, "device: device failed");

        luaTHRD_pushudata(L, device, fastnn_device_tname);

        return 1;
}

static int device_init(lua_State *L)
{
	Device *device = luaTHRD_checkudata(L, 1, fastnn_device_tname);
	Status status;
	Initialize(device, &status);
	NERV_LUA_CHECK_STATUS(L, status);
	return 0;
}	

static int device_select_gpu(lua_State *L)
{
	Device *device = luaTHRD_checkudata(L, 1, fastnn_device_tname);
	
	Status status;
	if(lua_gettop(L) == 2)
	{
		int gpuid = luaL_checkinteger(L, 2);
		SelectGPU(device, gpuid, &status);
		NERV_LUA_CHECK_STATUS(L, status);
		return 0;
	}
	else if(lua_gettop(L) == 1)
	{
		int gpuid = AutoSelectGPU(device, &status);
		NERV_LUA_CHECK_STATUS(L, status);
		lua_pushinteger(L, gpuid);
		return 1;
	}
	else
		luaL_error(L, "device: device select gpu failed");
}	

static int device_id(lua_State *L)
{
	Device *device = luaTHRD_checkudata(L, 1, fastnn_device_tname);
	lua_pushinteger(L, Device_id(device));
	return 1;
}

static int device_tostring(lua_State *L)
{
        char str[STRLEN];
        Device* device = luaTHRD_checkudata(L, 1, fastnn_device_tname);
        snprintf(str, STRLEN, "%s <%lx>", fastnn_device_tname, Device_id(device));
        lua_pushstring(L, str);
        return 1;
}

static int device_destroy(lua_State *L)
{
        Device *device = luaTHRD_checkudata(L, 1, fastnn_device_tname);
        Device_destroy(device);
        //printf("device_destroy ... end\n");
        return 0;
}


//////////////////////////////////////////////

static int context_new(lua_State *L) 
{
        CuContext *context = NULL;
        if(lua_gettop(L) == 0)
        {   
                context = CuContext_new();
        }   
        else if(lua_gettop(L) == 1)
        {   
                long id = luaL_checkinteger(L, 1); 
                context = CuContext_newWithId(id);
        }   
        else
                luaL_error(L, "device: context new invalid arguments");
        if (!context)
                luaL_error(L, "device: context failed");

        luaTHRD_pushudata(L, context, nerv_context_tname);

        return 1;
}


static int context_id(lua_State *L)
{
        CuContext *context = luaTHRD_checkudata(L, 1, nerv_context_tname);
        lua_pushinteger(L, CuContext_id(context));
        return 1;
}

static int context_tostring(lua_State *L)
{
        char str[STRLEN];
        CuContext* context = luaTHRD_checkudata(L, 1, nerv_context_tname);
        snprintf(str, STRLEN, "%s <%lx>", nerv_context_tname, CuContext_id(context));
        lua_pushstring(L, str);
        return 1;
}

static int context_destroy(lua_State *L)
{
	CuContext* context = luaTHRD_checkudata(L, 1, nerv_context_tname);
        CuContext_destroy(context);
        return 0;
}


static const struct luaL_Reg device__ [] = {
  {"new", device_new},
  {"__tostring", device_tostring},
  {"id", device_id},
  {"init", device_init},
  {"select_gpu", device_select_gpu},
  {"free", device_destroy},
  {NULL, NULL}
};

static const struct luaL_Reg context__ [] = {
  {"new", context_new},
  {"__tostring", context_tostring},
  {"id", context_id},
  {"free", context_destroy},
  {NULL, NULL}
};


void fastnn_init_device(lua_State *L)
{
        luaT_newmetatable(L, fastnn_device_tname, NULL, device_new, device_destroy, NULL);
        luaL_register(L, NULL, device__);
        lua_pop(L, 1);
}

void fastnn_init_context(lua_State *L)
{
	luaT_newmetatable(L, nerv_context_tname, NULL, context_new, context_destroy, NULL);
        luaL_register(L, NULL, context__);
        lua_pop(L, 1);
}