// TODO rules validator cli tool

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

#include "djb2.h"
#include "av_rules.h"

// TODO move to string.c
static char *trim(char *s) {
	while (isspace(*s)) {
		s++;
	}

	if (*s == '\0') {
		return s;
	}

	char *end = s + strlen(s) - 1;
	while (end > s && isspace(*end)) {
		*end-- = '\0';
	}

	return s;
}

rule_action_t parse_action(const char *s) {
	if (strcmp(s, "block") == 0)      return RULE_BLOCK;
	if (strcmp(s, "allow") == 0)      return RULE_ALLOW;
	if (strcmp(s, "quarantine") == 0) return RULE_QUARANTINE;
	if (strcmp(s, "info")       == 0) return RULE_INFORMATIONAL;

	return RULE_ALLOW; // default fallback
}

static condition_type_t parse_condition_type(const char *tok, char **out_id, size_t *out_offset, bool *has_offset) {
    condition_type_t type = COND_TYPE_REQUIRED;
    const char *p = tok;

    if (*p == '!') {
        type = COND_TYPE_NEGATED;
        p++;
    } else if (*p == '?') {
        type = COND_TYPE_OPTIONAL;
        p++;
    }

    char *colon = strchr(p, ':');
    if (colon) {
        *colon = '\0';
        *out_id = strdup(p);
        *out_offset = (size_t)atoi(colon + 1);
        *has_offset = true;
    } else {
        *out_id = strdup(p);
        *out_offset = 0;
        *has_offset = false;
    }

    return type;
}

char *strip_quotes(char *s) {
	size_t len = strlen(s);

	if (len >= 2 && s[0] == '"' && s[len - 1] == '"') {
		s[len - 1] = '\0'; // ending quote
		return s + 1; // leading quote
	}

	return s;
}

static int hex_char_to_nibble(char c) {
	if ('0' <= c && c <= '9') return c - '0';
	if ('a' <= c && c <= 'f') return c - 'a' + 10;
	if ('A' <= c && c <= 'F') return c - 'A' + 10;

	return -1;
}

size_t hex_to_bytes(const char *hex, uint8_t *out, size_t maxlen) {
	size_t len = 0;

	while (*hex && *(hex + 1) && len < maxlen) {
		while (*hex && isspace((unsigned char)*hex)) {
			hex++;
		}

		int hi = hex_char_to_nibble(*hex++);
		int lo = hex_char_to_nibble(*hex++);
		if (hi < 0 || lo < 0) {
			break;
		}

		out[len++] = (hi << 4) | lo;
	}

	return len;
}

void unescape_string(char *s) {
	char *src = s;
	char *dst = s;

	while (*src) {
		if (*src == '\\') {
			src++;
			switch(*src) {
			case 'n': *dst++ = '\n'; break;
			case 'r': *dst++ = '\r'; break;
			case 't': *dst++ = '\t'; break;
			case '0': *dst++ = '\0'; break;
			case '\\': *dst++ = '\\'; break;
			case '"': *dst++ = '\"'; break;
			case 'x':
				int hi = hex_char_to_nibble(*(++src));
				int lo = hex_char_to_nibble(*(++src));
				if (hi >= 0 && lo >= 0) {
					*dst++ = (hi << 4) | lo;
				}
				break;
			default:
				*dst++ = *src;
				break;
			}

			src++;
		} else {
			*dst++ = *src++;
		}
	}

	*dst = '\0';
}

int load_rules(const char *path, rule_set_t *out_ruleset) {
	FILE *f = fopen(path, "r");
	if (!f) {
		return -1;
	}

	char line[1024];
	rule_t *current_rule = NULL;

	while (fgets(line, sizeof(line), f)) {
		char *comment = strchr(line, '#');
		if (comment) {
			*comment = '\0';
		}

		char *s = trim(line);
		if (*s == '\0') {
			continue;
		}

		char *tok = strtok(s, " \t\n");
		if (!tok) {
			continue;
		}

		if (strcmp(tok, "rule") == 0) {
			if (out_ruleset->rule_count >= MAX_RULES) {
				break;
			}

			current_rule = &out_ruleset->rules[out_ruleset->rule_count++];
			memset(current_rule, 0, sizeof(rule_t));

			char *rule_id = strtok(NULL, " \t\n");
			char *action_str = strtok(NULL, " \t\n");
			if (!rule_id || !action_str) {
				continue;
			}

			current_rule->id = strdup(rule_id);
			current_rule->action = parse_action(action_str);

			while ((tok = strtok(NULL, " \t\n")) &&
				   current_rule->condition_count < MAX_RULE_CONDITIONS) {
				char *id = NULL;
				size_t offset = 0;
				bool has_offset = false;

				condition_type_t ctype = parse_condition_type(tok, &id, &offset, &has_offset);

				rule_condition_t *cond = &current_rule->conditions[current_rule->condition_count++];
				cond->pattern_id = id;
				cond->offset = offset;
				cond->has_offset = has_offset;
				cond->type = ctype;
			}
		} else if (strcmp(tok, "string") == 0) {
			char *id = strtok(NULL, " \t\n");
			char *value = strtok(NULL, "");

			if (!id || !value) {
				// malformed string declaration
				continue;
			}

			if (*value == '"') {
				// if value is quoted, remove quotes and unescape
				value = strip_quotes(value);
				unescape_string(value);
			}

			// TODO this silently fails
			pattern_table_add(&out_ruleset->patterns, id, (uint8_t *)value, strlen(value));
		} else if (strcmp(tok, "pattern") == 0) {
			char *id = strtok(NULL, " \t\n");
			char *hexstr = strtok(NULL, "\t\n");

			if (!id || !hexstr) {
				// malformed hex pattern
				continue;
			}

			uint8_t bytes[512]; // max pattern size
			size_t byte_len = hex_to_bytes(hexstr, bytes, sizeof(bytes));
			if (byte_len == 0) {
				continue;
			}

			// TODO this silently fails
			pattern_table_add(&out_ruleset->patterns, id, bytes, byte_len);
		}
	}

	fclose(f);

	return 0;
}

bool evaluate_rule(const rule_t *rule, pattern_table_t *patterns) {
	for (size_t i = 0; i < rule->condition_count; i++) {
		const rule_condition_t *cond = &rule->conditions[i];
		pattern_t *p = pattern_table_find(patterns, cond->pattern_id);

		bool matched = false;

		if (p) {
			if (cond->has_offset) {
				for (size_t j = 0; j < p->match_count; j++) {
					if (p->match_offsets[j] == cond->offset) {
						matched = true;
						break;
					}
				}
			} else {
				matched = (p->match_count > 0);
			}
		}

		switch (cond->type) {
		case COND_TYPE_REQUIRED:
			if (!matched) {
				return false;
			}
			break;
		case COND_TYPE_NEGATED:
			if (matched) {
				return false;
			}
			break;
		case COND_TYPE_OPTIONAL:
			break;
		}
	}

	return true;
}

void dump_rule(const rule_t *r) {
    printf("rule %s %s\n", r->id,
           r->action == RULE_BLOCK ? "block" :
           r->action == RULE_ALLOW ? "allow" :
		   r->action == RULE_QUARANTINE ? "quarantine" :
		   "info");

    for (size_t i = 0; i < r->condition_count; i++) {
        const rule_condition_t *c = &r->conditions[i];

		char offset_buf[16];
		if (c->has_offset) sprintf(offset_buf, "%zu", c->offset);
        printf("  %s%s%s%s\n",
               c->type == COND_TYPE_NEGATED ? "!" :
               c->type == COND_TYPE_OPTIONAL ? "?" : "",
               c->pattern_id,
               c->has_offset ? ":" : "",
               c->has_offset ? offset_buf : "");
    }
}

void free_rules(rule_set_t *rules) {
	for (size_t i = 0; i < rules->rule_count; i++) {
		rule_t *r = &rules->rules[i];
		free((void *)r->id);

		for (size_t j = 0; j < r->condition_count; j++) {
			free((void *)r->conditions[j].pattern_id);
		}
	}
}

pattern_t *pattern_table_find(pattern_table_t *table, const char *id) {
	uint32_t h = djb2(id) % PATTERN_TABLE_SIZE;

	for (pattern_bucket_t *b = table->buckets[h]; b; b = b->next) {
		if (strcmp(b->id, id) == 0) {
			return b->entry;
		}
	}

	return NULL;
}

int pattern_table_add(pattern_table_t *table, const char *id, const uint8_t *bytes, size_t len) {
	if (table->count >= MAX_PATTERNS) {
		return -1;
	}

	if (pattern_table_find(table, id)) {
		// id already exists. skip adding it.
		return -1;
	}

	pattern_t *p = &table->patterns[table->count++];
	p->id = strdup(id);
	p->bytes = malloc(len);
	if (!p->bytes) {
		return -1;
	}

	memcpy(p->bytes, bytes, len);
	p->len = len;
	p->match_count = 0;

	uint32_t h = djb2(id) % PATTERN_TABLE_SIZE;

	pattern_bucket_t *b = malloc(sizeof(*b));
	if (!b) {
		return -1;
	}

	b->id = p->id;
	b->entry = p;
	b->next = table->buckets[h];
	table->buckets[h] = b;

	return 0;
}

void pattern_table_clear_matches(pattern_table_t *table) {
    for (size_t i = 0; i < table->count; i++) {
        table->patterns[i].match_count = 0;
		memset(table->patterns[i].match_offsets, 0, sizeof(table->patterns[i].match_offsets));
    }
}

void pattern_table_free(pattern_table_t *table) {
	for (size_t i = 0; i < PATTERN_TABLE_SIZE; i++) {
		pattern_bucket_t *b = table->buckets[i];

		while (b) {
			pattern_bucket_t *next = b->next;
			free((void *)b->id);
			free(b);
			b = next;
		}
	}

	for (size_t i = 0; i < table->count; i++) {
		free(table->patterns[i].bytes);
	}

	memset(table, 0, sizeof(*table));
}