#include <pspkernel.h>
#include <pspdebug.h>
#include <pspctrl.h>
#include <string.h>
#include <fastmath.h>
#include <psputility.h>
#include <stdlib.h>
#include <unistd.h>
#include <psputility_netmodules.h>
#include <psputility_netparam.h>
#include <pspgu.h>
#include <pspgum.h>
#include <pspsdk.h>
#include <psphttp.h>
#include <pspnet.h>
#include <pspnet_inet.h>
#include <pspnet_apctl.h>
#include <pspwlan.h>
#include <pspnet_adhoc.h>
#include <pspnet_adhocctl.h>
#include <oslib/oslib.h>
#include "share.h"
#include "my_socket.h"
#include "adhoc.h"
extern int sceNetResolverCreate(int *rid, void *buf, SceSize buflen);
extern int sceNetResolverStartNtoA(int rid, const char *hostname, u32* in_addr, unsigned int timeout, int retry);
extern int sceNetResolverStartAtoN(int rid, const u32* in_addr, char *hostname, SceSize hostname_len, unsigned int timeout, int retry);
extern int sceNetResolverStop(int rid);
extern int sceNetInetInetAton(const char* host, u32* in_addr);

typedef struct
{
	SOCKET sock;
	struct sockaddr_in addrTo;
	bool serverSocket;
} Socket;
static int endmesswlan = 1;
static int netok = 1;
static char resolverBuffer[1024];
static int resolverId;
void netInit(void)
{
	sceNetInit(128*1024, 42, 4*1024, 42, 4*1024);

	sceNetInetInit();

	sceNetApctlInit(0x8000, 48);

}
static int Wlaninit(void)
{	
	if(endmesswlan == 1){
		netok = 0;
		oslInitNetDialog();
		endmesswlan = 0;
	}
	if(endmesswlan == 0){
		oslDrawDialog();
		if (oslGetDialogStatus() == PSP_UTILITY_DIALOG_NONE){
			oslEndDialog();
			endmesswlan = 1;
			return 1;
		}
		return 0;
	}
}
static int wlan(lua_State *L)
{
	netInit();
	sceNetResolverCreate(&resolverId, resolverBuffer, sizeof(resolverBuffer));
	lua_pushnumber(L, Wlaninit());
	return 1;
}
static int end(lua_State *L)
{
	oslNetTerm();
	netok = 1;
	return 0;
}

// hack: this should be moved to PSPSDK, but by someone who knows how to do the in_addr mess the right way :-)
static int Wlan_getIPAddress(lua_State* L)
{
	union SceNetApctlInfo szMyIPAddr;
	if (sceNetApctlGetInfo(8, &szMyIPAddr) != 0) return 0;
	lua_pushstring(L, szMyIPAddr.ip);
	return 1;
}
UserdataStubs(Socket, Socket*)
static int Socket_free(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	sceNetInetClose(socket->sock);
	free(socket);
	return 0;
}

unsigned short htons(unsigned short wIn)
{
    u8 bHi = (wIn >> 8) & 0xFF;
    u8 bLo = wIn & 0xFF;
    return ((unsigned short)bLo << 8) | bHi;
}

static int setSockNoBlock(SOCKET s, u32 val)
{ 
    return sceNetInetSetsockopt(s, SOL_SOCKET, 0x1009, (const char*)&val, sizeof(u32));
}

static int Socket_connect(lua_State *L)
{
	int isconnected;
	int argc = lua_gettop(L); 
	if (argc != 2) return luaL_error(L, "host and port expected"); 
	sceKernelDelayThread(50*1000); // without this delay it doesn't work sometime

	Socket** luaSocket = pushSocket(L);
	Socket* socket = (Socket*) malloc(sizeof(Socket));
	*luaSocket = socket;
	socket->serverSocket = 0;
	// resolve host
	const char *host = luaL_checkstring(L, 1);
	int port = luaL_checkint(L, 2);
	socket->addrTo.sin_family = AF_INET;
	socket->addrTo.sin_port = htons(port);
	int err = sceNetInetInetAton(host, &socket->addrTo.sin_addr);
	if (err == 0) {
		err = sceNetResolverStartNtoA(resolverId, host, &socket->addrTo.sin_addr, 2, 3);
		//if (err < 0) return luaL_error(L, "Socket:connect: DNS resolving failed.");
	}
	//create non-blocking socket	
	socket->sock = sceNetInetSocket(AF_INET, SOCK_STREAM, 0);
	if (socket->sock & 0x80000000) {
		return luaL_error(L, "invalid socket."); 
	}
	setSockNoBlock(socket->sock, 1);
	u32 err2;
	// connect
	err = sceNetInetConnect(socket->sock, &socket->addrTo, sizeof(socket->addrTo));
	//Timeout
	u32 timeout = 1000000; // in microseconds == 1 sec
        err2 = sceNetInetSetsockopt(socket->sock, SOL_SOCKET, SO_RCVTIMEO, (char*)&timeout, sizeof(timeout));
        if (err2 != 0)
            return luaL_error(L, "set SO_RCVTIMEO failed\n");
	//end timeout
	int inetErr = sceNetInetGetErrno();
	if (err == -1 /*&& sceNetInetGetErrno() != 0x77 */) {  // returns 0x7d, I don't know why
		lua_pushnumber(L, inetErr);
	} else {
		lua_pushnumber(L, 0);
	}
	return 2;
}

static int Socket_isConnected(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	if (socket->serverSocket) {
		lua_pushboolean(L, 1);
		return 1;
	}
	// try connect again, which should always fail
	// look at why it failed to figure out if it is connected
	//REVIEW: a conceptually cleaner way to poll this?? (accept?)
	int err = sceNetInetConnect(socket->sock, &socket->addrTo, sizeof(socket->addrTo));
	if (err == 0 || (err == -1 && sceNetInetGetErrno() == 0x7F)) {
		// now connected - I hope
		lua_pushboolean(L, 1);
		return 1;
	}
	lua_pushboolean(L, 0);
	return 1;
}

static int Socket_createServerSocket(lua_State *L)
{
	int port = luaL_checkint(L, 1);
	Socket** luaSocket = pushSocket(L);
	Socket* socket = (Socket*) malloc(sizeof(Socket));
	*luaSocket = socket;
	socket->serverSocket = 1;
        socket->sock = sceNetInetSocket(AF_INET, SOCK_STREAM, 0);
        socket->addrTo.sin_family = AF_INET;
        socket->addrTo.sin_port = htons(port);
        socket->addrTo.sin_addr = 0;
        int err = sceNetInetBind(socket->sock, &socket->addrTo, sizeof(socket->addrTo));
	setSockNoBlock(socket->sock, 1);
        err = sceNetInetListen(socket->sock, 1);
        return 1;
}

static int Socket_accept(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	socket->serverSocket;
	// check for waiting incoming connections
        struct sockaddr_in addrAccept;
        int cbAddrAccept = sizeof(addrAccept);
        SOCKET sockClient = sceNetInetAccept(socket->sock, &addrAccept, &cbAddrAccept);
        if (sockClient <= 0) {
        	return 0;
        }
	// create new lua socket
	Socket** luaSocket = pushSocket(L);
	Socket* incomingSocket = (Socket*) malloc(sizeof(Socket));
	*luaSocket = incomingSocket;
	incomingSocket->serverSocket = 0;
	incomingSocket->sock = sockClient;
	incomingSocket->addrTo = addrAccept;
	return 1;
}

static int Socket_recv(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	socket->serverSocket;
	int datasize = luaL_checkint(L, 2);
	char data[datasize];
	int count = sceNetInetRecv(socket->sock, (u8*) &data, datasize, 0);
	if (count > 0) {
		lua_pushlstring(L, data, count);
	} else {
		lua_pushstring(L, "");
	}
	return 1;
}

static int Socket_send(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	socket->serverSocket;
	size_t size;
	const char *string = luaL_checklstring(L, 2, &size);
	int result = sceNetInetSend(socket->sock, string, size, 0);
	lua_pushnumber(L, result);
	return 1;
}

static int Socket_close(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	sceNetInetClose(socket->sock);
	return 0;
}

static int Socket_tostring(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	char buf[128];
	sprintf(buf, "%i.%i.%i.%i",
		socket->addrTo.sin_addr & 255,
		(socket->addrTo.sin_addr >> 8) & 255,
		(socket->addrTo.sin_addr >> 16) & 255,
		(socket->addrTo.sin_addr >> 24) & 255);
	lua_pushstring(L, buf);  // pushfstring doesn't work, returns userdata object
	return 1;
}


////UDP Functions Included in Socket 


static int Socket_connect_udp(lua_State *L)
{
	sceKernelDelayThread(50*1000); // without this delay it doesn't work sometime
	Socket** luaSocket = pushSocket(L);
	Socket* socket = (Socket*) malloc(sizeof(Socket));
	*luaSocket = socket;
	socket->serverSocket = 0;
	// resolve host
	const char *host = luaL_checkstring(L, 1);
	int port = luaL_checkint(L, 2);
	socket->addrTo.sin_family = AF_INET;
	socket->addrTo.sin_port = htons(port);
	int err = sceNetInetInetAton(host, &socket->addrTo.sin_addr);
	if (err == 0) {
		err = sceNetResolverStartNtoA(resolverId, host, &socket->addrTo.sin_addr, 2, 3);
		//if (err < 0) return luaL_error(L, "Socket:connect: DNS resolving failed.");
	}
	// create non-blocking socket	
	socket->sock = sceNetInetSocket(AF_INET, SOCK_DGRAM, 0);
	if (socket->sock & 0x80000000) {
	}
	setSockNoBlock(socket->sock, 1);
    	err = sceNetInetBind(socket->sock, &socket->addrTo, sizeof(socket->addrTo));
	return 1;
}
static int Socket_recv_udp(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	socket->serverSocket;
	    char buffer[256];
            struct sockaddr_in addrFrom;
            int cbAddrFrom = sizeof(addrFrom);
            u32 cb;
            cb = sceNetInetRecvfrom(socket->sock, (u8*) &buffer, sizeof(buffer), 0,&addrFrom, &cbAddrFrom);                   
	if (cb > 0) {
		lua_pushlstring(L, buffer, cb);
		lua_pushnumber(L,cb);
	} else {
		lua_pushstring(L, "");
	}
	return 2;
}
static int Socket_send_udp(lua_State *L)
{
	Socket* socket = *toSocket(L, 1);
	socket->serverSocket;
	size_t size;
	const char *string = luaL_checklstring(L, 2, &size);
	int cbAddrTo = sizeof(socket->addrTo);  
	int result = sceNetInetSendto(socket->sock, string, size, 0,&socket->addrTo, cbAddrTo);
	lua_pushnumber(L, result);
	return 1;
}
/////end UDP


//////Adhoc By PiCkDaT
static int Adhoc_init(lua_State *L){
	initAll();
	return 0;
}
static int Adhoc_term(lua_State *L){
	termAll();
	return 0;
}

static int adhoc_Connect(lua_State* L)
{
	connectAdhoc();
	return 0;
}


static int adhoc_GetState(lua_State* L)
{
	int err;
	int state;
	err = sceNetAdhocctlGetState(&state);
	if(err < 0)
		lua_pushnumber(L, -1);
	else
		lua_pushnumber(L, state);
	return 1;
}


static int adhoc_sendPDP(lua_State *L)
{
	size_t size;
	const char *string = luaL_checklstring(L, 1, &size);
	sendPDP((char*)string, strlen(string));
	return 0;
}


static int adhoc_recvPDP(lua_State *L)
{
	char recvData[4096] = "";
	unsigned int recvLEN;
	recvPDP(&recvData, &recvLEN);
	if(recvData != "")
		lua_pushlstring(L, recvData, strlen(recvData));
	else
		lua_pushnil(L);
	return 1;
}
static int adhoc_getMac(lua_State *L)
{
	char mac[18] = "";
	getmac(&mac);
	lua_pushstring(L, mac);
	return 1;
}
//end Matching
static const luaL_reg Socket_methods[] = {
	{"connect", Socket_connect},
	{"createServerSocket", Socket_createServerSocket},
	{"udpConnect", Socket_connect_udp},
	{"isConnected", Socket_isConnected},
	{"accept", Socket_accept},
	{"send", Socket_send},
	{"recv", Socket_recv},
	{"close", Socket_close},
	{"udpRecv", Socket_recv_udp},
	{"udpSend", Socket_send_udp},
	{0,0}
};

static const luaL_reg Socket_meta[] = {
	{"__gc", Socket_free},
	{"__tostring", Socket_tostring},
	{0,0}
};

UserdataRegister(Socket, Socket_methods, Socket_meta)

static const luaL_reg net[] = {
	{"init", wlan},
	{"term", end},
	{"getIP", Wlan_getIPAddress},
	{0, 0}
};

static const luaL_reg Adhoc_functions[] = {
	{"init", Adhoc_init},
	{"connect", adhoc_Connect},
	{"getState", adhoc_GetState},
	{"send", adhoc_sendPDP},
	{"recv", adhoc_recvPDP},
	{"term", Adhoc_term},
	{"getMac", adhoc_getMac},
	{0, 0}
};
void net_init(lua_State *L) {
	luaL_register(L, "Wlan", net);
	luaL_openlib(L, "Adhoc", Adhoc_functions, 0);
	Socket_register(L);
}