#include "postgres.h"
#include "fmgr.h"
#include "libpq/pqformat.h"
#include "utils/builtins.h"

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

typedef int32 ssn_t;

#define SSNGetDatum(x)	 Int32GetDatum(x)
#define DatumGetSSN(x)	 DatumGetInt32(x)
#define PG_GETARG_SSN(n) DatumGetSSN(PG_GETARG_DATUM(n))
#define PG_RETURN_SSN(x) return SSNGetDatum(x)

#define DIGITTOINT(n)	(((n) >= '0' && (n) <= '9') ? (int32) ((n) - '0') : 0)

Datum ssn_in(PG_FUNCTION_ARGS);
Datum ssn_out(PG_FUNCTION_ARGS);
Datum ssn_to_text(PG_FUNCTION_ARGS);
Datum text_to_ssn(PG_FUNCTION_ARGS);
Datum ssn_send(PG_FUNCTION_ARGS);
Datum ssn_recv(PG_FUNCTION_ARGS);
Datum ssn_lt(PG_FUNCTION_ARGS);
Datum ssn_le(PG_FUNCTION_ARGS);
Datum ssn_eq(PG_FUNCTION_ARGS);
Datum ssn_ne(PG_FUNCTION_ARGS);
Datum ssn_ge(PG_FUNCTION_ARGS);
Datum ssn_gt(PG_FUNCTION_ARGS);
Datum ssn_cmp(PG_FUNCTION_ARGS);

static ssn_t cstring_to_ssn(char *ssn_string);
static char *ssn_to_cstring(ssn_t ssn);
static bool ssn_is_valid(int32 area, int32 group, int32 serial);



/* generic input/output functions */
PG_FUNCTION_INFO_V1(ssn_in);
Datum 
ssn_in(PG_FUNCTION_ARGS)
{
	ssn_t	result;
	
	char   *ssn_str = PG_GETARG_CSTRING(0);
	result = cstring_to_ssn(ssn_str);
	
	PG_RETURN_SSN(result);
}

PG_FUNCTION_INFO_V1(ssn_out);
Datum 
ssn_out(PG_FUNCTION_ARGS)
{
	ssn_t	packed_ssn;
	char   *ssn_string;
	
	packed_ssn = PG_GETARG_SSN(0);
	ssn_string = ssn_to_cstring(packed_ssn);
	
	PG_RETURN_CSTRING(ssn_string);	
}

/* type to/from text conversion routines */
PG_FUNCTION_INFO_V1(ssn_to_text);
Datum 
ssn_to_text(PG_FUNCTION_ARGS) 
{
	char	*ssn_string;
	text	*ssn_text;
	
	ssn_t	packed_ssn =  PG_GETARG_SSN(0);
	
	ssn_string = ssn_to_cstring(packed_ssn);
	ssn_text = DatumGetTextP(DirectFunctionCall1(textin, CStringGetDatum(ssn_string)));
	
	PG_RETURN_TEXT_P(ssn_text);	
}

PG_FUNCTION_INFO_V1(text_to_ssn);
Datum 
text_to_ssn(PG_FUNCTION_ARGS) 
{
	text  *ssn_text = PG_GETARG_TEXT_P(0);
	char  *ssn_str = DatumGetCString(DirectFunctionCall1(textout, PointerGetDatum(ssn_text)));
	ssn_t ssn = cstring_to_ssn(ssn_str);
	
	PG_RETURN_SSN(ssn);
}

/* Functions to convert to/from bytea */
PG_FUNCTION_INFO_V1(ssn_send);
Datum 
ssn_send(PG_FUNCTION_ARGS)
{
	StringInfoData buffer;
	ssn_t ssn = PG_GETARG_SSN(0);
	
	pq_begintypsend(&buffer);
	pq_sendint(&buffer, (int32)ssn, 4);
	PG_RETURN_BYTEA_P(pq_endtypsend(&buffer));
}

PG_FUNCTION_INFO_V1(ssn_recv);
Datum ssn_recv(PG_FUNCTION_ARGS)
{
	ssn_t	ssn;
	StringInfo	buffer = (StringInfo) PG_GETARG_POINTER(0);
	
	ssn = pq_getmsgint(buffer, 4);
	PG_RETURN_SSN(ssn);
}

/* functions to support btree opclass */
PG_FUNCTION_INFO_V1(ssn_lt);
Datum
ssn_lt(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	PG_RETURN_BOOL(a < b);
}

PG_FUNCTION_INFO_V1(ssn_le);
Datum
ssn_le(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	PG_RETURN_BOOL (a <= b);
}

PG_FUNCTION_INFO_V1(ssn_eq);
Datum
ssn_eq(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	PG_RETURN_BOOL(a == b);
}

PG_FUNCTION_INFO_V1(ssn_ne);
Datum
ssn_ne(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	PG_RETURN_BOOL(a != b);
}

PG_FUNCTION_INFO_V1(ssn_ge);
Datum
ssn_ge(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	PG_RETURN_BOOL(a >= b);
}

PG_FUNCTION_INFO_V1(ssn_gt);
Datum
ssn_gt(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	PG_RETURN_BOOL(a > b);
}

PG_FUNCTION_INFO_V1(ssn_cmp);
Datum
ssn_cmp(PG_FUNCTION_ARGS)
{
	ssn_t a = PG_GETARG_SSN(0);
	ssn_t b = PG_GETARG_SSN(1);
	
	PG_RETURN_INT32(a - b);
}


/* 
 * Convert a cstring to an SSN, validating the input.
 * Input in forms AAA-BB-CCCC or AAABBCCCC is accepted.
 */ 
static ssn_t 
cstring_to_ssn(char *ssn_str)
{
	char   *c;
	ssn_t 	result;
	char 	dashes[] = {3, 6};
	int		dashes_no = 0;
	int		ndigits = 0;
	int32	area,
			group,
			serial;
	
	area = group = serial = 0;
	
	for (c = ssn_str; *c != 0; c++) 
	{
		if (isdigit(*c)) 
		{
			if (ndigits < 3) area = area * 10 + DIGITTOINT(*c);
			else if (ndigits < 5) group = group * 10 + DIGITTOINT(*c);
			else if (ndigits < 9) serial = serial * 10 + DIGITTOINT(*c);
			ndigits++;
		}
		else if (*c == '-') 
		{
			int pos = c - ssn_str;
			if (dashes_no < 2 && pos == dashes[dashes_no])
				dashes_no++;
			else
				ereport(ERROR,
						(errmsg("invalid format of input data %s", ssn_str),
						 errhint("Valid formats are: AAA-BB-CCCC or AAABBCCCC")));
		}
		else
			ereport(ERROR,
					(errmsg("unexpected character '%c' in input data %s", *c, ssn_str),
					 errhint("Valid SSN consists of digits and optional dashes")));
		
	}
	if (ndigits != 9)
		ereport(ERROR,
				(errmsg("invalid number of digits (%d) in input data %s", ndigits, ssn_str),
				 errhint("Valid SSN consists of 9 digits")));
				
	if (!ssn_is_valid(area, group, serial))
		ereport(ERROR, 
				(errmsg("SSN number %s is invalid", ssn_str)));

	result = serial + (group << 14) + (area << 21);
	PG_RETURN_SSN(result);

}

/* Convert the internal representation to the AAA-BB-CCCC output string */
static char *
ssn_to_cstring(ssn_t ssn)
{
	int32	area,
			group,
			serial;
	int32	ndigits;
	char   *ssn_str = palloc(12);
			
	if (ssn < 0)
		/* Error out */;
	
	serial  = ssn & ((1 << 14) - 1); /* first 14 bits */
	group = (ssn >> 14) & 127; /* next 7 bits */ 
	area = (ssn >> 21);  /* last 10 bits (highest one is always 0) */
	
	if (!ssn_is_valid(area, group, serial))
		ereport(ERROR, 
				(errmsg("SSN number %s is invalid", ssn_str)));
	
	if ((ndigits = snprintf(ssn_str, 12, "%03d-%02d-%04d", area, group, serial)) != 11)
		ereport(ERROR,
				(errmsg("invalid size (%d) of in input data %s", ndigits-2, ssn_str),
				 errhint("Valid SSN consists of 9 digits")));
		
	if (ssn_str[3] != '-' || ssn_str[6] != '-')
		ereport(ERROR,
				(errmsg("invalid format of input data %s", ssn_str),
				 errhint("Valid formats are: AAA-BB-CCCC or AAABBCCCC")));
	

	return ssn_str;
}

/* 
 * Check whether 3 components of the SSN are valid. Invalid SSNs are:
 * - containing 0 in one of the components.
 * - those with are code equal to 666 or greater or equal to 900.
 * Note: Area is not validated against area number groups due to 
 * SSN randomization process, which will take place after June 25th, 2011.
 */
static bool 
ssn_is_valid(int32 area, int32 group, int32 serial)
{
	if (area == 0 || group == 0 || serial == 0)
		return false;
	
	if (area == 666 || area >= 900)
		return false;
		
	return true;
}