#define _XOPEN_SOURCE

#include "tabular.h"
#include "util.h"

#include <sys/ioctl.h>
#include <wchar.h>
#include <errno.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <stdint.h>
#include <stddef.h>

#define BIT(i) (1U << (i))

struct col_state {
	size_t width;
	size_t written;
	size_t maxlen;
	size_t lpad, rpad;
	bool hidden;
};

struct fmt_state {
	struct col_state *columns;
	size_t column_count;

	size_t hseplen;
	size_t vseplen;
	size_t xseplen;
	size_t mseplen;

	size_t line_limit;
	size_t row_limit;

	bool header_line;
};

static void calc_word_aware(struct tabular_entry *entry, size_t maxwidth,
	size_t *offset, size_t *width, size_t *lines);
static size_t recalc_col_word_aware(const struct tabular_cfg *cfg,
	struct tabular_row *rows, size_t limit,
	size_t col, size_t maxwidth);
static size_t calc_row_width(struct fmt_state *fmt);
static size_t calc_output_lines(const struct tabular_cfg *cfg,
	struct fmt_state *fmt, struct tabular_row *rows, size_t limit);
static bool recalc_cols(const struct tabular_cfg *cfg,
	struct tabular_stats *stats, struct fmt_state *fmt,
	struct tabular_row *rows, size_t limit);

static int calc_params_row(const struct tabular_cfg *cfg,
	struct tabular_stats *stats, struct tabular_row *rows,
	struct tabular_row *row, struct fmt_state *fmt, size_t limit);
static int calc_params(const struct tabular_cfg *cfg,
	struct tabular_stats *stat, struct tabular_row **rows,
	struct fmt_state *fmt);

static void output_header(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stat, const struct fmt_state *fmt);
static void output_row(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stat, const struct tabular_row *row,
	struct fmt_state *fmt);
static int output_rows(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stats, struct tabular_row *rows,
	struct fmt_state *fmt);

const char *
calc_word_aware_line(const char *str, size_t maxwidth,
	size_t *out_offset, size_t *out_width)
{
	const char *sep, *nstr;
	size_t width, nwidth;
	size_t offset, wlen;

	if (!str) {
		if (out_offset) *out_offset = 0;
		if (out_width) *out_width = 0;
		return NULL;
	}

	offset = strspn(str, " \v\t\r\n");
	str += offset;

	nstr = str;
	nwidth = width = 0;
	while (nwidth <= maxwidth) {
		width = nwidth;
		str = nstr;
		if (!str) break;
		sep = strchr(str, ' ');
		if (!sep) {
			wlen = u8strlen(str);
			nstr = NULL;
		} else {
			wlen = u8strnlen(str, (size_t) (sep - str));
			nstr = sep + 1;
		}
		nwidth = width + (width > 0) + wlen;
	}

	if (out_offset) *out_offset = offset;
	if (out_width) {
		/* single word > maxwidth */
		if (!width && nwidth) {
			*out_width = maxwidth;
			str = nstr; /* skip single word */
		} else {
			*out_width = width;
		}
	}

	return str;
}

void
calc_word_aware(struct tabular_entry *entry, size_t maxwidth,
	size_t *out_offset, size_t *out_width, size_t *out_lines)
{
	const char *str;
	size_t width, mwidth;
	size_t lines;

	if (out_offset) *out_offset = 0;

	lines = 0;
	mwidth = 0;
	str = entry->str;
	while (str) {
		str = calc_word_aware_line(str, maxwidth,
			str == entry->str ? out_offset : NULL, &width);
		lines++;
		mwidth = MAX(mwidth, width);
	}

	if (out_width) *out_width = mwidth;
	if (out_lines) *out_lines = lines;
}

size_t
recalc_col_word_aware(const struct tabular_cfg *cfg,
	struct tabular_row *rows, size_t limit, size_t col,
	size_t maxwidth)
{
	size_t off, wwidth, max_wwidth, rowcnt;
	struct tabular_row *row;

	max_wwidth = cfg->columns[col].minwidth;

	rowcnt = 0;
	for (row = rows; row && rowcnt < limit; row = row->next) {
		calc_word_aware(&row->entries[col],
			maxwidth, &off, &wwidth, NULL);
		max_wwidth = MAX(max_wwidth, wwidth);
		rowcnt += 1;
	}

	return max_wwidth;
}

size_t
calc_row_width(struct fmt_state *fmt)
{
	size_t sum, i;

	sum = 0;
	for (i = 0; i < fmt->column_count; i++) {
		if (fmt->columns[i].hidden) continue;
		sum += fmt->columns[i].width + (sum ? fmt->hseplen : 0);
	}

	return sum;
}

size_t
calc_output_lines(const struct tabular_cfg *cfg, struct fmt_state *fmt,
	struct tabular_row *rows, size_t limit)
{
	struct tabular_row *row;
	size_t row_index, row_lines;
	size_t entry_lines;
	size_t lines, i;
	size_t width;

	lines = 0;
	row_index = 0;
	for (row = rows; row; row = row->next) {
		if (row_index == limit) break;

		row_lines = 0;
		for (i = 0; i < cfg->column_count; i++) {
			if (fmt->columns[i].hidden) continue;
			width = fmt->columns[i].width
				- fmt->columns[i].lpad - fmt->columns[i].rpad;
			switch (cfg->columns[i].strategy) {
			case TABULAR_TRUNC:
				entry_lines = 1;
				break;
			case TABULAR_WRAP:
				entry_lines = CEILDIV(row->entries[i].ulen, width);
				break;
			case TABULAR_WRAP_WORDAWARE:
				calc_word_aware(&row->entries[i],
					width, NULL, NULL, &entry_lines);
				break;
			}
			row_lines = MAX(row_lines, entry_lines);
		}
		lines += row_lines;

		row_index += 1;
	}

	return lines;
}

bool
recalc_cols(const struct tabular_cfg *cfg, struct tabular_stats *stats,
	struct fmt_state *fmt, struct tabular_row *rows, size_t limit)
{
	size_t width, fullwidth, remaining;
	ssize_t i;

	fullwidth = cfg->outw - cfg->lpad - cfg->rpad;

	/* reset widths to minimum requirement */
	for (i = 0; i < cfg->column_count; i++) {
		if (fmt->columns[i].hidden) continue;
		if (cfg->columns[i].squashable) {
			fmt->columns[i].width = cfg->columns[i].minwidth;
		} else {
			width = MIN(fmt->columns[i].maxlen,
				cfg->columns[i].maxwidth - cfg->columns[i].lpad
				- cfg->columns[i].rpad);
			if (cfg->columns[i].strategy == TABULAR_WRAP_WORDAWARE)
				width = recalc_col_word_aware(cfg, rows,
					limit, (size_t) i, width);
			fmt->columns[i].width = width + cfg->columns[i].lpad
				+ cfg->columns[i].rpad;
		}
	}

	/* could not fit all necessary columns at minimum width */
	width = calc_row_width(fmt);
	while (width > fullwidth) {
		/* remove non-essential columns */
		for (i = ((ssize_t) cfg->column_count) - 1; i >= 0; i--) {
			if (!cfg->columns[i].essential
					&& !fmt->columns[i].hidden) {
				fmt->columns[i].hidden = true;
				break;
			}
		}

		/* failed to remove any more */
		if (i < 0) return false;

		stats->cols_truncated = true;
		width = calc_row_width(fmt);
	}

	/* redistribute excess width to columns, left to right */
	remaining = fullwidth - width;
	for (i = 0; remaining > 0 && i < cfg->column_count; i++) {
		if (fmt->columns[i].hidden) continue;
		if (!cfg->columns[i].squashable) continue;
		width = MIN(fmt->columns[i].maxlen, cfg->columns[i].maxwidth
			- fmt->columns[i].lpad - fmt->columns[i].rpad);
		width = MIN(width, fmt->columns[i].width + remaining
			- fmt->columns[i].lpad - fmt->columns[i].rpad);
		if (cfg->columns[i].strategy == TABULAR_WRAP_WORDAWARE)
			width = recalc_col_word_aware(cfg, rows,
				limit, (size_t) i, width);
		width = MAX(width + fmt->columns[i].lpad + fmt->columns[i].rpad,
			fmt->columns[i].width);
		remaining -= width - fmt->columns[i].width;
		fmt->columns[i].width = width;
	}

	return true;
}

int
calc_params_row(const struct tabular_cfg *cfg, struct tabular_stats *stats,
	struct tabular_row *rows, struct tabular_row *row,
	struct fmt_state *fmt, size_t limit)
{
	struct colinfo *copy;
	size_t i, height;
	int rc;

	/* make copy of current width in case the next row does not fit */
	copy = cfg->allocator->alloc(cfg->allocator,
		sizeof(struct col_state) * cfg->column_count, &rc);
	if (!copy) return -rc;
	memcpy(copy, fmt->columns, sizeof(struct col_state) * cfg->column_count);

	/* unhide cols */
	for (i = 0; i < cfg->column_count; i++) {
		tabular_load_row_entry_hidden(cfg, row, i);
		if (fmt->columns[i].hidden) {
			fmt->columns[i].hidden =
				(row->entries[i].flags & BIT(TABULAR_ENTRY_HIDDEN));
		}
		if (fmt->columns[i].hidden) continue;
	}

	/* update maxlen for visible cols */
	for (i = 0; i < cfg->column_count; i++) {
		tabular_load_row_entry_str(cfg, row, i);
		fmt->columns[i].maxlen = MAX(fmt->columns[i].maxlen,
			row->entries[i].ulen);
	}

	/* recalc column sizes to fit new column */
	if (!recalc_cols(cfg, stats, fmt, rows, limit))
		goto undo;

	height = calc_output_lines(cfg, fmt, rows, limit);

	/* check the line limit if applicable */
	if (fmt->line_limit > 0 && height > fmt->line_limit)
		goto undo;

	cfg->allocator->free(cfg->allocator, copy);

	return 0;

undo:
	memcpy(fmt->columns, copy, sizeof(struct col_state) * fmt->column_count);
	cfg->allocator->free(cfg->allocator, copy);

	return 1;
}

int
calc_params(const struct tabular_cfg *cfg, struct tabular_stats *stats,
	struct tabular_row **rows, struct fmt_state *fmt)
{
	struct tabular_row **row;
	ssize_t first, last;
	size_t i;
	int rc;

	fmt->header_line = ((cfg->vsep != NULL && *cfg->vsep)
		&& (cfg->xsep != NULL && *cfg->xsep));

	fmt->line_limit = 0;
	if (cfg->fit_rows) {
		fmt->line_limit = MAX(0, cfg->outh
			- cfg->skip_lines - 1 - fmt->header_line);
	}

	fmt->hseplen = u8strlen(cfg->hsep);
	fmt->vseplen = u8strlen(cfg->vsep);
	fmt->xseplen = u8strlen(cfg->xsep);
	fmt->mseplen = MAX(fmt->hseplen, fmt->xseplen);

	fmt->column_count = cfg->column_count;
	fmt->columns = cfg->allocator->alloc(cfg->allocator,
		fmt->column_count * sizeof(struct col_state), &rc);
	if (!fmt->columns) return -rc;

	for (i = 0; i < cfg->column_count; i++) {
		fmt->columns[i].lpad = cfg->columns[i].lpad;
		fmt->columns[i].rpad = cfg->columns[i].rpad;
		fmt->columns[i].hidden = true;
		fmt->columns[i].width = cfg->columns[i].minwidth;
		fmt->columns[i].maxlen = u8strlen(cfg->columns[i].name)
			+ fmt->columns[i].lpad + fmt->columns[i].rpad;
		if (fmt->columns[i].maxlen > cfg->columns[i].minwidth)
			return 1;
		fmt->columns[i].written = 0;
	}

	fmt->row_limit = 0;

	if (!*rows && cfg->row_gen) *rows = cfg->row_gen(&cfg->user);
	if (!*rows) return 0;

	if (!recalc_cols(cfg, stats, fmt, *rows, 0)) {
		stats->cols_truncated = true;
		stats->rows_truncated = true;
		return 0;
	}

	for (row = rows; ; row = &(*row)->next) {
		if (!*row && cfg->row_gen) *row = cfg->row_gen(&cfg->user);
		if (!*row) break;
		rc = calc_params_row(cfg, stats, *rows,
			*row, fmt, fmt->row_limit + 1);
		if (rc < 0) return rc;
		if (rc > 0) {
			stats->rows_truncated = true;
			break;
		}
		fmt->row_limit++;
	}

	first = last = -1;
	for (i = 0; i < cfg->column_count; i++) {
		if (fmt->columns[i].hidden) continue;
		if (first < 0) first = (ssize_t) i;
		last = (ssize_t) i;
	}

	if (first >= 0 && last >= 0) {
		fmt->columns[first].lpad += cfg->lpad;
		fmt->columns[first].width += cfg->lpad;
		fmt->columns[last].rpad += cfg->rpad;
		fmt->columns[last].width += cfg->rpad;
	}

	return 0;
}

void
output_header(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stats, const struct fmt_state *fmt)
{
	size_t i, k, width;
	bool dirty, first;

	first = true;
	for (i = 0; i < cfg->column_count; i++) {
		if (fmt->columns[i].hidden) continue;
		if (!first) {
			dirty = cfg->print_style(file, cfg, NULL, NULL);
			fprintf(file, "%s%*.s", cfg->hsep,
				(int) (fmt->mseplen - fmt->hseplen), "");
			if (dirty) fprintf(file, "\x1b[0m");
		}
		first = false;

		width = fmt->columns[i].width
			- fmt->columns[i].lpad - fmt->columns[i].rpad;

		dirty = cfg->print_style(file, cfg,
			NULL, &cfg->columns[i]);
		print_pad(file, fmt->columns[i].lpad);
		print_left(file, cfg->columns[i].name, width, width);
		print_pad(file, fmt->columns[i].rpad);
		if (dirty) fprintf(file, "\x1b[0m");
	}

	fprintf(file, "\n");
	stats->lines_used += 1;

	if (fmt->header_line) {
		first = true;
		dirty = cfg->print_style(file, cfg, NULL, NULL);
		for (i = 0; i < cfg->column_count; i++) {
			if (fmt->columns[i].hidden) continue;
			if (!first) {
				fprintf(file, "%s%*.s", cfg->xsep,
					(int) (fmt->mseplen - fmt->xseplen), "");
			}
			for (k = 0; k < fmt->columns[i].width; k++)
				fprintf(file, "%s", cfg->vsep);
			first = false;
		}
		if (dirty) fprintf(file, "\x1b[0m");
		fprintf(file, "\n");
		stats->lines_used += 1;
	}
}

void
output_row(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stat, const struct tabular_row *row,
	struct fmt_state *fmt)
{
	size_t wwidth, padwidth, outlen;
	size_t i, off, width;
	bool first, done, dirty;
	char *entry;

	for (i = 0; i < cfg->column_count; i++) {
		fmt->columns[i].written = 0;

		if (cfg->columns[i].strategy == TABULAR_TRUNC) {
			width = fmt->columns[i].width
				- fmt->columns[i].lpad - fmt->columns[i].rpad;
			if (row->entries[i].ulen > width) {
				width = u8rawlen(row->entries[i].str, width - 2);
				row->entries[i].str[width] = '.';
				row->entries[i].str[width+1] = '.';
				row->entries[i].str[width+2] = '\0';
			}
		}
	}

	do {
		done = true;
		first = true;
		for (i = 0; i < cfg->column_count; i++) {
			if (fmt->columns[i].hidden) continue;

			if (!first) {
				dirty = cfg->print_style(file, cfg, NULL, NULL);
				fprintf(file, "%s", cfg->hsep);
				if (dirty) fprintf(file, "\x1b[0m");
			}
			first = false;

			if (!row->entries[i].str) {
				print_pad(file, fmt->columns[i].width);
				continue;
			}

			entry = row->entries[i].str + fmt->columns[i].written;

			dirty = cfg->print_style(file, cfg, row, &cfg->columns[i]);

			off = 0;
			width = fmt->columns[i].width
				- fmt->columns[i].lpad - fmt->columns[i].rpad;
			padwidth = width;
			if (cfg->columns[i].strategy == TABULAR_WRAP_WORDAWARE) {
				calc_word_aware_line(entry, width, &off, &wwidth);
				entry += off;
				width = wwidth;
			}

			print_pad(file, fmt->columns[i].lpad);
			if (cfg->columns[i].align == TABULAR_ALIGN_LEFT) {
				print_left(file, entry, width, padwidth);
			} else if (cfg->columns[i].align == TABULAR_ALIGN_RIGHT) {
				print_right(file, entry, width, padwidth);
			} else {
				print_center(file, entry, width, padwidth);
			}
			print_pad(file, fmt->columns[i].rpad);

			if (dirty) fprintf(file, "\x1b[0m");

			outlen = MIN(u8strlen(entry), width + off);
			fmt->columns[i].written += u8rawlen(entry, outlen);

			/* check if anything other than whitespace left */
			entry += u8rawlen(entry, outlen);
			if (strspn(entry, " \t\v\r\n") != u8strlen(entry))
				done = false;
		}
		fprintf(file, "\n");
		stat->lines_used += 1;
	} while (!done);

	stat->rows_displayed += 1;
}

int
output_rows(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stats, struct tabular_row *rows,
	struct fmt_state *fmt)
{
	struct tabular_row *row;
	size_t count;

	if (fmt->row_limit == 0) return 0;

	output_header(file, cfg, stats, fmt);

	count = 0;
	for (row = rows; row; row = row->next) {
		if (count == fmt->row_limit)
			break;
		output_row(file, cfg, stats, row, fmt);
		count += 1;
	}

	return 0;
}

int
tabular_format(FILE *file, const struct tabular_cfg *cfg,
	struct tabular_stats *stats, struct tabular_row **rows)
{
	struct fmt_state fmt;
	size_t i;
	int rc;

	stats->lines_used = 0;
	stats->rows_displayed = 0;
	stats->rows_truncated = false;
	stats->cols_truncated = false;

	if (!rows) return 1;

	if (!cfg->column_count) return 0;

	for (i = 0; i < cfg->column_count; i++) {
		if (cfg->columns[i].minwidth > cfg->columns[i].maxwidth)
			return 1;
		if (cfg->columns[i].lpad + cfg->columns[i].rpad
				>= cfg->columns[i].minwidth)
			return 1;
		if (cfg->columns[i].strategy != TABULAR_TRUNC
				&& cfg->columns[i].strategy != TABULAR_WRAP
				&& cfg->columns[i].strategy != TABULAR_WRAP_WORDAWARE)
			return 1;
	}

	rc = calc_params(cfg, stats, rows, &fmt);
	if (rc) return rc;

	rc = output_rows(file, cfg, stats, *rows, &fmt);
	if (rc) return rc;

	rc = cfg->allocator->free(cfg->allocator, fmt.columns);
	if (rc) return -rc;

	return 0;
}

struct tabular_row *
tabular_alloc_row(const struct tabular_cfg *cfg,
	int *rc, struct tabular_user user)
{
	struct tabular_row *row;

	row = cfg->allocator->alloc(cfg->allocator,
		sizeof(struct tabular_row), rc);
	if (!row && rc) *rc = -*rc;
	if (!row) return NULL;

	row->next = NULL;
	row->user = user;
	row->entries = cfg->allocator->alloc(cfg->allocator,
		sizeof(struct tabular_entry) * cfg->column_count, rc);
	if (!row->entries && rc) *rc = -*rc;
	if (!row->entries) return NULL;
	memset(row->entries, 0, sizeof(struct tabular_entry) * cfg->column_count);

	return row;
}

int
tabular_free_rows(const struct tabular_cfg *cfg, struct tabular_row *rows)
{
	struct tabular_row *iter, *next;
	size_t i;
	int rc;

	for (iter = rows; iter; iter = next) {
		next = iter->next;
		for (i = 0; i < cfg->column_count; i++) {
			rc = cfg->allocator->free(cfg->allocator,
				iter->entries[i].str);
			if (rc) return -rc;
		}
		rc = cfg->allocator->free(cfg->allocator, iter->entries);
		if (rc) return -rc;
		rc = cfg->allocator->free(cfg->allocator, iter);
		if (rc) return -rc;
	}

	return 0;
}

void
tabular_load_row_entry_hidden(const struct tabular_cfg *cfg,
	struct tabular_row *row, size_t col)
{
	bool hidden;

	if (row->entries[col].flags & BIT(TABULAR_ENTRY_HIDDEN_SET)) return;

	if (cfg->columns[col].is_hidden) {
		hidden = cfg->columns[col].is_hidden(
			&row->user, &cfg->columns[col].user);
	} else {
		hidden = false;
	}

	if (hidden) {
		row->entries[col].flags |= BIT(TABULAR_ENTRY_HIDDEN);
	} else {
		row->entries[col].flags &= ~BIT(TABULAR_ENTRY_HIDDEN);
	}

	row->entries[col].flags |= BIT(TABULAR_ENTRY_HIDDEN_SET);
}

void
tabular_load_row_entry_str(const struct tabular_cfg *cfg,
	struct tabular_row *row, size_t col)
{
	if (row->entries[col].flags & BIT(TABULAR_ENTRY_STR_SET)) return;

	row->entries[col].str = cfg->columns[col].to_str(
		&row->user, &cfg->columns[col].user);
	row->entries[col].ulen = (uint32_t) u8strlen(row->entries[col].str);
	row->entries[col].flags |= BIT(TABULAR_ENTRY_STR_SET);
}