#include "lib/mlrutil.h"
#include "lib/string_builder.h"
#include "containers/lhmss.h"
#include "containers/sllv.h"
#include "mapping/mappers.h"
#include "cli/argparse.h"

#define RENAME_SB_ALLOC_LENGTH 16

typedef struct _regex_pair_t {
	regex_t regex;
	char*   replacement;
} regex_pair_t;

typedef struct _mapper_rename_state_t {
	lhmss_t* pold_to_new;
	sllv_t*  pregex_pairs;
	string_builder_t* psb;
	int      do_gsub;
} mapper_rename_state_t;

static sllv_t*   mapper_rename_process(lrec_t* pinrec, context_t* pctx, void* pvstate);
static sllv_t*   mapper_rename_regex_process(lrec_t* pinrec, context_t* pctx, void* pvstate);
static void      mapper_rename_free(void* pvstate);
static mapper_t* mapper_rename_alloc(lhmss_t* pold_to_new, int do_regexes, int do_gsub);
static void      mapper_rename_usage(FILE* o, char* argv0, char* verb);
static mapper_t* mapper_rename_parse_cli(int* pargi, int argc, char** argv);

// ----------------------------------------------------------------
mapper_setup_t mapper_rename_setup = {
	.verb = "rename",
	.pusage_func = mapper_rename_usage,
	.pparse_func = mapper_rename_parse_cli
};

// ----------------------------------------------------------------
static void mapper_rename_usage(FILE* o, char* argv0, char* verb) {
	fprintf(o, "Usage: %s %s [options] {old1,new1,old2,new2,...}\n", argv0, verb);
	fprintf(o, "Renames specified fields.\n");
	fprintf(o, "Options:\n");
	fprintf(o, "-r         Treat old field  names as regular expressions. \"ab\", \"a.*b\"\n");
	fprintf(o, "           will match any field name containing the substring \"ab\" or\n");
	fprintf(o, "           matching \"a.*b\", respectively; anchors of the form \"^ab$\",\n");
	fprintf(o, "           \"^a.*b$\" may be used. New field names may be plain strings,\n");
	fprintf(o, "           or may contain capture groups of the form \"\\1\" through\n");
	fprintf(o, "           \"\\9\". Wrapping the regex in double quotes is optional, but\n");
	fprintf(o, "           is required if you wish to follow it with 'i' to indicate\n");
	fprintf(o, "           case-insensitivity.\n");
	fprintf(o, "-g         Do global replacement within each field name rather than\n");
	fprintf(o, "           first-match replacement.\n");
	fprintf(o, "Examples:\n");
	fprintf(o, "%s %s -r 'Date_[0-9]+,Date,'  Rename all such fields to be \"Date\"\n", argv0, verb);
	fprintf(o, "%s %s -r '\"Date_[0-9]+\",Date' Same\n", argv0, verb);
	fprintf(o, "%s %s -r 'Date_([0-9]+).*,\\1' Rename all such fields to be of the form 20151015\n", argv0, verb);
	fprintf(o, "%s %s -r '\"name\"i,Name'       Rename \"name\", \"Name\", \"NAME\", etc. to \"Name\"\n", argv0, verb);
}

static mapper_t* mapper_rename_parse_cli(int* pargi, int argc, char** argv) {
	int do_regexes = FALSE;
	int do_gsub = FALSE;

	char* verb = argv[(*pargi)++];

	ap_state_t* pstate = ap_alloc();
	ap_define_true_flag(pstate, "-r", &do_regexes);
	ap_define_true_flag(pstate, "-g", &do_gsub);

	if (!ap_parse(pstate, verb, pargi, argc, argv)) {
		mapper_rename_usage(stderr, argv[0], verb);
		return NULL;
	}

	if ((argc - *pargi) < 1) {
		mapper_rename_usage(stderr, argv[0], verb);
		return NULL;
	}

	if (do_gsub)
		do_regexes = TRUE;

	slls_t* pnames = slls_from_line(argv[*pargi], ',', FALSE);
	if ((pnames->length % 2) != 0) {
		mapper_rename_usage(stderr, argv[0], verb);
		return NULL;
	}
	lhmss_t* pold_to_new = lhmss_alloc();
	for (sllse_t* pe = pnames->phead; pe != NULL; pe = pe->pnext->pnext) {
		lhmss_put(pold_to_new, pe->value, pe->pnext->value);
	}

	*pargi += 1;
	return mapper_rename_alloc(pold_to_new, do_regexes, do_gsub);
}

// ----------------------------------------------------------------
static mapper_t* mapper_rename_alloc(lhmss_t* pold_to_new, int do_regexes, int do_gsub) {
	mapper_t* pmapper = mlr_malloc_or_die(sizeof(mapper_t));

	mapper_rename_state_t* pstate = mlr_malloc_or_die(sizeof(mapper_rename_state_t));

	if (do_regexes) {
		pmapper->pprocess_func = mapper_rename_regex_process;
		pstate->pold_to_new    = NULL;
		pstate->pregex_pairs   = sllv_alloc();

		for (lhmsse_t* pe = pold_to_new->phead; pe != NULL; pe = pe->pnext) {
			char* regex_string = pe->key;
			char* replacement  = pe->value;

			regex_pair_t* ppair = mlr_malloc_or_die(sizeof(regex_pair_t));
			regcomp_or_die_quoted(&ppair->regex, regex_string, 0);
			ppair->replacement = replacement;
			sllv_add(pstate->pregex_pairs, ppair);
		}

		pstate->psb     = sb_alloc(RENAME_SB_ALLOC_LENGTH);
		pstate->do_gsub = do_gsub;
	} else {
		pmapper->pprocess_func = mapper_rename_process;
		pstate->pold_to_new    = pold_to_new;
		pstate->pregex_pairs   = NULL;
		pstate->psb            = NULL;
		pstate->do_gsub        = FALSE;
	}
	pmapper->pfree_func = mapper_rename_free;

	pmapper->pvstate = (void*)pstate;
	return pmapper;
}

static void mapper_rename_free(void* pvstate) {
	mapper_rename_state_t* pstate = (mapper_rename_state_t*)pvstate;
	lhmss_free(pstate->pold_to_new);
}

// ----------------------------------------------------------------
static sllv_t* mapper_rename_process(lrec_t* pinrec, context_t* pctx, void* pvstate) {
	if (pinrec != NULL) {
		mapper_rename_state_t* pstate = (mapper_rename_state_t*)pvstate;
		for (lhmsse_t* pe = pstate->pold_to_new->phead; pe != NULL; pe = pe->pnext) {
			char* old_name = pe->key;
			char* new_name = pe->value;
			if (lrec_get(pinrec, old_name) != NULL) {
				lrec_rename(pinrec, old_name, new_name, FALSE);
			}
		}
		return sllv_single(pinrec);
	}
	else {
		return sllv_single(NULL);
	}
}

static sllv_t* mapper_rename_regex_process(lrec_t* pinrec, context_t* pctx, void* pvstate) {
	if (pinrec != NULL) {
		mapper_rename_state_t* pstate = (mapper_rename_state_t*)pvstate;

		for (sllve_t* pe = pstate->pregex_pairs->phead; pe != NULL; pe = pe->pnext) {
			regex_pair_t* ppair = pe->pvdata;
			regex_t* pregex = &ppair->regex;
			char* replacement = ppair->replacement;
			for (lrece_t* pf = pinrec->phead; pf != NULL; pf = pf->pnext) {
				int matched = FALSE;
				int all_captured = FALSE;
				char* old_name = pf->key;
				char* new_name = pstate->do_gsub
					? regex_gsub(old_name, pregex, pstate->psb, replacement, &matched, &all_captured)
					:  regex_sub(old_name, pregex, pstate->psb, replacement, &matched, &all_captured);

				if (matched)
					lrec_rename(pinrec, old_name, new_name, TRUE);
			}
		}

		return sllv_single(pinrec);
	}
	else {
		return sllv_single(NULL);
	}
}
