#include <TeX/node.h>
#include <TeX/flow.h>
#include <TeX/env.h>

#include <TeX/vector.h>

#include <stdlib.h>
#include <string.h>

#define fill_prototype(env, prefix) {			\
	env->free = prefix ## _free;			\
	env->add_node = prefix ## _add_node;		\
	env->add_separator = prefix ## _add_separator;	\
	env->add_break = prefix ## _add_break;		\
	env->layout = prefix ## _layout;		\
	env->render = prefix ## _render;		\
}

//---
//	Primary environment: a simple flow
//
//	This environment ignores separators and breaks and concatenates all
//	nodes in a single flow. It is the type of the outermost environment of
//	any formula.
//---

struct TeX_Env_Primary
{
	struct TeX_Env env;

	struct TeX_Flow *flow;
};

static void primary_free(struct TeX_Env *env)
{
	struct TeX_Env_Primary *p = (void *)env;
	TeX_flow_free(p->flow);
	free(env);
}

static void primary_add_node(struct TeX_Env *env, struct TeX_Node *node)
{
	struct TeX_Env_Primary *p = (void *)env;
	p->flow = TeX_flow_add_node(p->flow, node);
}

static void primary_add_separator(__attribute__((unused)) struct TeX_Env *env)
{
}

static void primary_add_break(__attribute__((unused)) struct TeX_Env *env)
{
}

static void primary_layout(struct TeX_Env *env)
{
	struct TeX_Env_Primary *p = (void *)env;
	TeX_flow_layout(p->flow);

	env->width = p->flow->width;
	env->height = p->flow->height;
	env->line = p->flow->line;
}

static void primary_render(struct TeX_Env *env, int x, int y, int color)
{
	struct TeX_Env_Primary *p = (void *)env;
	TeX_flow_render(p->flow, x, y, color);
}

struct TeX_Env *TeX_env_primary(void)
{
	struct TeX_Env_Primary *p = malloc(sizeof *p);
	struct TeX_Env *env = (void *)p;

	env->name = "primary";
	fill_prototype(env, primary);

	p->flow = NULL;

	return env;
}

//---
//	Matrix environment.
//---

struct TeX_Env_Matrix
{
	struct TeX_Env env;

	vector_type(struct TeX_Flow *, elements);
	vector_type(int, rowstart);

	int *rowdepth;
	int *rowline;
	int *colwidth;

	int rows;
	int cols;
};

static void matrix_free(struct TeX_Env *env)
{
	struct TeX_Env_Matrix *m = (void *)env;

	free(m->elements);
	free(m->rowstart);
	free(m->rowdepth);
	free(m->rowline);
	free(m->colwidth);

	free(env);
}

static void matrix_add_node(struct TeX_Env *env, struct TeX_Node *node)
{
	struct TeX_Env_Matrix *m = (void *)env;

	int last = m->elements_len - 1;
	m->elements[last] = TeX_flow_add_node(m->elements[last], node);
}

static void matrix_add_separator(struct TeX_Env *env)
{
	struct TeX_Env_Matrix *m = (void *)env;

	/* Create a new flow at the end of the element array */
	vector_append(m->elements, NULL);

	/* Current row length is automatically increased. Set column count */
	int row_length = m->elements_len - m->rowstart[m->rows - 1];
	m->cols = max(m->cols, row_length);
}

static void matrix_add_break(struct TeX_Env *env)
{
	struct TeX_Env_Matrix *m = (void *)env;

	/* First add a new flow */
	vector_append(m->elements, NULL);

	/* Then add a row and its start index */
	vector_append(m->rowstart, m->elements_len - 1);
	m->rows++;
}

static void matrix_layout(struct TeX_Env *env)
{
	struct TeX_Env_Matrix *m = (void *)env;

	int row = 0;
	int col = 0;

	env->width = 0;
	env->height = 0;
	env->line = 0;

	m->rowdepth = calloc(m->rows, sizeof *m->rowdepth);
	m->rowline  = calloc(m->rows, sizeof *m->rowline);
	m->colwidth = calloc(m->cols, sizeof *m->colwidth);
	if(!m->rowdepth || !m->rowline || !m->colwidth) return;

	/* Compute the layout of the grid */
	for(int i = 0; i < m->elements_len; i++)
	{
		if(row + 1 < m->rows && i == m->rowstart[row + 1])
		{
			col = 0;
			row++;
		}

		struct TeX_Flow *f = m->elements[i];
		TeX_flow_layout(f);

		/* Update the current row and column */
		m->rowline[row] = max(m->rowline[row], f->line);
		m->rowdepth[row] = max(m->rowdepth[row], f->height - f->line);
		m->colwidth[col] = max(m->colwidth[col], f->width);

		col++;
	}

	/* TODO: Add width spacing and height spacing */

	for(int i = 0; i < m->rows; i++)
	{
		env->height += m->rowline[i] + m->rowdepth[i];
	}
	for(int j = 0; j < m->cols; j++)
	{
		env->width += m->colwidth[j];
	}

	env->width  += TEX_MATRIX_COL_SPACING * (m->cols - 1);
	env->height += TEX_MATRIX_ROW_SPACING * (m->rows - 1);
	env->line = env->height >> 1;
}

static void matrix_render(struct TeX_Env *env, int x, int y, int color)
{
	struct TeX_Env_Matrix *m = (void *)env;

	int n = 0, dy = 0;

	for(int row = 0; row < m->rows; row++)
	{
		int col = 0, dx = 0;

		while((row + 1 <  m->rows && n < m->rowstart[row + 1]) ||
		      (row + 1 >= m->rows && n < m->elements_len))
		{
			struct TeX_Flow *flow = m->elements[n++];
			int cw = m->colwidth[col];
			int rh = m->rowline[row] + m->rowdepth[row];

			TeX_flow_render(flow,
				x + dx + ((cw - flow->width) >> 1),
				y + dy + ((rh - flow->height) >> 1),
				color);

			dx += cw + TEX_MATRIX_COL_SPACING;
			col++;
		}

		dy += m->rowline[row]+m->rowdepth[row]+TEX_MATRIX_ROW_SPACING;
	}
}

struct TeX_Env *TeX_env_matrix(void)
{
	struct TeX_Env_Matrix *m = malloc(sizeof *m);
	struct TeX_Env *env = (void *)m;

	env->name = "matrix";
	fill_prototype(env, matrix);

	vector_init(m->elements);
	vector_init(m->rowstart);
	m->rowdepth = NULL;
	m->rowline = NULL;
	m->colwidth = NULL;

	vector_extend(m->elements);
	m->elements[0] = NULL;
	m->elements_len = 1;

	vector_extend(m->rowstart);
	m->rowstart[0] = 0;
	m->rowstart_len = 1;

	m->rows = 1;
	m->cols = 0;

	return env;
}

//---
//	Environment printers
//---

#ifdef TEX_PRINT

#include <stdio.h>
#include <TeX/print.h>

/* TeX_print_env_primary(): Print a primary environment to stdout */
void TeX_print_env_primary(struct TeX_Env *env, int indent)
{
	struct TeX_Env_Primary *p = (void *)env;

	printf("%*senv:primary", indent, "");
	printf(GRAY " %dx%d,%d" END "\n", env->width, env->height, env->line);

	TeX_print_flow(p->flow, indent + 4);
}

/* TeX_print_env_matrix(): Print a matrix environment to stdout */
void TeX_print_env_matrix(struct TeX_Env *env, int indent)
{
	struct TeX_Env_Matrix *m = (void *)env;

	printf("%*senv:matrix %dx%d", indent, "", m->rows, m->cols);
	printf(GRAY " %dx%d,%d", env->width, env->height, env->line);

	for(int i = 0; i < m->cols; i++)
	{
		printf("%c%d", (i ? ';' : ' '), m->colwidth[i]);
	}
	printf(END "\n");

	int row = 0;
	for(int i = 0; i < m->elements_len; i++)
	{
		if(row < m->rows && i == m->rowstart[row])
		{
			printf("%*s%d: " GRAY "%d,%d" END "\n", indent + 4, "",
				row, m->rowline[row] + m->rowdepth[row],
				m->rowline[row]);
			row++;
		}

		TeX_print_flow(m->elements[i], indent + 8);
	}
}

/* TeX_print_env(): Recursively print an environment */
void TeX_print_env(struct TeX_Env *env, int indent)
{
	if(!strcmp(env->name, "primary"))
	{
		TeX_print_env_primary(env, indent);
	}
	else if(!strcmp(env->name, "matrix"))
	{
		TeX_print_env_matrix(env, indent);
	}
	else
	{
		printf("%*senv:%s (?)\n", indent, "", env->name);
	}
}

#endif /* TEX_PRINT */