/* See LICENSE file for copyright and license details. */
#include "liblss16.h"
#include <arpa/inet.h>
#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#if defined(__clang__)
# pragma clang diagnostic ignored "-Wunsafe-buffer-usage"
# pragma clang diagnostic ignored "-Wimplicit-fallthrough"
# pragma clang diagnostic ignored "-Wcast-align" /* clang generates false positives, missing in its analysis that
                                                  * `buf` is aligned due to use of malloc(3) and realloc(3) */
#endif


static const char *argv0 = "ppmtolss16";

static int eof = 0;
static int in_comment = 0;
static int maptype;
static char *buf;
static size_t buf_size;
static size_t len = 0;
static size_t width = 0;
static size_t height = 0;
static uint16_t maxvalue;
static int invert;


static void
writeall(char *data, size_t n)
{
	ssize_t r;

	while (n) {
		r = write(STDOUT_FILENO, data, n);
		if (r < 0) {
			if (errno == EINTR)
				continue;
			fprintf(stderr, "%s: write <stdout>: %s\n", argv0, strerror(errno));
			exit(2);
		}
		data += (size_t)r;
		n -= (size_t)r;
	}
}


#if defined(__GNUC__)
__attribute__((__noreturn__))
#endif
static void
not_pnm(void)
{
	fprintf(stderr, "%s: <stdin> is no a P1-P6 Portable Anymap file\n", argv0);
	exit(1);
}


static void
fill_buffer(void)
{
	ssize_t r;
	while (len < buf_size) {
		r = read(STDIN_FILENO, &buf[len], buf_size - len);
		if (r <= 0) {
			if (!r) {
				eof = 1;
				break;
			}
			if (errno == EINTR)
				continue;
			fprintf(stderr, "%s: read <stdin>: %s\n", argv0, strerror(errno));
			exit(3);
		}
		len += (size_t)r;
	}
}


static size_t
skip_whitespace(size_t p, int break_on_newline)
{
	for (;;) {
		if (in_comment) {
			while (p < len && buf[p] != '\n')
				p++;
			if (p < len) {
				if (break_on_newline)
					return ++p;
				goto toggle_comment;
			}
		} else {
			while (p < len && isspace(buf[p])) {
				if (break_on_newline && buf[p] == '\n')
					return ++p;
				p++;
			}
			if (p < len) {
				if (buf[p] != '#')
					break;
			toggle_comment:
				in_comment ^= 1;
				p++;
				continue;
			}
		}
		if (eof)
			not_pnm();
		len = 0;
		fill_buffer();
	}
	return p;
}


static size_t
read_dimension(size_t p, size_t *value)
{
	size_t digit;
	for (;;) {
		if (p == len) {
			if (eof)
				not_pnm();
			len = 0;
			fill_buffer();
		}
		if (isspace(buf[p]))
			break;
		if (!isdigit(buf[p]))
			not_pnm();
		digit = (size_t)(buf[p++] & 15);
		if (*value > (UINT16_MAX - digit) / 10U) {
			*value = SIZE_MAX;
			break;
		}
		*value = *value * 10U + digit;
	}
	return p;
}


static void
read_header(void)
{
	size_t p;

	fill_buffer();
	if (len < 3U || buf[0] != 'P' || buf[1] < '1' || buf[1] > '6' || !isspace(buf[2]))
		not_pnm();
	maptype = buf[1] - '0';

	p = skip_whitespace(3U, 0);
	p = read_dimension(p, &width);
	if (!width)
		not_pnm();
	if (width > (size_t)UINT16_MAX) {
		fprintf(stderr, "%s: image from <stdin> is too wide for LSS16 encoding\n", argv0);
		exit(2);
	}
	p = skip_whitespace(p, 0);
	p = read_dimension(p, &height);
	if (!height)
		not_pnm();
	if (height > (size_t)(UINT16_MAX >> 1)) {
		fprintf(stderr, "%s: image from <stdin> is too tall for LSS16 encoding\n", argv0);
		exit(2);
	}

	if (maptype == 1 || maptype == 4) {
		maxvalue = 1U;
		invert = 1;
	} else {
		size_t t;
		p = skip_whitespace(p, 0);
		p = read_dimension(p, &t);
		if (!t || t > UINT16_MAX)
			not_pnm();
		maxvalue = (uint16_t)t;
		invert = 0;
	}

	p = skip_whitespace(p, maptype > 3);

	memmove(&buf[0], &buf[p], len -= p);
}


static int
read_image_raw(size_t bytes_per_pixel, int decode_big_endian)
{
	size_t bytes = width * height;
	size_t i;
	int extra;

	if (bytes > SIZE_MAX / bytes_per_pixel) {
		fprintf(stderr, "%s: image too large to for the address space\n", argv0);
		return 3;
	}
	bytes *= bytes_per_pixel;

	if (buf_size < bytes + 1U) {
		buf = realloc(buf, buf_size = bytes + 1U);
		if (!buf) {
			fprintf(stderr, "%s: out of memory\n", argv0);
			exit(3);
		}
	}

	if (len < bytes + 1U) {
		if (!eof)
			fill_buffer();
		if (len < bytes) {
			fprintf(stderr, "%s: input file truncated\n", argv0);
			exit(1);
		}
	}

	if (decode_big_endian)
		for (i = 0; i < bytes; i += 2)
			*(uint16_t *)&buf[i] = ntohs(*(const uint16_t *)&buf[i]);

	extra = len > bytes;
	len = bytes;
	return extra;

}


static int
read_image_plain(size_t words_per_pixel, size_t bytes_per_word)
{
	size_t r, w, bytes, words = width * height;
	uint16_t value = 0, digit;
	int extra;

	if (words > SIZE_MAX / words_per_pixel) {
		fprintf(stderr, "%s: image too large to for the address space\n", argv0);
		return 3;
	}
	words *= words_per_pixel;
	if (words > (SIZE_MAX - 16U) / bytes_per_word) {
		fprintf(stderr, "%s: image too large to for the address space\n", argv0);
		return 3;
	}
	bytes = bytes_per_word * words;

	if (buf_size < bytes + 16U) {
		buf = realloc(buf, buf_size = bytes + 16U);
		if (!buf) {
			fprintf(stderr, "%s: out of memory\n", argv0);
			exit(3);
		}
	}

	if (!len) {
		if (eof) {
		truncated:
			fprintf(stderr, "%s: input file truncated\n", argv0);
			exit(1);
		}
		fill_buffer();
		if (!len)
			goto truncated;
	}

	r = w = 0;
	while (words--) {
		if (!isdigit(buf[r]))
			not_pnm();
		value = 0;
		do {
			digit = (uint16_t)(buf[r] & 15);
			if (value > (UINT16_MAX - digit) / 10U)
				not_pnm();
			value = (uint16_t)(value * 10U + digit);
			if (++r == len) {
				r = len = w;
				if (eof)
					break;
				fill_buffer();
				if (r == len)
					break;
			}
		} while (isdigit(buf[r]));
		if (bytes_per_word == 1U) {
			if (value > UINT8_MAX)
				not_pnm();
			buf[w++] = (char)(uint8_t)value;
		} else {
			*(uint16_t *)&buf[w] = value;
			w += 2U;
		}
		if (r == len) {
			if (words)
				goto truncated;
			break;
		}

		while (isspace(buf[r])) {
			if (++r == len) {
				r = len = w;
				if (eof)
					break;
				fill_buffer();
				if (r == len)
					break;
			}
		}
		if (r == len) {
			if (words)
				goto truncated;
			break;
		}
	}

	extra = r < len;
	len = w;
	return extra;
}


static int
read_image_bits(void)
{
	size_t bytes = width * height;
	size_t input = ((width + 7U) / 8U) * height;
	char *inbuf;
	size_t r, w, x, y;
	int extra = 0;
	ssize_t rd;

	if (buf_size < bytes) {
		buf = realloc(buf, buf_size = bytes);
		if (!buf) {
			fprintf(stderr, "%s: out of memory\n", argv0);
			exit(3);
		}
	}

	inbuf = malloc(input);
	if (!inbuf) {
		fprintf(stderr, "%s: out of memory\n", argv0);
		exit(3);
	}

	if (len > input) {
		extra = 1;
		len = input;
	}

	memcpy(inbuf, buf, len);
	if (len < input && eof) {
	truncated:
		fprintf(stderr, "%s: input file truncated\n", argv0);
		exit(1);
	}
	while (len < input) {
		rd = read(STDIN_FILENO, &inbuf[len], input - len);
		if (rd <= 0) {
			if (!rd)
				goto truncated;
			if (errno == EINTR)
				continue;
			fprintf(stderr, "%s: read <stdin>: %s\n", argv0, strerror(errno));
			exit(3);
		}
		len += (size_t)rd;
	}

	r = w = 0;
	for (y = 0; y < height; y++) {
		for (x = 0; x + 8U < width; x += 8) {
			buf[w++] = (char)((inbuf[r] >> 7) & 1);
			buf[w++] = (char)((inbuf[r] >> 6) & 1);
			buf[w++] = (char)((inbuf[r] >> 5) & 1);
			buf[w++] = (char)((inbuf[r] >> 4) & 1);
			buf[w++] = (char)((inbuf[r] >> 3) & 1);
			buf[w++] = (char)((inbuf[r] >> 2) & 1);
			buf[w++] = (char)((inbuf[r] >> 1) & 1);
			buf[w++] = (char)((inbuf[r] >> 0) & 1);
			r++;
		}
		if (width & 7U) {
			inbuf[r] >>= 8U - (width & 7U);
			switch (width & 7U) {
			case 7:
				buf[w++] = (char)((inbuf[r] >> 6) & 1);
				/* fall through */
			case 6:
				buf[w++] = (char)((inbuf[r] >> 5) & 1);
				/* fall through */
			case 5:
				buf[w++] = (char)((inbuf[r] >> 4) & 1);
				/* fall through */
			case 4:
				buf[w++] = (char)((inbuf[r] >> 3) & 1);
				/* fall through */
			case 3:
				buf[w++] = (char)((inbuf[r] >> 2) & 1);
				/* fall through */
			case 2:
				buf[w++] = (char)((inbuf[r] >> 1) & 1);
				/* fall through */
			case 1:
				buf[w++] = (char)((inbuf[r] >> 0) & 1);
				break;
			default:
				abort();
			}
			r++;
		}
	}

	free(inbuf);

	if (!extra && !eof) {
	read_extra_again:
		rd = read(STDIN_FILENO, &(char){0}, 1U);
		if (rd < 0) {
			if (errno == EINTR)
				goto read_extra_again;
			fprintf(stderr, "%s: read <stdin>: %s\n", argv0, strerror(errno));
			exit(3);
		} else if (rd > 0) {
			extra = 1;
		} else {
			eof = 1;
		}
	}

	len = w;
	return extra;
}


int
main(int argc, char *argv[])
{
	struct liblss16_header header;
	struct liblss16_encoder encoder;
	enum liblss16_encode_error error;
	enum liblss16_encode_state state;
	uint16_t half_maxvalue;
	size_t i, off, consumed, written;
	char wbuf[1024];
	int extra, have_stdout = 0;

	/* cmdline syntax is inherited from the SYSLINUX implementation,
	 * it's bad, but it's the way it is, and the way it will remain,
	 * however removed palette specification operands */

	if (argc) {
		argv0 = *argv++;
		argc--;
		if (argc) {
			fprintf(stderr, "%s: Unknown option: %s\n", argv0, *argv);
			return 3;
		}
	}

	buf_size = 1024U;
	buf = malloc(buf_size);
	if (!buf) {
		fprintf(stderr, "%s: out of memory\n", argv0);
		return 3;
	}

	read_header();

	switch (maptype) {
	case 1:
		extra = read_image_plain(1U, 1U);
		break;
	case 2:
		extra = read_image_plain(1U, maxvalue > 255 ? 2U : 1U);
		break;
	case 3:
		extra = read_image_plain(3U, maxvalue > 255 ? 2U : 1U);
		break;
	case 4:
		extra = read_image_bits();
		break;
	case 5:
		extra = read_image_raw(maxvalue > 255 ? 2U : 1U, maxvalue > 255);
		break;
	case 6:
		extra = read_image_raw(maxvalue > 255 ? 6U : 3U, maxvalue > 255);
		break;
	default:
#if defined(__GNUC__)
		__builtin_unreachable();
#else
		abort();
#endif
	}

	if (extra)
		fprintf(stderr, "%s: <stdin> contained extraneous data\n", argv0);

	if (maptype > 3)
		maptype -= 3;

	half_maxvalue = maxvalue >> 1;
	if (maxvalue > 255) {
		uint16_t value;
		len >>= 1;
		for (i = 0; i < len; i++) {
			value = *(uint16_t *)&buf[i << 1];
			if (value > maxvalue)
				not_pnm();
			buf[i] = (char)(uint8_t)(((uint32_t)value * 63U + half_maxvalue) / maxvalue);
		}
	} else if (maxvalue == 1) {
		for (i = 0; i < len; i++)
			if (buf[i] & 0xFE)
				not_pnm();
		if (maptype == 3) {
			for (i = 0; i < len; i++)
				buf[i] = (char)(uint8_t)((uint8_t)buf[i] * 63U);
		} else {
			maptype = 1;
		}
	} else if (maxvalue != 63) {
		for (i = 0; i < len; i++) {
			if ((uint8_t)buf[i] > maxvalue)
				not_pnm();
			buf[i] = (char)(uint8_t)(((uint8_t)buf[i] * 63U + half_maxvalue) / maxvalue);
		}
	}

	memset(&header, 0, sizeof(header));
	header.width = (uint16_t)width;
	header.height = (uint16_t)height;

	if (maptype == 1) {
		header.colour_map[!invert].r = 63U;
		header.colour_map[!invert].g = 63U;
		header.colour_map[!invert].b = 63U;
	} else if (maptype == 2) {
		/* TODO convert to 16 colours -- greyscale values [0, 63], quantize to 16 colours and map Y to (Y, Y, Y) */
	} else {
		/* TODO convert to 16 colours -- [0, 63]-triplet values, quantize to 16 colours */
	}

	liblss16_optimise(&header, (uint8_t *)buf);
	if (liblss16_encoder_init(&encoder, &header, 1, &error)) {
		fprintf(stderr, "%s:liblss16_encoder_init: %s\n", argv0, liblss16_encode_strerror(error));
		exit(3);
	}
	off = 0;
	do {
		state = liblss16_encode_from_colour_index(&encoder, wbuf, sizeof(buf), &written,
		                                          (const uint8_t *)&buf[off], len - off, &consumed);
		off += consumed;
		writeall(wbuf, written);
	} while (state == LIBLSS16_ENCODE_RUNNING);
	if (off != len)
		abort();

	free(buf);

	while (close(STDOUT_FILENO)) {
		if (have_stdout && errno == EBADF)
			break;
		if (errno == EINTR) {
			have_stdout = 1;
		} else {
			fprintf(stderr, "%s: write <stdout>: %s\n", argv0, strerror(errno));
			return 3;
		}
	}
	return 0;
}