#include <solegmp.h>

#define stream_inp_is_aligned(strm) (!strm->offset)
#define stream_inp_set_aligned(strm) (strm->offset = 0)

#define strm_update_in(strm, read)		\
	do {					\
		strm->inp += read;		\
		strm->total_in += read;		\
		strm->avail_in -= read;		\
	} while(0)

#define strm_update_out(strm, written)			\
	do {						\
		strm->outp += written;			\
		strm->total_out += written;		\
		strm->avail_out -= written;		\
	} while(0)

/**
 * Internal Declarations
 */
void soleg_pass1_compose(mpz_t rop, mpz_t x, mpz_t y);
void soleg_pass1_decompose(mpz_t q, mpz_t rem, mpz_t op, mpz_t b3i);
void soleg_pass2_regroup(mpz_t rop, mpz_t x, mpz_t y, mpz_t i3);

void soleg_pass2_rev_regroup(blk256_t *, blk128_t *, blk256_t *, blk128_t *);
void soleg_pass1_rev_decompose(blk256_t *, blk256_t *, blk128_t *, blk128_t *);
void soleg_pass1_rev_compose(blk256_t *, blk256_t *, blk256_t *);

int soleg_encode_blk256(soleg_stream *);
int soleg_decode_blk256(soleg_stream *);

int soleg_stream_init(soleg_stream *strm)
{
	strm->flags = strm->total_in = strm->total_out = 0;

	mpz_inits(strm->i3, strm->b3i, strm->r, strm->q, NULL);
	mpz_setbit(strm->b3i, 128); /* 1 << 128 */
	mpz_setbit(strm->i3, 128); /* 1 << 128 */

	mpz_init2(strm->b1, 128);
	mpz_init2(strm->b2, 128);

	return SOLEG_OK;
}

#define stream_stat(strm)						\
	/* do {								\ */
	/* 	printf("strm: total_in=%lu, total_out=%lu, avail_in=%lu, avail_out=%lu, offset=%u, in=%p, out=%p\n", \ */
	/* 	       strm->total_in, strm->total_out, strm->avail_in, strm->avail_out, strm->offset, strm->inp, strm->outp); \ */
	/* } while(0) */

/**
 * Output a 256-bit block
 */
int soleg_put_blk256(soleg_stream *strm, mpz_t b)
{
	if (strm->avail_out < BLK256SIZ)
		return SOLEG_NO_BUF_SPC;

	/* gmp_printf("db: %Zd\n", b); */
	mpz_export(strm->outp,
		   NULL,	/* Discard number of bytes output */
		   1,		/* Order: Most significant byte first */
		   BLK256SIZ,
		   0,		/* Endian: native */
		   0,		/* Produce full words */
		   b);
	strm_update_out(strm, BLK256SIZ);

	stream_stat(strm);
	return SOLEG_OK;
}


/**
 * Rarely used
 */
int soleg_put_blk128(soleg_stream *strm, mpz_t b)
{
	if (strm->avail_out < BLK128SIZ)
		return SOLEG_NO_BUF_SPC;

	/* gmp_printf("r:  %Zd\n"); */
	mpz_export(strm->outp, NULL, 1, BLK128SIZ, 0, 0, b);
	strm_update_out(strm, BLK128SIZ);

	return SOLEG_OK;
}

/**
 * Encode the available bytes
 *
 * This should always consume all available input, to allow the user to specify
 * the next input buffer (if present).
 */
int soleg_encode(soleg_stream *strm)
{
	int ret;
	stream_stat(strm);

	if (!strm->avail_in)
		return SOLEG_OK;

	/* TODO Check avail_out > avail_in. We can actually compute how much
	 *      bigger this value should be */
	if (strm->avail_in > strm->avail_out)
		return SOLEG_NO_BUF_SPC;


	if (!stream_inp_is_aligned(strm)) {
		/**
		 * In case the given buffer was not aligned to 256 bits, realign on subsequent
		 * invocations of soleg_encode().
		 */
		size_t to_align = BLK256SIZ - strm->offset;
		if (strm->avail_in < to_align) {
			memcpy(strm->buf + strm->offset, strm->inp, strm->avail_in);
			strm->offset += strm->avail_in;
			strm_update_in(strm, strm->avail_in);
			return SOLEG_OK;
		}

		/* Can complete a double-block */
		memcpy(strm->buf + strm->offset, strm->inp, to_align);
		strm_update_in(strm, to_align);
		stream_inp_set_aligned(strm);

		mpz_import(strm->b1, 1, 1, BLK128SIZ, 0, 0, strm->buf);
		mpz_import(strm->b2, 1, 1, BLK128SIZ, 0, 0, strm->buf + BLK128SIZ);
		if ((ret = soleg_encode_blk256(strm)) != SOLEG_OK)
			return ret;
	}


	/* Function:
	 * void mpz_import (mpz_t rop, size_t count, int order, size_t size, int endian, size_t nails, const void *op) */

	/* The parameters specify the format of the data. count many words
	 * are read, each size bytes. order can be 1 for most significant
	 * word first or -1 for least significant first. Within each word
	 * endian can be 1 for most significant byte first, -1 for least
	 * significant first, or 0 for the native endianness of the host
	 * CPU. The most significant nails bits of each word are skipped,
	 * this can be 0 to use the full words. */

	/* There is no sign taken from the data, rop will simply be a
	 * positive integer. An application can handle any sign itself,
	 * and apply it for instance with mpz_neg. */

	/* Here's an example converting an array of unsigned long data,
	 * most significant element first, and host byte order within each
	 * value. */

	/*           unsigned long  a[20]; */
	/*           /\* Initialize z and a *\/ */
	/*           mpz_import (z, 20, 1, sizeof(a[0]), 0, 0, a); */

	/* This example assumes the full sizeof bytes are used for data in
	 * the given type, which is usually true, and certainly true for
	 * unsigned long everywhere we know of. However on Cray vector
	 * systems it may be noted that short and int are always stored in
	 * 8 bytes (and with sizeof indicating that) but use only 32 or 46
	 * bits. The nails feature can account for this, by passing for
	 * instance 8*sizeof(int)-INT_BIT.  */

	/* Main body - Read into 128-bit blocks */
	for (; strm->avail_in >= BLK256SIZ; ) {
		mpz_import(strm->b1, 1, 1, BLK128SIZ, 0, 0, strm->inp);
		strm_update_in(strm, BLK128SIZ);

		mpz_import(strm->b2, 1, 1, BLK128SIZ, 0, 0, strm->inp);
		strm_update_in(strm, BLK128SIZ);

		ret = soleg_encode_blk256(strm);
		if (ret != SOLEG_OK)
			return ret;
	}

	/* Read the remaining bytes */
	memcpy(strm->buf, strm->inp, strm->avail_in);
	strm->offset = strm->avail_in;
	strm_update_in(strm, strm->avail_in);

	return SOLEG_OK;
}

/* Function: void mpz_mul_2exp (mpz_t rop, const mpz_t op1, mp_bitcnt_t op2)

   Set rop to op1 times 2 raised to op2. This operation can also be defined as a left shift by op2 bits.
 */
int soleg_encode_blk256(soleg_stream *strm)
{
	mpz_t db;
	mpz_init(db);

	/* Pass 1 */
	soleg_pass1_compose(db, strm->b1, strm->b2);
	soleg_pass1_decompose(db, strm->r, db, strm->b3i);

	if (!mpz_sgn(strm->i3)) {
		soleg_put_blk128(strm, strm->r);

		mpz_set(strm->q, db);
	} else {
		/* Pass 2 */
		mpz_t res;
		mpz_init(res);

		/* soleg_pass2_regroup(&res, strm->q, strm->r, strm->i3); */
		soleg_pass2_regroup(res, strm->q, strm->r, strm->i3);
		soleg_put_blk256(strm, res);

		mpz_set(strm->q, db);
		mpz_clear(res);
	}

	/* Adjust counters */
	mpz_add_ui(strm->i3, strm->i3, 3);
	mpz_sub_ui(strm->b3i, strm->b3i, 3);

	mpz_clear(db);
	return SOLEG_OK;
}

/**
 * Instruct the given stream that it is finished. This takes care of finalising
 * the output and encoding EOF.
 */
void soleg_encode_end(soleg_stream *strm)
{
	stream_stat(strm);

	/* Writes the last two-three blocks */
	mpz_t db, res, eof;
	mpz_inits(db, res, eof, NULL);

	int have_even_blocks = strm->offset > BLK128SIZ;

	if (have_even_blocks) {
		soleg_encode_blk256(strm);

		/* Composed double EOF */
		mpz_setbit(eof, 256);
		mpz_setbit(eof, 129);

		soleg_pass1_decompose(db, strm->r, eof, strm->b3i);
	} else {
		mpz_setbit(eof, 128);
		soleg_pass1_compose(db, strm->b1, eof);
		soleg_pass1_decompose(db, strm->r, db, strm->b3i);
	}

	soleg_pass2_regroup(res, strm->q, strm->r, strm->i3);
	soleg_put_blk256(strm, res);

	/* Cheeky implementation of final Pass 2 */
	soleg_put_blk256(strm, db);

	mpz_clears(db, res, eof, NULL);
}

/**
 * Drop and release all resources associated with the given stream. This does
 * not free() the stream itself.
 */
void soleg_stream_free(soleg_stream *strm)
{
	if (!strm)
		return;
	mpz_clears(strm->i3, strm->b3i, strm->r, strm->q, strm->b1, strm->b2, NULL);
}

/**
 * Compose two 128-bit numbers
 */
void soleg_pass1_compose(mpz_t rop,
			 mpz_t x,
			 mpz_t y)
{
	mpz_mul_2exp(rop, x, 128);
	mpz_add(rop, rop, x);
	mpz_add(rop, rop, y);
}

/**
 * Decompose a 256-bit block by (B-3i)
 */
void soleg_pass1_decompose(mpz_t q,
			   mpz_t rem,
			   mpz_t op,
			   mpz_t b3i)
{
	mpz_tdiv_qr(q, rem, op, b3i);
}

/**
 * In pass 2 we first compose a 256-bit and 128-bit number then extract the relevant bits
 *
 * say x is a 256-bit number (a quotient from a div)
 *     y is a 128-bit number (a remainder)
 *
 * We do y(B+3i) + x
 */
void soleg_pass2_regroup(mpz_t rop,
			 mpz_t x,
			 mpz_t y,
			 mpz_t i3)
{
	mpz_set(rop, x);
	mpz_addmul(rop, y, i3);
}