#include "../debug.h"
#include "cmd.h"
#include "mie/select/builder.h"

#include <blue/cmd.h>
#include <blue/core/error.h>
#include <blue/term.h>
#include <errno.h>
#include <ivy/asm/mie.h>
#include <ivy/diag.h>
#include <ivy/file.h>
#include <ivy/lang/ast.h>
#include <ivy/lang/codegen.h>
#include <ivy/lang/diag.h>
#include <ivy/lang/lex.h>
#include <mie/ctx.h>
#include <mie/ir/block.h>
#include <mie/ir/convert.h>
#include <mie/ir/func.h>
#include <mie/ir/module.h>
#include <mie/select/builder.h>
#include <mie/select/graph.h>
#include <stdio.h>

enum {
	ARG_SOURCE_FILE = 200,
	OPT_SHOW_LEX_TOKENS,
	OPT_SHOW_AST_NODES,
	OPT_SHOW_IR,
	OPT_SHOW_ISEL_PRE,
	OPT_SHOW_ISEL_POST,
};

enum compile_flags {
	FLAG_SHOW_LEX_TOKENS = 0x01u,
	FLAG_SHOW_AST_NODES = 0x02u,
	FLAG_SHOW_IR = 0x04u,
	FLAG_SHOW_PRE_SELECT_GRAPH = 0x08u,
	FLAG_SHOW_POST_SELECT_GRAPH = 0x10u,
};

struct compile_ctx {
	enum compile_flags flags;
	const char *src_path;
	struct ivy_file *src;
	struct ivy_diag_ctx *diag;
	struct ivy_diag_stream diag_stream;

	struct ivy_lexer *lex;
	struct ivy_parser *parser;
	struct ivy_codegen *codegen;

	struct mie_ctx *mie_ctx;
	struct mie_select_builder *select;

	struct mie_module *mod;
};

static enum ivy_status node_codegen(
	struct ivy_ast_node *node, enum ivy_ast_iteration_type iteration_type,
	struct ivy_ast_node_iterator *it, void *arg)
{
	if (iteration_type != IVY_AST_ITERATION_PRE) {
		return IVY_OK;
	}

	struct compile_ctx *ctx = arg;
	return ivy_codegen_push_node(ctx->codegen, node, node->n_it.it_depth);
}

static b_result compile_ctx_init(struct compile_ctx *ctx, const b_arglist *args)
{
	memset(ctx, 0x0, sizeof *ctx);

	b_arglist_get_count(args, OPT_SHOW_LEX_TOKENS, B_COMMAND_INVALID_ID)
		&& (ctx->flags |= FLAG_SHOW_LEX_TOKENS);
	b_arglist_get_count(args, OPT_SHOW_AST_NODES, B_COMMAND_INVALID_ID)
		&& (ctx->flags |= FLAG_SHOW_AST_NODES);
	b_arglist_get_count(args, OPT_SHOW_IR, B_COMMAND_INVALID_ID)
		&& (ctx->flags |= FLAG_SHOW_IR);
	b_arglist_get_count(args, OPT_SHOW_ISEL_PRE, B_COMMAND_INVALID_ID)
		&& (ctx->flags |= FLAG_SHOW_PRE_SELECT_GRAPH);
	b_arglist_get_count(args, OPT_SHOW_ISEL_POST, B_COMMAND_INVALID_ID)
		&& (ctx->flags |= FLAG_SHOW_POST_SELECT_GRAPH);

	enum mie_status mie_status = MIE_SUCCESS;
	enum ivy_status ivy_status = IVY_OK;

	ctx->mie_ctx = mie_ctx_create();

	ivy_status = ivy_diag_ctx_create(&ctx->diag);
	ivy_status = ivy_lexer_create(&ctx->lex);
	ivy_status = ivy_parser_create(&ctx->parser);
	ivy_status = ivy_codegen_create(ctx->mie_ctx, &ctx->codegen);
	ctx->select
		= mie_select_builder_create(ctx->mie_ctx, ivy_asm_mie_target());

	ivy_lang_diag_ctx_init(ctx->diag);
	ivy_diag_stream_init_tty(&ctx->diag_stream, b_stdtty);
	ivy_lexer_set_diag_ctx(ctx->lex, ctx->diag);
	ivy_parser_set_diag_ctx(ctx->parser, ctx->diag);

	return B_RESULT_SUCCESS;
}

static b_result compile_ctx_cleanup(struct compile_ctx *ctx)
{
	if (ctx->mod) {
		mie_value_destroy(MIE_VALUE(ctx->mod));
	}

	if (ctx->src) {
		ivy_diag_ctx_set_line_source(ctx->diag, NULL);
		ivy_lexer_set_source(ctx->lex, NULL);
		ivy_file_close(ctx->src);
	}

	if (ctx->select) {
		mie_select_builder_destroy(ctx->select);
	}

	if (ctx->codegen) {
		ivy_codegen_destroy(ctx->codegen);
	}

	if (ctx->parser) {
		ivy_parser_destroy(ctx->parser, NULL);
	}

	if (ctx->lex) {
		ivy_lexer_destroy(ctx->lex);
	}

	if (ctx->diag) {
		ivy_diag_ctx_destroy(ctx->diag);
	}

	if (ctx->mie_ctx) {
		mie_ctx_destroy(ctx->mie_ctx);
	}

	return B_RESULT_SUCCESS;
}

static b_result open_file(struct compile_ctx *ctx, const char *path)
{
	enum ivy_status status = ivy_file_open(path, &ctx->src);

	ivy_diag_ctx_set_line_source(ctx->diag, &ctx->src->f_base);
	ivy_lexer_set_source(ctx->lex, &ctx->src->f_base);

	ctx->src_path = path;

	return B_RESULT_SUCCESS;
}

static b_result parse_file_tokens(struct compile_ctx *ctx)
{
	enum ivy_status status = IVY_OK;

	while (true) {
		struct ivy_token *tok = ivy_lexer_read(ctx->lex);
		status = ivy_lexer_get_status(ctx->lex);

		if (status != IVY_OK && status != IVY_ERR_EOF) {
			ivy_diag_ctx_write(
				ctx->diag, IVY_DIAG_FORMAT_PRETTY,
				&ctx->diag_stream);
			return b_error_with_code(
				IVY_ERROR_VENDOR, IVY_ERR_BAD_SYNTAX);
		}

		if (tok && (ctx->flags & FLAG_SHOW_LEX_TOKENS)) {
			print_lex_token(tok);
		}

		status = ivy_parser_push_token(ctx->parser, tok);

		if (status == IVY_ERR_BAD_SYNTAX) {
			ivy_diag_ctx_write(
				ctx->diag, IVY_DIAG_FORMAT_PRETTY,
				&ctx->diag_stream);
			ivy_token_destroy(tok);
			return b_error_with_code(
				IVY_ERROR_VENDOR, IVY_ERR_BAD_SYNTAX);
		}

		if (status != IVY_OK) {
			return b_error_caused_by_code(
				IVY_ERROR_VENDOR, IVY_ERR_PARSE_FAILURE,
				IVY_ERROR_VENDOR, status);
		}

		if (!tok) {
			break;
		}
	}

	return B_RESULT_SUCCESS;
}

static b_result generate_mie_ir(struct compile_ctx *ctx)
{
	ivy_codegen_start_module(ctx->codegen);

	struct ivy_ast_node *node = ivy_parser_root_node(ctx->parser);
	struct ivy_ast_node_iterator it = {0};

	enum ivy_status status
		= ivy_ast_node_iterate(node, &it, node_codegen, ctx);
	if (status != IVY_OK) {
		return b_error_caused_by_code(
			IVY_ERROR_VENDOR, IVY_ERR_CODEGEN_FAILURE,
			IVY_ERROR_VENDOR, status);
	}

	status = ivy_codegen_push_eof(ctx->codegen);
	if (status != IVY_OK) {
		return b_error_caused_by_code(
			IVY_ERROR_VENDOR, IVY_ERR_CODEGEN_FAILURE,
			IVY_ERROR_VENDOR, status);
	}

	ivy_codegen_end_module(ctx->codegen, &ctx->mod);

	return B_RESULT_SUCCESS;
}

static b_result build_block_isel_graph(
	struct compile_ctx *ctx, struct mie_func *func, struct mie_block *block)
{
	b_queue_entry *entry = b_queue_first(&block->b_phi);
	while (entry) {
		struct mie_value *instr_v
			= b_unbox(struct mie_value, entry, v_entry);
		mie_select_builder_push_instr(
			ctx->select, (struct mie_instr *)instr_v);
		entry = b_queue_next(entry);
	}

	entry = b_queue_first(&block->b_instr);
	while (entry) {
		struct mie_value *instr_v
			= b_unbox(struct mie_value, entry, v_entry);
		mie_select_builder_push_instr(
			ctx->select, (struct mie_instr *)instr_v);
		entry = b_queue_next(entry);
	}

	if (block->b_terminator) {
		mie_select_builder_push_instr(ctx->select, block->b_terminator);
	}

	struct mie_select_graph *graph = mie_select_builder_finish(ctx->select);

	if (ctx->flags & FLAG_SHOW_PRE_SELECT_GRAPH) {
		printf("%s.%s instruction graph:\n", func->f_base.v_name.n_str,
		       block->b_base.v_name.n_str);
		mie_select_graph_dump_dot(graph);
		mie_select_graph_destroy(graph);
	}

	return B_RESULT_SUCCESS;
}

static b_result build_func_isel_graph(struct compile_ctx *ctx, struct mie_func *func)
{
	b_queue_entry *entry = b_queue_first(&func->f_blocks);
	while (entry) {
		struct mie_value *block_v
			= b_unbox(struct mie_value, entry, v_entry);
		struct mie_block *block = (struct mie_block *)block_v;

		b_result result = build_block_isel_graph(ctx, func, block);
		if (b_result_is_error(result)) {
			return b_result_propagate(result);
		}

		entry = b_queue_next(entry);
	}

	return B_RESULT_SUCCESS;
}

static b_result build_isel_graph(struct compile_ctx *ctx)
{
	b_queue_entry *entry = b_queue_first(&ctx->mod->m_func);
	while (entry) {
		struct mie_value *func_v
			= b_unbox(struct mie_value, entry, v_entry);
		struct mie_func *func = (struct mie_func *)func_v;

		b_result result = build_func_isel_graph(ctx, func);
		if (b_result_is_error(result)) {
			return b_result_propagate(result);
		}

		entry = b_queue_next(entry);
	}

	return B_RESULT_SUCCESS;
}

static b_result dump_ast(struct compile_ctx *ctx)
{
	struct ivy_ast_node *node = ivy_parser_root_node(ctx->parser);
	struct ivy_ast_node_iterator it = {0};
	struct print_ast_node_args args = {
		.indent = true,
	};

	ivy_ast_node_iterate(node, &it, print_ast_node, &args);

	return B_RESULT_SUCCESS;
}

static b_result dump_ir(struct compile_ctx *ctx)
{
	struct mie_ir_converter *convert
		= mie_ir_converter_create(ctx->mie_ctx, MIE_IR_MEM, MIE_IR_TEXT);
	mie_ir_converter_set_src_value(convert, MIE_VALUE(ctx->mod));
	mie_ir_converter_set_dest_file(convert, stdout);
	mie_ir_converter_process(convert);

	mie_ir_converter_destroy(convert);

	return B_RESULT_SUCCESS;
}

static int compile_file(const char *path, const b_arglist *args)
{
#define THROW_AND_RETURN(result, error_code)                                     \
	do {                                                                     \
		if (b_result_is(result, IVY_ERROR_VENDOR, IVY_ERR_BAD_SYNTAX)) { \
			return error_code;                                       \
		}                                                                \
		if (b_result_is_error(result)) {                                 \
			b_throw(result);                                         \
			return error_code;                                       \
		}                                                                \
	} while (0)

	struct compile_ctx ctx;
	b_result result = compile_ctx_init(&ctx, args);
	THROW_AND_RETURN(result, -1);

	int progress = 0;
	(ctx.flags & FLAG_SHOW_LEX_TOKENS) && (progress = 1);
	(ctx.flags & FLAG_SHOW_AST_NODES) && (progress = 1);
	(ctx.flags & FLAG_SHOW_IR) && (progress = 2);
	(ctx.flags & FLAG_SHOW_PRE_SELECT_GRAPH) && (progress = 3);
	(ctx.flags & FLAG_SHOW_POST_SELECT_GRAPH) && (progress = 4);

	result = open_file(&ctx, path);
	THROW_AND_RETURN(result, -1);

	result = parse_file_tokens(&ctx);
	THROW_AND_RETURN(result, -1);

	if (ctx.flags & FLAG_SHOW_AST_NODES) {
		dump_ast(&ctx);
	}

	if (progress == 1) {
		return 0;
	}

	result = generate_mie_ir(&ctx);
	THROW_AND_RETURN(result, -1);

	if (ctx.flags & FLAG_SHOW_IR) {
		dump_ir(&ctx);
	}

	if (progress == 2) {
		return 0;
	}

	result = build_isel_graph(&ctx);
	THROW_AND_RETURN(result, -1);

	if (progress == 3) {
		return 0;
	}

end:
	compile_ctx_cleanup(&ctx);

#undef THROW_AND_RETURN
	return 0;
}

static int compile(const b_command *cmd, const b_arglist *args, const b_array *_)
{
	b_arglist_iterator it;
	b_arglist_foreach_filtered(&it, args, B_COMMAND_INVALID_ID, ARG_SOURCE_FILE)
	{
		b_arglist_value *path = it.value;
		if (path->val_type != B_COMMAND_ARG_STRING) {
			continue;
		}

		int r = compile_file(path->val_str, args);
		if (r != 0) {
			return r;
		}
	}

	return 0;
}

B_COMMAND(CMD_COMPILE, CMD_ROOT)
{
	B_COMMAND_NAME("compile");
	B_COMMAND_SHORT_NAME('C');
	B_COMMAND_DESC(
		"compile one or more Ivy source files into Ivy object files.");
	B_COMMAND_FLAGS(B_COMMAND_SHOW_HELP_BY_DEFAULT);
	B_COMMAND_FUNCTION(compile);

	B_COMMAND_OPTION(OPT_SHOW_LEX_TOKENS)
	{
		B_OPTION_LONG_NAME("show-lex");
		B_OPTION_SHORT_NAME('l');
		B_OPTION_DESC(
			"print the lexical tokens generated from the input "
			"files.");
	}

	B_COMMAND_OPTION(OPT_SHOW_AST_NODES)
	{
		B_OPTION_LONG_NAME("show-ast");
		B_OPTION_SHORT_NAME('a');
		B_OPTION_DESC(
			"print the abstract syntax tree generated from the "
			"input files.");
	}

	B_COMMAND_OPTION(OPT_SHOW_IR)
	{
		B_OPTION_LONG_NAME("show-ir");
		B_OPTION_SHORT_NAME('i');
		B_OPTION_DESC(
			"print the Mie IR generated from the "
			"input files.");
	}

	B_COMMAND_OPTION(OPT_SHOW_ISEL_PRE)
	{
		B_OPTION_LONG_NAME("show-isel-pre");
		B_OPTION_SHORT_NAME('s');
		B_OPTION_DESC(
			"print the instruction selection graph before "
			"selection has taken place.");
	}

	B_COMMAND_OPTION(OPT_SHOW_ISEL_POST)
	{
		B_OPTION_LONG_NAME("show-isel-post");
		B_OPTION_SHORT_NAME('m');
		B_OPTION_DESC(
			"print the instruction selection graph after "
			"selection has taken place.");
	}

	B_COMMAND_ARG(ARG_SOURCE_FILE)
	{
		B_ARG_NAME("source file");
		B_ARG_DESC("the .im source files to compile.");
		B_ARG_NR_VALUES(B_ARG_1_OR_MORE_VALUES);
	}

	B_COMMAND_HELP_OPTION();
}