/*

LKS Sockets interface v0.03

LSOCK 0.04

TODO: Concurrency

LSOCK 0.03

Added data buffering to lsockets to enable more advanced blocking functions:

  LSock_ReadLine
    Reads up to a new line character
  LSock_ReadUpTo
    Reads up to a specified character pattern
  LSock_WriteLine
    Sends a buffer and then a newline character
  LSock_WriteString
    Sends a null-terminated string (without sending the null character)

TODO: Error checking everywhere
TODO: Check if there can be any concurrency problems...
TODO: Make GetSocketLocal and GetSocketRemote actually work!!
TODO: Make a GetSocketStats function
TODO: Make sure bytes_send and bytes_recv are updated everywhere they need to be.
TODO: SetSocketUserData and GetSocketUserData
TODO: Make the string searching function work properly! (eg. find "ab" in "aabb" would fail at the moment)

Made LSock_Startup and LSock_Shutdown call WSAStartup and WSACleanup.

LSOCK 0.02

  Basic lsocket functionality (blocking only)

*/

#ifndef __INCLUDE_LSOCK003
#define __INCLUDE_LSOCK003

#include <windows.h>
#include <winsock.h>
#include <stdlib.h>
#include <stdio.h>
#include "lmem003.h"



//---- DATA TYPES ------------------------------------------------------------//

typedef LMEM_HANDLE LSOCK;

#define LSOCK_TYPE_NONE        0x0000
#define LSOCK_TYPE_UNCONNECTED 0x0001
#define LSOCK_TYPE_LISTENING   0x0002
#define LSOCK_TYPE_SERVER      0x0003
#define LSOCK_TYPE_CLIENT      0x0004

#define LSOCK_TYPE_MASK        0x00FF

#define LSOCK_PROTO_TCPIP      0x0100
#define LSOCK_PROTO_UDPIP      0x0200
#define LSOCK_PROTO_MASK       0xFF00

#define LSOCK_ERROR            0xFFFFFFFF



//---- PRIVATE DATA TYPES ----------------------------------------------------//

//---- LSOCK_INFO - Socket Info structure

typedef struct {
  SOCKET s;                      //  WinSock socket aociated with this LSock

  UINT type;                     //  Combination of the TYPE and PROTO flags defined above

  ULONG ip_local, ip_remote;     //  Local and remote IP addresses and ports
  UINT port_local, port_remote;

  ULONG bytes_send, bytes_recv;  //  Send and receive total bytes for this connection

  ULONG buffer_size;             //  Current allocation size of read buffer
  ULONG buffer_min, buffer_max;  //  Minimum and maximum allowed allocation sizes
  ULONG buffer_delta;            //  Buffer resize ammount
  ULONG buffer_len;              //  Current number of byte in buffer

  LMEM_HANDLE buffer_handle;     //  LMem heap block handle of read buffer

  LPVOID userdata;
  CRITICAL_SECTION sync;         //  For synchronising accesses to this socket

} LSOCK_INFO, *LSOCK_PINFO;

//  Default values

#define LSOCK_BUFFER_MAX   8192
#define LSOCK_BUFFER_MIN   64
#define LSOCK_BUFFER_DELTA 256



//---- PRIVATE DATA ----------------------------------------------------------//

//#define LSOCK_DEBUG

LMEM_HANDLE lsock_heap;

UINT lsock_startup_count = 0;



//---- DEBUG ROUTINES --------------------------------------------------------//

/*
Functions for recording debug messages and events
In a non-debug build these all equate to null routines

LSock_Debug_Startup
LSock_Debug_Shutdown
LSock_Debug_Write
LSock_LogWrite
*/

#ifdef LSOCK_DEBUG

HANDLE LSock_log_output;

void LSock_Debug_Startup() {
  AllocConsole();
  LSock_log_output = GetStdHandle(STD_OUTPUT_HANDLE);
}

void LSock_Debug_Shutdown() {
}

void LSock_Debug_Write(char *msg) {
  int i;
  unsigned long j;
  for(i=0; msg[i]; i++);
  WriteConsole(LSock_log_output, msg, i, &j, NULL);
}

void LSock_Debug_WriteInt(int val) {
  char msg[12];
  sprintf(msg, "%d", val);
  LSock_Debug_Write(msg);
}

void LSock_Debug_WriteHex(char *buf, int len) {
  char msg[12];
  int i;
  for(i=0; i<len; i++) {
    sprintf(msg, "%02X ", (int)buf[i]);
    LSock_Debug_Write(msg);
  }
}

#define LSock_LogWrite(a) LSock_Debug_Write(a)

#else
#define LSock_Debug_Startup(a)
#define LSock_Debug_Shutdown(a)
#define LSock_Debug_Write(a)
#define LSock_Debug_WriteInt(a)
#define LSock_Debug_WriteHex(a, b)
#define LSock_LogWrite(a)
#endif



//---- INITIALISATION ROUTINES -----------------------------------------------//

/*
LSock_Startup
LSock_Shutdown
*/



//---- LSock_Startup - Initialises LSock library

//  Returns TRUE on success, FALSE on error.

BOOL LSock_Startup() {
  WSADATA info;

  if(!lsock_startup_count) {

    if(WSAStartup(0x0101, &info) == SOCKET_ERROR) {
      return FALSE;
    }

    LMem_Startup();           //  TODO: Make this return BOOL and check it here
    LSock_Debug_Startup();
    LSock_LogWrite("LSock_Startup - v0.003 Build B\n");
    lsock_heap = LMem_Heap_Create(); //  TODO: Error check this
  }
  lsock_startup_count ++;
  return TRUE;
}



//---- LSock_Shutdown - Closes down the LSock library

//  Returns TRUE if the library has been shut down, FALSE if this function
// needs to be called more times to shut down properly.

BOOL LSock_Shutdown() {
  if(!lsock_startup_count) return TRUE;
  lsock_startup_count--;
  if(!lsock_startup_count) {
    LSock_LogWrite("LSock_Shutdown\n");
    LSock_Debug_Shutdown();
    LMem_Shutdown();
    WSACleanup();
    return TRUE;
  }
  return FALSE;
}



//---- CONNECTION ROUTINES ---------------------------------------------------//

/*
Routines for creating and managing connections

LSock_Listen
LSock_Accept
LSock_Connect
LSock_CloseSocket
LSock_GetSocketLocal
LSock_GetSocketRemote
*/



//---- LSock_Listen - Listens on a port.

//  Returns the listening lsock handle

LSOCK LSock_Listen(UINT port) {

  LSOCK s;
  LSOCK_PINFO p;

  SOCKET sockServer;
  SOCKADDR_IN saServ;

  LSock_LogWrite("LSock_Listen - start\n");

  //- Allocate lsock structure

  s = LMem_Heap_Alloc(lsock_heap, sizeof(LSOCK_INFO));
  p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, s);

  //- Create a listening socket

  sockServer = socket(AF_INET, SOCK_STREAM, 0);
  if(sockServer == INVALID_SOCKET) {
    LSock_LogWrite("LSock_Listen - error #1\n");
    LMem_Heap_Free(lsock_heap, s);
    return LSOCK_ERROR;
  }

  //  Bind the socket to the internet address

  saServ.sin_family = AF_INET;
  saServ.sin_port = htons(port);
  saServ.sin_addr.s_addr = INADDR_ANY;

  if(bind(sockServer, (LPSOCKADDR)&saServ, sizeof(SOCKADDR_IN)) == SOCKET_ERROR) {
    LSock_LogWrite("LSock_Listen - error #2\n");
    closesocket(sockServer);
    LMem_Heap_Free(lsock_heap, s);
    return LSOCK_ERROR;
  }

  //  Listen for connections

  if(listen(sockServer,0) == SOCKET_ERROR) {
    LSock_LogWrite("LSock_Listen - error #3\n");
    closesocket(sockServer);
    LMem_Heap_Free(lsock_heap, s);
    return LSOCK_ERROR;
  }

  //  Store details

  p->type = LSOCK_TYPE_LISTENING | LSOCK_PROTO_TCPIP;
  p->s = sockServer;

  p->buffer_handle = LMEM_ERROR;

  LSock_LogWrite("LSock_Listen - end\n");

  return s;
}



//---- LSock_Accept - Accepts a connection on a listening socket.  Blocks if one not available.

LSOCK LSock_Accept(LSOCK lsServer) {

  SOCKADDR_IN sa;
  int i = sizeof(SOCKADDR_IN);

  LSOCK ls;
  LSOCK_PINFO p;

  LSOCK_PINFO pServer;

  LSock_LogWrite("LSock_Accept - start\n");

  //  Allocate lsock structure

  ls = LMem_Heap_Alloc(lsock_heap, sizeof(LSOCK_INFO));
  p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  pServer = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, lsServer);

  //  Accept connection

  p->s = accept(pServer->s, (LPSOCKADDR)&sa, NULL);
  if(p->s == INVALID_SOCKET) {
    LMem_Heap_Free(lsock_heap, ls);
    return INVALID_SOCKET;
  }

  //  Create data buffer

  p->buffer_min    = LSOCK_BUFFER_MIN;
  p->buffer_max    = LSOCK_BUFFER_MAX;
  p->buffer_delta  = LSOCK_BUFFER_DELTA;
  p->buffer_size   = p->buffer_min;
  p->buffer_handle = LMem_Heap_Alloc(lsock_heap, p->buffer_size);
  p->buffer_len    = 0;

  //  Store details

  p->type = LSOCK_TYPE_SERVER | LSOCK_PROTO_TCPIP;

  p->ip_remote = ntohl(sa.sin_addr.s_addr);
  p->port_remote = ntohs(sa.sin_port);

  getsockname(p->s, (LPSOCKADDR)&sa, &i);

  p->ip_local = ntohl(sa.sin_addr.s_addr);
  p->port_local = ntohs(sa.sin_port);
  
  LSock_LogWrite("LSock_Accept - end\n");

  return ls;
}




//---- LSock_Connect - Connects to a server

LSOCK LSock_Connect(char *server, ULONG ulServer, UINT port) {

  LSOCK ls;
  LSOCK_PINFO p;

  SOCKET s;
  SOCKADDR_IN sa;
  LPHOSTENT lpHost;
  ULONG saddr;
  int i;

  LSock_LogWrite("LSock_Connect - start\n");

  //- Allocate lsock structure

  ls = LMem_Heap_Alloc(lsock_heap, sizeof(LSOCK_INFO));
  p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);

  //- Create a socket

  s = socket(AF_INET, SOCK_STREAM, 0);
  if(s == INVALID_SOCKET) {
    return LSOCK_ERROR;
  }

  //- Find the server

  if(server) {
    if((lpHost = gethostbyname(server)) == NULL) {
      closesocket(s);
      return LSOCK_ERROR;
    }
    sa.sin_family = AF_INET;
    sa.sin_addr = *((LPIN_ADDR)*lpHost->h_addr_list);
    sa.sin_port = htons(port);
    saddr = ntohl(sa.sin_addr.s_addr);
  } else {
    sa.sin_family = AF_INET;
    sa.sin_port = htons(port);
    sa.sin_addr.s_addr = htonl(ulServer);
    saddr = ulServer;
  }

  //- Connect to server

  if(connect(s, (LPSOCKADDR)&sa, sizeof(SOCKADDR_IN)) == SOCKET_ERROR) {
    closesocket(s);
    return LSOCK_ERROR;
  }

  //- Get local details

  getsockname(s, (LPSOCKADDR)&sa, &i);

  //  Create data buffer

  p->buffer_min    = LSOCK_BUFFER_MIN;
  p->buffer_max    = LSOCK_BUFFER_MAX;
  p->buffer_delta  = LSOCK_BUFFER_DELTA;
  p->buffer_size   = p->buffer_min;
  p->buffer_handle = LMem_Heap_Alloc(lsock_heap, p->buffer_size);
  p->buffer_len    = 0;

  //  Store details

  p->s = s;
  p->type = LSOCK_TYPE_CLIENT | LSOCK_PROTO_TCPIP;

  p->ip_remote = saddr;
  p->port_remote = port;

  p->ip_local = ntohl(sa.sin_addr.s_addr);
  p->port_local = ntohs(sa.sin_port);

  LSock_LogWrite("LSock_Connect - end\n");

  return ls;
}



//---- LSock_CloseSocket - Closes a socket

int LSock_CloseSocket(LSOCK ls) {
  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  int i = closesocket(p->s);

  LSock_LogWrite("LSock_CloseSocket - start\n");

  if(p->buffer_handle != LMEM_ERROR) {
    LMem_Heap_Free(lsock_heap, p->buffer_handle);
  }
  LMem_Heap_Free(lsock_heap, ls);

  LSock_LogWrite("LSock_CloseSocket - end\n");

  if(i) return LSOCK_ERROR;
  return 0;
}



//---- LSock_GetSocketLocal

void LSock_GetSocketLocal(LSOCK ls, ULONG *ip, UINT *port) {
  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);

  LSock_LogWrite("LSock_GetSocketLocal - called\n");

  *port = p->port_local;
  *ip = p->ip_local;
}



//---- LSock_GetSocketRemote

void LSock_GetSocketRemote(LSOCK ls, ULONG *ip, UINT *port) {
  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);

  LSock_LogWrite("LSock_GetSocketRemote - called\n");

  *port = p->port_remote;
  *ip = p->ip_remote;
}



//---- SOCKET INFO MANIPULATION ROUTINES (Private) ---------------------------//

/*
LSock_Buffer_Find
LSock_Buffer_Get
LSock_Buffer_Remove
*/



//---- LSock_Buffer_Find - Find a character string in the socket's buffer

//  Returns -1 if bytes not found.  Starts at given start poition

int LSock_Buffer_Find(LSOCK ls, int start, char *find, int findlen) {
  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  char *buffer = (char *)LMem_Heap_GetPtr(lsock_heap, p->buffer_handle);
  int i, j;

  //  Throw out any error cases

  if(!find) return -1;
  if(findlen <= 0) return -1;
  if(findlen > p->buffer_len) return -1;
  if(start < 0) return -1;

  //  Find characters in buffer

  for(i=start, j=0; i<p->buffer_len && j<findlen; i++) {
    if(buffer[i] != find[j]) j=0;
    if(buffer[i] == find[j]) j++;
  }

  //  Return -1 if not found or the location if it is

  if(j < findlen) return -1;

  return i-j;
}



//---- LSock_Buffer_Get - Copies a block of bytes from a socket's buffer

//  Returns the ammount actually copied
//  If buf is NULL then it returns 0

int LSock_Buffer_Get(LSOCK ls, char *buf, int len) {
  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  char *buffer = (char *)LMem_Heap_GetPtr(lsock_heap, p->buffer_handle);
  int copylen = len;
  int i;

  //  Throw out error cases

  if(len <= 0) return 0;
  if(!buf) return 0;

  //  Copy data up to len bytes

  if(copylen > p->buffer_len) copylen = p->buffer_len;
  for(i=0; i<copylen; i++) buf[i] = buffer[i];

  return copylen;
}



//---- LSock_Buffer_Remove - Remove a block of bytes from a socket's buffer

// Returns the size of the remaining data in the buffer

int LSock_Buffer_Remove(LSOCK ls, int len) {
  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  char *buffer = (char *)LMem_Heap_GetPtr(lsock_heap, p->buffer_handle);
  int i;

  if(len <= 0) return p->buffer_len;

  if(p->buffer_len < len) {
    p->buffer_len = 0;
  } else {
    for(i=len; i<p->buffer_len; i++) buffer[i-len] = buffer[i];
    p->buffer_len -= len;
  }

  return p->buffer_len;
}



//---- DATA TRANSFER ROUTINES ------------------------------------------------//

/*
Reading and writing data

LSock_ReadData
LSock_WriteData
LSock_WriteString
LSock_ReadLine
LSock_WriteLine
*/



//---- LSock_ReadData - Reads data from a socket

//  Blocks until it can read some data or socket is closed by another thread.

//  TODO: Use LSock_Buffer_* functions in here instead of for loops :)

int LSock_ReadData(LSOCK ls, char *buf, int len) {

  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  int i;

  if(p->buffer_handle != LMEM_ERROR && p->buffer_len > 0) {

    //  Return data from socket buffer if there is any

    char *buffer = (char *)LMem_Heap_GetPtr(lsock_heap, p->buffer_handle);

    if(p->buffer_len > len) {

      //  If the buffer contains more than the user is requesting read that much
      // and move the rest of the data along in the buffer

      for(i=0; i<len; i++) buf[i] = buffer[i];
      for(i=len; i<p->buffer_len; i++) buffer[i-len] = buffer[i];
      p->buffer_len -= len;
      i = len;

    } else {

      //  Otherwise return all the data inthe buffer

      for(i=0; i<p->buffer_len; i++) buf[i] = buffer[i];
      p->buffer_len = 0;
    }
    
  } else {

    //  No data in socket buffer -> receive from network

    i = recv(p->s, buf, len, 0);
  }

  p->bytes_recv += i;

  if(i == 0) i = LSOCK_ERROR;

  return i;
}



//---- LSock_WriteData - Sends data on a socket.

//  Unlike send this will make sure all the data is sent.  Will block until data sent.

int LSock_WriteData(LSOCK ls, char *buf, int len) {

  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  int i, j=0;

  while(j<len) {
    i = send(p->s, buf+j, len-j, 0);
    if(i == SOCKET_ERROR) return LSOCK_ERROR;
    j += i;
    p->bytes_send += i;
  }

  return j;
}



//---- LSock_WriteString - Sends a string (not including the null terminator)

int LSock_WriteString(LSOCK ls, char *buf) {
  int i;
  for(i=0; buf[i]; i++);
  return LSock_WriteData(ls, buf, i);
}



//---- LSock_ReadLine - Reads data up to the next newline character

//  Reads up to a new line character or len bytes
//  Passing a NULL buf removes up to len chars of the next line

//  Returns LSOCK_ERROR on error - either socket error or buffer filled and no newline
//  Otherwise the number of characters in the line is returned - even if there were too many
//  The discard flag says whether this function hould throw away any excess chars on the line

int LSock_ReadLine(LSOCK ls, char *buf, int len, BOOL discard) {

  LSOCK_PINFO p = (LSOCK_PINFO)LMem_Heap_GetPtr(lsock_heap, ls);
  char *buffer = (char *)LMem_Heap_GetPtr(lsock_heap, p->buffer_handle);
  int linelen;
  char nl[2] = {0x0D, 0x0A};
  int i, j;

  //  Check if there's already a newline in the buffer

  linelen = LSock_Buffer_Find(ls, 0, nl, 2);

  //  Otherwise read data untilwe get a newline

  while(linelen == -1 && p->buffer_len < p->buffer_max) {

    //  Expand the buffer if there's less than delta free and the buffer isn't already at max size

    if(p->buffer_size-p->buffer_len < p->buffer_delta && p->buffer_size < p->buffer_max) {
      p->buffer_size += p->buffer_delta;
      if(p->buffer_size > p->buffer_max) p->buffer_size = p->buffer_max;
      buffer = (char *)LMem_Heap_ReAlloc(lsock_heap, p->buffer_handle, p->buffer_size);
    }

    //  Read more data if buffer not full

    if(p->buffer_size > p->buffer_len) {
      i = recv(p->s, buffer+p->buffer_len, p->buffer_size-p->buffer_len, 0);
      if(i <= 0) return LSOCK_ERROR;
      p->buffer_len += i;

//;;      LSock_Debug_WriteHex(buffer, p->buffer_len);
//;;      LSock_Debug_Write("\n");

      //  Check if there's a newline in the new data

      linelen = LSock_Buffer_Find(ls, 0 /*TODO: p->buffer_len-i*/, nl, 2);
      
//;;      LSock_Debug_WriteInt(linelen);
//;;      LSock_Debug_Write("\n");

    } else {
      return LSOCK_ERROR;
    }

  }

  //  Return the line if we've found one

  if(linelen > -1) {
    if(linelen > len) {
      if(buf) LSock_Buffer_Get(ls, buf, len);
      if(discard) {
        LSock_Buffer_Remove(ls, linelen+2);
      } else {
        LSock_Buffer_Remove(ls, len);
      }
      return linelen;
    } else {
      if(buf) LSock_Buffer_Get(ls, buf, linelen);
      LSock_Buffer_Remove(ls, linelen+2);
      return linelen;
    }
  }

  //  Otherwise error

  return LSOCK_ERROR;
}



//---- LSock_WriteLine - Writes data from a null terminated string

//  If this function returns an error you don't know how much data was actually sent.

int LSock_WriteLine(LSOCK ls, char *buf) {
  int i = LSock_WriteString(ls, buf);
  int j;
  char newline[2] = {0x0D, 0x0A};

  if(i == LSOCK_ERROR) return LSOCK_ERROR;

  j = LSock_WriteData(ls, newline, 2);

  if(j == LSOCK_ERROR) return LSOCK_ERROR;

  return i+j;
}



//---- END OF CODE -----------------------------------------------------------//

#endif