#include <pspkernel.h>
#include <pspdebug.h>
#include <pspctrl.h>
#include <string.h>
#include <math.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 <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 <curl/curl.h>
#include "share.h"
#include "my_socket.h"
#include "adhoc.h"
int sceNetResolverCreate(int *rid, void *buf, SceSize buflen);
int sceNetResolverStartNtoA(int rid, const char *hostname, u32* in_addr, unsigned int timeout, int retry);
int sceNetResolverStartAtoN(int rid, const u32* in_addr, char *hostname, SceSize hostname_len, unsigned int timeout, int retry);
int sceNetResolverStop(int rid);
int sceNetInetInetAton(const char* host, u32* in_addr);

typedef struct
{
	SOCKET sock;
	struct sockaddr_in addrTo;
	int serverSocket;
} Socket;
static const char* wlanNotInitialized = "WLAN not initialized.";
static int wlanInitialized = 0;
static char resolverBuffer[1024*10];
static int resolverId;

static int isconnected = 1;
// hack: this should be moved to PSPSDK, but by someone who knows how to do the in_addr mess the right way :-)
static void netInit(void)
{
	sceNetInit(128*1024, 42, 4*1024, 42, 4*1024);
	sceNetInetInit();
	sceNetApctlInit(0x8000, 48);
}
static void netTerm(void)
{
	sceNetApctlTerm();
	sceNetInetTerm();
	sceNetTerm();
	sceUtilityNetconfShutdownStart();
}

static int netDialog()
{
	oslInitNetDialog();
	OSL_IMAGE *img = oslCreateSwizzledImage(OSL_SECONDARY_BUFFER, OSL_IN_RAM);
	int endmess = 0;
while(endmess == 0){
		oslStartDrawing();
		oslClearScreen(0);
		oslDrawImageXY(img, 0, 0);
			oslDrawDialog();
			if (oslGetDialogStatus() == PSP_UTILITY_DIALOG_NONE){
				oslEndDialog();
				endmess = 1;
			}
	        oslEndDrawing();
        	oslSwapBuffers();
    }
	oslFreeImageData(img);
	return 1;
}
static int connectTo(int connectNum)
{
	int stateLast = -1;
	// WLAN
	pspSdkInetInit();
	while (1)
	{
		int state;
		sceNetApctlGetState(&state);
		if (state > stateLast)
		{
			printf("State %i \n", state);
			stateLast = state;
		}
		if (state == 4)
			break;  // connected with static IP
		// wait a little before polling again
		sceKernelDelayThread(200*1000); // 200 ms
	}
	return 1;
}
static int Wlan_init(lua_State *L)
{
	int argc = lua_gettop(L); 
	if (argc == 1){
		connectTo(luaL_checknumber(L, 1));
		isconnected = sceWlanGetSwitchState();
		if (isconnected != 0){
			wlanInitialized = 1;
			lua_pushnumber(L, 1);
		}
		if (isconnected == 0){
			netTerm();
			wlanInitialized = 0;
			lua_pushnumber(L, 0);
		}
		union SceNetApctlInfo szMyIPAddr;
		if (sceNetApctlGetInfo(8, &szMyIPAddr) != 0){
			netTerm();
			wlanInitialized = 0;
			return 0;
		}
		return 1;
	}
	if(wlanInitialized != 1){
		netInit();
		if (sceNetResolverInit() < 0) ShowMessageDialog("ERROR", 0);
		if (sceNetResolverCreate(&resolverId, resolverBuffer, sizeof(resolverBuffer)) < 0) ShowMessageDialog("ERROR", 0);
		netDialog();
		sceUtilityNetconfShutdownStart();
		isconnected = sceWlanGetSwitchState();
		if (isconnected != 0){
			wlanInitialized = 1;
			lua_pushnumber(L, 1);
		}
		if (isconnected == 0){
			netTerm();
			wlanInitialized = 0;
			lua_pushnumber(L, 0);
		}
		union SceNetApctlInfo szMyIPAddr;
		if (sceNetApctlGetInfo(8, &szMyIPAddr) != 0){
			netTerm();
			wlanInitialized = 0;
			lua_pushnumber(L, 0);
		}
	}
	return 1;
}
static int Wlan_term(lua_State *L)
{
	if(wlanInitialized != 0){
		wlanInitialized = 0;
		netTerm();
	}
	return 0;
}
static int Wlan_getIPAddress(lua_State* L)
{
	if (!wlanInitialized) return luaL_error(L, "wlanNotInitialized");
	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));
}
extern void ShowMessageDialog(const char *message, int enableYesno);
static int Socket_connect(lua_State *L)
{
	int isconnected;
	int err = -1;
	if (!wlanInitialized) return luaL_error(L, "wlanNotInitialized or connected to Wifi");
	if (lua_gettop(L) != 2) return luaL_error(L, "host and port expected");
	sceKernelDelayThread(100*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);
	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);
	// connect
	err = sceNetInetConnect(socket->sock, &socket->addrTo, sizeof(socket->addrTo));
	//Timeout
	int inetErr = sceNetInetGetErrno();
	
	if (err == -1) {  // returns 0x7d, I don't know why
		lua_pushnumber(L, inetErr);
	} else {
		lua_pushnumber(L, 0);
	}
	isconnected = sceWlanGetSwitchState();
		if (isconnected == 0){
			netTerm();
			wlanInitialized = 0;
			return luaL_error(L, "Socket Connect cancled by User");
		}
	return 2;
}

static int Socket_isConnected(lua_State *L)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 1) return luaL_error(L, "no argument expected.");
	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;
	}
	int isconnected = sceWlanGetSwitchState();
		if (isconnected == 0){
			netTerm();
			wlanInitialized = 0;
			return luaL_error(L, "Socket Connect cancled by User");
		}
	lua_pushboolean(L, 0);
	return 1;
}

static int Socket_createServerSocket(lua_State *L)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 1) return luaL_error(L, "port expected."); 
	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);
        if (socket->sock <= 0) {
		return luaL_error(L, "invalid socket."); 
        }
        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));
        if (err != 0) {
		return luaL_error(L, "bind error."); 
        }
	setSockNoBlock(socket->sock, 1);
        err = sceNetInetListen(socket->sock, 1);
        if (err != 0) {
		return luaL_error(L, "listen error."); 
        }
        return 1;
}

static int Socket_accept(lua_State *L)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 1) return luaL_error(L, "no argument expected.");
	Socket* socket = *toSocket(L, 1);
	if (!socket->serverSocket) return luaL_error(L, "accept allowed for server sockets only.");
	// 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)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 2) return luaL_error(L, "Byte size expected");
	Socket* socket = *toSocket(L, 1);
	if (socket->serverSocket) return luaL_error(L, "recv not allowed for server sockets.");
	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)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 2) return luaL_error(L, "one argument expected.");
	Socket* socket = *toSocket(L, 1);
	if (socket->serverSocket) return luaL_error(L, "send not allowed for server sockets.");
	size_t size;
	const char *string = luaL_checklstring(L, 2, &size);
	if (!string) return luaL_error(L, "Socket:write expected a string.");
	int result = sceNetInetSend(socket->sock, string, size, 0);
	lua_pushnumber(L, result);
	return 1;
}

static int Socket_close(lua_State *L)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 1) return luaL_error(L, "no argument expected.");
	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)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 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_DGRAM, 0);
	if (socket->sock & 0x80000000) {
		return luaL_error(L, "invalid socket."); 
	}
	setSockNoBlock(socket->sock, 1);
    	err = sceNetInetBind(socket->sock, &socket->addrTo, sizeof(socket->addrTo));
	return 1;
}
static int Socket_recv_udp(lua_State *L)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 1) return luaL_error(L, "no argument expected.");
	Socket* socket = *toSocket(L, 1);
	if (socket->serverSocket) return luaL_error(L, "recv not allowed for server sockets.");
	    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 == -1) return luaL_error(L, "something catastrophic happened, socket receive is locked up");
	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)
{
	if (!wlanInitialized) return luaL_error(L, wlanNotInitialized);
	if (lua_gettop(L) != 2) return luaL_error(L, "one argument expected.");
	Socket* socket = *toSocket(L, 1);
	if (socket->serverSocket) return luaL_error(L, "send not allowed for server sockets.");
	size_t size;
	const char *string = luaL_checklstring(L, 2, &size);
	if (!string) return luaL_error(L, "Socket:write expected a string.");
	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 bool adhocInitialized = 0;
static bool adhocConnected = 0;

static int Adhoc_init(lua_State *L){
	if(adhocInitialized != 1){
		if (sceWlanGetSwitchState() !=1) return luaL_error(L, "Turn Wifi Switch on");
		initAll();
		adhocInitialized = 1;
	}
	return 0;
}
static int Adhoc_term(lua_State *L){
	if(adhocInitialized != 0){
		termAll();
		adhocInitialized = 0;
	}	
	return 0;
}

static int adhoc_Connect(lua_State* L)
{
	//Check for arguments
	if(adhocInitialized != 1) return luaL_error(L, "Adhoc not Init.");
	connectAdhoc();
	adhocConnected = 1;
	return 0;
}


static int adhoc_GetState(lua_State* L)
{
	int err;
	int state;
	if(adhocInitialized)
	{
		//Check for arguments
		if(lua_gettop(L) != 0) return luaL_error(L, "no arguments expected for Ad-Hoc getstate.");
		//Get the state of adhoc control connect returns 0 if still connecting 1 if connected
		err = sceNetAdhocctlGetState(&state);
		//Checks for errors... returns a -1 if there was an error else 0 or 1
		if(err < 0)
			lua_pushnumber(L, -1);
		else
			lua_pushnumber(L, state);

		if(state == 1)
			adhocConnected = 1;
	}
	else
	lua_pushnumber(L, 0);
	return 1;
}


static int adhoc_sendPDP(lua_State *L)
{
if(lua_gettop(L) != 1) return luaL_error(L, "String expected");
if(adhocInitialized != 1) return luaL_error(L, "Adhoc not Init.");
	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)
{
if(lua_gettop(L) != 0) return luaL_error(L, "no arguments expected");
if(adhocInitialized != 1) return luaL_error(L, "Adhoc not Init.");
	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);
	if(mac)
		lua_pushstring(L, mac);
	else
		lua_pushstring(L, "MAC unavaliable");
	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 Wlan_functions[] = {
	{"init", Wlan_init},
	{"term", Wlan_term},
	{"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_openlib(L, "Wlan", Wlan_functions, 0);
	luaL_openlib(L, "Adhoc", Adhoc_functions, 0);
	Socket_register(L);
}
