#include "command.h"

#include <blue/cmd.h>
#include <blue/object/string.h>
#include <blue/term.h>
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define OUTPUT_STREAM stderr

static struct b_btree command_list = {};

B_BTREE_DEFINE_SIMPLE_GET(struct b_command, unsigned int, b_node, b_id, get_command)
B_BTREE_DEFINE_SIMPLE_INSERT(struct b_command, b_node, b_id, put_command)

enum item_type {
	ITEM_OPTION,
	ITEM_ARG,
	ITEM_SUBCOMMAND,
};

struct b_command *b_command_create(unsigned int id)
{
	struct b_command *out = malloc(sizeof *out);
	if (!out) {
		return NULL;
	}

	memset(out, 0x0, sizeof *out);

	out->b_id = id;
	out->b_parent_id = B_COMMAND_INVALID_ID;
	return out;
}

void b_command_destroy(struct b_command *cmd)
{
}

b_status b_command_register(struct b_command *cmd)
{
	struct b_command *tmp = get_command(&command_list, cmd->b_id);
	if (tmp) {
		return B_ERR_NAME_EXISTS;
	}

	put_command(&command_list, cmd);
	return B_SUCCESS;
}

b_status b_command_set_name(struct b_command *cmd, const char *name)
{
	char *n = b_strdup(name);
	if (!n) {
		return B_ERR_NO_MEMORY;
	}

	if (cmd->b_name) {
		free(cmd->b_name);
		cmd->b_name = NULL;
	}

	cmd->b_name = n;
	return B_SUCCESS;
}

b_status b_command_set_long_name(struct b_command *cmd, const char *name)
{
	char *n = b_strdup(name);
	if (!n) {
		return B_ERR_NO_MEMORY;
	}

	if (cmd->b_long_name) {
		free(cmd->b_long_name);
		cmd->b_long_name = NULL;
	}

	cmd->b_long_name = n;
	return B_SUCCESS;
}

b_status b_command_set_short_name(struct b_command *cmd, char name)
{
	cmd->b_short_name = name;
	return B_SUCCESS;
}

b_status b_command_set_flags(struct b_command *cmd, b_command_flags flags)
{
	cmd->b_flags = flags;
	return B_SUCCESS;
}

b_status b_command_set_description(struct b_command *cmd, const char *description)
{
	char *desc = b_strdup(description);
	if (!desc) {
		return B_ERR_NO_MEMORY;
	}

	if (cmd->b_description) {
		free(cmd->b_description);
		cmd->b_description = NULL;
	}

	cmd->b_description = desc;
	return B_SUCCESS;
}

b_status b_command_set_callback(struct b_command *cmd, b_command_callback callback)
{
	cmd->b_callback = callback;
	return B_SUCCESS;
}

b_status b_command_set_parent(b_command *cmd, unsigned int parent_id)
{
	cmd->b_parent_id = parent_id;
	cmd->b_parent = NULL;
	return B_SUCCESS;
}

struct b_command_option *b_command_add_option(struct b_command *cmd, int id)
{
	struct b_command_option *opt = b_command_option_create();
	if (!opt) {
		return NULL;
	}

	opt->opt_id = id;
	b_queue_push_back(&cmd->b_opt, &opt->opt_entry);
	return opt;
}

struct b_command_arg *b_command_add_arg(struct b_command *cmd, int id)
{
	struct b_command_arg *arg = b_command_arg_create();
	if (!arg) {
		return NULL;
	}

	arg->arg_id = id;
	b_queue_push_back(&cmd->b_arg, &arg->arg_entry);
	return arg;
}

struct b_command_usage *b_command_add_usage(struct b_command *cmd)
{
	struct b_command_usage *usage = malloc(sizeof *usage);
	if (!usage) {
		return NULL;
	}

	memset(usage, 0x0, sizeof *usage);
	b_queue_push_back(&cmd->b_usage, &usage->u_entry);

	return usage;
}

b_status b_command_usage_add_option(
	struct b_command_usage *usage, struct b_command_option *opt)
{
	struct b_command_usage_opt *u_opt = malloc(sizeof *u_opt);
	if (!u_opt) {
		return B_ERR_NO_MEMORY;
	}

	memset(u_opt, 0x0, sizeof *u_opt);

	u_opt->opt = opt;

	b_queue_push_back(&usage->u_opt, &u_opt->opt_entry);
	return B_SUCCESS;
}

b_status b_command_usage_add_arg(
	struct b_command_usage *usage, struct b_command_arg *arg)
{
	struct b_command_usage_arg *u_arg = malloc(sizeof *u_arg);
	if (!u_arg) {
		return B_ERR_NO_MEMORY;
	}

	memset(u_arg, 0x0, sizeof *u_arg);

	u_arg->arg = arg;

	b_queue_push_back(&usage->u_arg, &u_arg->arg_entry);
	return B_SUCCESS;
}

static void prepend_command_name(struct b_command *cmd, b_string *out)
{
	int nr_names = 0;
	cmd->b_name &&nr_names++;
	cmd->b_long_name &&nr_names++;
	cmd->b_short_name &&nr_names++;

	if (nr_names > 1) {
		b_string_prepend_cstr(out, "}");
		int r = 0;

		if (cmd->b_name) {
			b_string_prepend_cstrf(out, "%s", cmd->b_name);
			r++;
		}

		if (r == 1) {
			b_string_prepend_cstr(out, "|");
		}

		if (cmd->b_long_name) {
			b_string_prepend_cstrf(out, "--%s", cmd->b_long_name);
			r++;
		}

		if (r == 2) {
			b_string_prepend_cstr(out, "|");
		}

		if (cmd->b_short_name) {
			b_string_prepend_cstrf(out, "-%c", cmd->b_short_name);
			r++;
		}

		b_string_prepend_cstr(out, "{");
	} else {
		b_string_prepend_cstr(out, cmd->b_name);
	}
}

static void get_qualified_command_name(struct b_command *cmd, b_string *out)
{
	prepend_command_name(cmd, out);
	cmd = cmd->b_parent;

	while (cmd) {
		b_string_prepend_cstr(out, " ");
		prepend_command_name(cmd, out);
		cmd = cmd->b_parent;
	}
}

static void get_usage_string(
	struct b_command *cmd, struct b_command_usage *usage, b_string *out)
{
	get_qualified_command_name(cmd, out);

	b_queue_iterator it;
	b_queue_foreach (&it, &usage->u_opt) {
		struct b_command_usage_opt *opt = b_unbox(
			struct b_command_usage_opt, it.entry, opt_entry);

		if (!opt) {
			break;
		}

		b_string_append_cstr(out, " ");
		z__b_get_option_usage_string(opt->opt, CMD_STR_DIRECT_USAGE, out);
	}

	b_queue_foreach (&it, &usage->u_arg) {
		struct b_command_usage_arg *arg = b_unbox(
			struct b_command_usage_arg, it.entry, arg_entry);

		if (!arg) {
			break;
		}

		b_string_append_cstr(out, " ");
		z__b_get_arg_usage_string(arg->arg, false, out);
	}
}

b_string *z__b_command_default_usage_string(
	struct b_command *cmd, struct b_command_option *with_opt)
{
	b_string *str = b_string_create();
	get_qualified_command_name(cmd, str);

	if (with_opt) {
		b_string_append_cstr(str, " ");
		z__b_get_option_usage_string(with_opt, CMD_STR_DIRECT_USAGE, str);
	} else if (!b_queue_empty(&cmd->b_opt)) {
		b_string_append_cstr(str, " [OPTIONS]");
	}

	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_arg) {
		struct b_command_arg *arg
			= b_unbox(struct b_command_arg, it.entry, arg_entry);

		b_string_append_cstr(str, " ");
		z__b_get_arg_usage_string(arg, false, str);
	}

	if (!b_queue_empty(&cmd->b_subcommands)) {
		b_string_append_cstr(str, " [COMMAND]");
	}

	return str;
}

static void get_command_string(struct b_command *cmd, b_string *out)
{
	int r = 0;

	if (cmd->b_name) {
		b_string_append_cstrf(out, F_GREEN "%s" F_RESET, cmd->b_name);
		r++;
	}

	if (r == 1) {
		b_string_append_cstr(out, ", ");
	}

	if (cmd->b_short_name) {
		b_string_append_cstrf(out, F_GREEN "-%c" F_RESET, cmd->b_short_name);
		r++;
	}

	if (r == 2) {
		b_string_append_cstr(out, ", ");
	}

	if (cmd->b_long_name) {
		b_string_append_cstrf(out, F_GREEN "--%s" F_RESET, cmd->b_long_name);
		r++;
	}
}

static void get_command_description(struct b_command *cmd, b_string *out)
{
	b_string_append_cstr(out, cmd->b_description);
}

static void print_options_list(struct b_command *cmd)
{
	b_fprintf(OUTPUT_STREAM, "\n" F_YELLOW "OPTIONS:" F_RESET "\n");

	size_t desb_margin = 0;
	b_string *str = b_string_create();

	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_opt) {
		struct b_command_option *opt
			= b_unbox(struct b_command_option, it.entry, opt_entry);
		if (!opt) {
			continue;
		}

		b_string_clear(str);
		z__b_get_option_usage_string(opt, CMD_STR_COLOUR, str);
		size_t len = b_string_get_size(str, B_STRLEN_IGNORE_ESC) + 4;

		if (len > desb_margin) {
			desb_margin = len;
		}
	}

	b_paragraph_format format = {};
	format.p_flags = B_PARAGRAPH_DONT_INDENT_FIRST_LINE;
	format.p_left_margin = desb_margin + 4;
	format.p_right_margin = 4;

	size_t i = 0;
	b_queue_foreach (&it, &cmd->b_opt) {
		struct b_command_option *opt
			= b_unbox(struct b_command_option, it.entry, opt_entry);
		if (!opt) {
			continue;
		}

		b_string_clear(str);
		z__b_get_option_usage_string(opt, CMD_STR_COLOUR, str);

		fputs("    ", OUTPUT_STREAM);
		b_fputs(b_string_ptr(str), OUTPUT_STREAM);
		unsigned int len = b_string_get_size(str, B_STRLEN_IGNORE_ESC) + 4;
		while (len < format.p_left_margin) {
			fputc(' ', OUTPUT_STREAM);
			len++;
		}

		b_string_clear(str);
		z__b_get_option_description(opt, str);

		b_print_paragraph(b_string_ptr(str), OUTPUT_STREAM, &format);
	}

	b_string_release(str);
}

static void print_args_list(struct b_command *cmd)
{
	b_fprintf(OUTPUT_STREAM, "\n" F_YELLOW "ARGS:" F_RESET "\n");

	size_t desb_margin = 0;
	b_string *str = b_string_create();

	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_arg) {
		struct b_command_arg *arg
			= b_unbox(struct b_command_arg, it.entry, arg_entry);

		b_string_clear(str);
		z__b_get_arg_usage_string(arg, true, str);
		size_t len = b_string_get_size(str, B_STRLEN_IGNORE_ESC) + 4;

		if (len > desb_margin) {
			desb_margin = len;
		}
	}

	b_paragraph_format format = {};
	format.p_flags = B_PARAGRAPH_DONT_INDENT_FIRST_LINE;
	format.p_left_margin = desb_margin + 4;
	format.p_right_margin = 4;

	size_t i = 0;
	b_queue_foreach (&it, &cmd->b_arg) {
		struct b_command_arg *arg
			= b_unbox(struct b_command_arg, it.entry, arg_entry);

		b_string_clear(str);
		z__b_get_arg_usage_string(arg, true, str);

		fputs("    ", OUTPUT_STREAM);
		b_fputs(b_string_ptr(str), OUTPUT_STREAM);
		unsigned int len = b_string_get_size(str, B_STRLEN_IGNORE_ESC) + 4;
		while (len < format.p_left_margin) {
			fputc(' ', OUTPUT_STREAM);
			len++;
		}

		b_string_clear(str);
		z__b_get_arg_description(arg, str);

		b_print_paragraph(b_string_ptr(str), OUTPUT_STREAM, &format);
	}

	b_string_release(str);
}

static void print_commands_list(struct b_command *cmd)
{
	b_fprintf(OUTPUT_STREAM, "\n" F_YELLOW "COMMANDS:" F_RESET "\n");

	size_t desb_margin = 0;
	b_string *str = b_string_create();

	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_subcommands) {
		struct b_command *sub
			= b_unbox(struct b_command, it.entry, b_entry);

		b_string_clear(str);
		get_command_string(sub, str);
		size_t len = b_string_get_size(str, B_STRLEN_IGNORE_ESC) + 4;

		if (len > desb_margin) {
			desb_margin = len;
		}
	}

	b_paragraph_format format = {};
	format.p_flags = B_PARAGRAPH_DONT_INDENT_FIRST_LINE;
	format.p_left_margin = desb_margin + 4;
	format.p_right_margin = 4;

	size_t i = 0;
	b_queue_foreach (&it, &cmd->b_subcommands) {
		struct b_command *sub
			= b_unbox(struct b_command, it.entry, b_entry);

		b_string_clear(str);
		get_command_string(sub, str);

		fputs("    ", OUTPUT_STREAM);
		b_fputs(b_string_ptr(str), OUTPUT_STREAM);
		unsigned int len = b_string_get_size(str, B_STRLEN_IGNORE_ESC) + 4;
		while (len < format.p_left_margin) {
			fputc(' ', OUTPUT_STREAM);
			len++;
		}

		b_string_clear(str);
		get_command_description(sub, str);

		b_print_paragraph(b_string_ptr(str), OUTPUT_STREAM, &format);
	}

	b_string_release(str);
}

struct b_command *b_command_get_subcommand_with_name(
	struct b_command *cmd, const char *name)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_subcommands) {
		struct b_command *subcmd
			= b_unbox(struct b_command, it.entry, b_entry);
		if (!subcmd || !subcmd->b_name) {
			continue;
		}

		if (!strcmp(subcmd->b_name, name)) {
			return subcmd;
		}
	}

	return NULL;
}

struct b_command *b_command_get_subcommand_with_long_name(
	struct b_command *cmd, const char *long_name)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_subcommands) {
		struct b_command *subcmd
			= b_unbox(struct b_command, it.entry, b_entry);
		if (!subcmd || !subcmd->b_long_name) {
			continue;
		}

		if (!strcmp(subcmd->b_name, long_name)) {
			return subcmd;
		}
	}

	return NULL;
}

struct b_command *b_command_get_subcommand_with_short_name(
	struct b_command *cmd, char short_name)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_subcommands) {
		struct b_command *subcmd
			= b_unbox(struct b_command, it.entry, b_entry);
		if (!subcmd || !subcmd->b_short_name) {
			continue;
		}

		if (subcmd->b_short_name == short_name) {
			return subcmd;
		}
	}

	return NULL;
}

struct b_command_option *b_command_get_option_with_long_name(
	struct b_command *cmd, const char *long_name)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_opt) {
		struct b_command_option *opt
			= b_unbox(struct b_command_option, it.entry, opt_entry);
		if (!opt || !opt->opt_long_name) {
			continue;
		}

		if (!strcmp(opt->opt_long_name, long_name)) {
			return opt;
		}
	}

	return NULL;
}

struct b_command_option *b_command_get_option_with_short_name(
	struct b_command *cmd, char short_name)
{
	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_opt) {
		struct b_command_option *opt
			= b_unbox(struct b_command_option, it.entry, opt_entry);
		if (!opt || !opt->opt_long_name) {
			continue;
		}

		if (opt->opt_short_name == short_name) {
			return opt;
		}
	}

	return NULL;
}

static void print_usage(struct b_command *cmd)
{
	b_paragraph_format format = {};
	format.p_left_margin = format.p_right_margin = 4;

	b_fprintf(OUTPUT_STREAM, F_YELLOW "USAGE:" F_RESET "\n");

	if (b_queue_empty(&cmd->b_usage)) {
		b_string *usage = z__b_command_default_usage_string(cmd, NULL);
		b_print_paragraph(b_string_ptr(usage), OUTPUT_STREAM, &format);
		b_string_release(usage);
		return;
	}

	b_string *str = b_string_create();
	b_queue_iterator it;
	b_queue_foreach (&it, &cmd->b_usage) {
		struct b_command_usage *usage
			= b_unbox(struct b_command_usage, it.entry, u_entry);
		if (!usage) {
			break;
		}

		b_string_clear(str);
		get_usage_string(cmd, usage, str);
		b_print_paragraph(b_string_ptr(str), OUTPUT_STREAM, &format);
	}

	b_string_release(str);
}

static void print_help(struct b_command *cmd)
{
	b_paragraph_format format = {};

	if (!cmd->b_parent) {
		b_fprintf(OUTPUT_STREAM, F_GREEN "%s" F_RESET "\n", cmd->b_name);
	}

	if (cmd->b_description) {
		b_print_paragraph(cmd->b_description, OUTPUT_STREAM, &format);
	}

	b_fprintf(OUTPUT_STREAM, "\n");

	print_usage(cmd);

	if (!b_queue_empty(&cmd->b_opt)) {
		print_options_list(cmd);
	}

	if (!b_queue_empty(&cmd->b_arg)) {
		print_args_list(cmd);
	}

	if (!b_queue_empty(&cmd->b_subcommands)) {
		print_commands_list(cmd);
	}
}

static int execute_command(struct b_command *cmd, struct b_arglist *args)
{
	if (!cmd) {
		return -1;
	}

	if (cmd->b_flags & B_COMMAND_SHOW_HELP_BY_DEFAULT
	    && b_arglist_get_count(args, B_COMMAND_INVALID_ID, B_COMMAND_INVALID_ID)
	               == 0) {
		print_help(cmd);
		return 0;
	}

	if (b_arglist_get_count(args, B_COMMAND_OPTION_HELP, B_COMMAND_INVALID_ID)
	    > 0) {
		print_help(cmd);
		return 0;
	}

	if (cmd->b_callback) {
		return cmd->b_callback(cmd, args, NULL);
	}

	return -1;
}

static b_status add_subcommand(struct b_command *parent, struct b_command *child)
{
	b_queue_push_back(&parent->b_subcommands, &child->b_entry);
	child->b_parent = parent;

	return B_SUCCESS;
}

static int resolve_command_parents(struct b_btree *commands)
{
	struct b_btree_iterator it;
	b_btree_foreach (&it, commands) {
		struct b_command *cmd = b_unbox(struct b_command, it.node, b_node);

		if (cmd->b_parent_id == B_COMMAND_INVALID_ID) {
			continue;
		}

		cmd->b_parent = get_command(commands, cmd->b_parent_id);
		if (!cmd->b_parent) {
			return -1;
		}

		add_subcommand(cmd->b_parent, cmd);
	}

	return 0;
}

int b_command_dispatch(unsigned int cmd_id, int argc, const char **argv)
{
	if (resolve_command_parents(&command_list) != 0) {
		return -1;
	}

	struct b_command *cmd = get_command(&command_list, cmd_id);
	if (!cmd) {
		return -1;
	}

	struct b_arglist *args = b_arglist_create();
	b_status status = b_arglist_parse(args, &cmd, argc, argv);
	if (B_ERR(status)) {
		b_arglist_destroy(args);
		return -1;
	}

	int ret = execute_command(cmd, args);

	b_arglist_destroy(args);
	// print_help(cmd);
	return ret;
}