#include <setjmp.h>
#include <getopt.h>
#include <stdlib.h>
#include <stdio.h>
#include <inttypes.h>

#include "xtest.h"
#include <assert.h>

#ifdef XTEST_PRNG

#include <time.h>
#include <math.h>
#include "pcg_basic.h"

#endif


#ifdef XTEST_PRNG
#define PRNG_SEED_LENGTH 22
pcg32_random_t rng;
typedef char prng_seed[PRNG_SEED_LENGTH + 1];
void xtest_rand_init(prng_seed seed);
#endif

struct xtest_assert_info {
    const char *file;
    int line;
    const char *func;
    const char *expr;
};

struct xtest_assert_info assertion_info;

#ifndef XTEST_MAX_GROUP_NESTING
#define XTEST_MAX_GROUP_NESTING 16
#endif

#ifndef XTEST_TEST_NAME_MAX_LENGTH
#define XTEST_TEST_NAME_MAX_LENGTH 1023
#endif

#ifndef XTEST_INDENT
#define XTEST_INDENT 2
#endif


enum output_format {
    format_default,
    format_subunit,
};

enum output_format output_format;

const char *name_filter;

const char *group_nesting[XTEST_MAX_GROUP_NESTING];
int group_nesting_pos = 0;

jmp_buf xtest_jmp;

#ifndef XTEST_ASSERT_MESSAGE_MAX_LEN
#define XTEST_ASSERT_MESSAGE_MAX_LEN 1023
#endif
char assert_message[XTEST_ASSERT_MESSAGE_MAX_LEN + 1];
const char *skip_reason;

int xtest_indent = 0;

int expecting_assertion = 0;
int num_tests = 0;
int successful_tests = 0;
int failed_tests = 0;
int skipped_tests = 0;

_Bool list_tests = 0;


#define USAGE_BASE "[-?|--help] [-l|--list] [-m|--match filter] [-f|--output-format format]"
#define USAGE_PRNG "[-s|--seed prng-seed]"

#define USAGE_HELP_LIST     "-l|--list                  list the available tests"
#define USAGE_HELP_MATCH    "-m|--match filter          only run/list tests whose name matches the filter"
#define USAGE_HELP_FORMAT   "-f|--output-format format  specify the output format; supported formats are \"default\" and \"subunit\""
#define USAGE_HELP_SEED     "-s|--seed prng-seed        specify a seed for the prng"
#define USAGE_HELP_HELP     "-?|--help                  show this message"
#ifdef XTEST_PRNG
#define USAGE "Usage: %s " USAGE_BASE " " USAGE_PRNG "\n  " USAGE_HELP_LIST "\n  " USAGE_HELP_MATCH "\n  " USAGE_HELP_FORMAT "\n  " USAGE_HELP_SEED "\n  " USAGE_HELP_HELP
#else
#define USAGE "Usage: %s " USAGE_BASE "\n  " USAGE_HELP_LIST "\n  " USAGE_HELP_MATCH "\n  " USAGE_HELP_FORMAT "\n  " USAGE_HELP_HELP
#endif

static _Noreturn void usage(const char *progname) {
    printf(USAGE, progname);
    exit(1);
}

static _Bool matches(const char *name, const char *filter) {
    // eat up the matching parts of name and filter until the first * or end of either string
    while (*name == *filter && *name != '\0' && *filter != '\0' && *filter != '*') {
        name++;
        filter++;
    }
    // if the next character in filter is not a * then there is no match (otherwise the previous loop would have continued)
    if (*filter != '\0' && *filter != '*') return 0;

    while (1) {
        // if the filter ends before the name then they don't match
        if (*filter == '\0') {
            return *name == '\0';
        }
        // the filter points to a * at this point because everything else was eaten up either before the loop or in the previous iteration of the loop
        assert (*filter == '*');
        while (*filter == '*') filter++; // eat up consecutive *s
        size_t part_len = strcspn(filter, "*");
        if (part_len == 0) return 1; // filter ends with a * so every remaining input in name is accepted
        char part[part_len + 1];
        strncpy(part, filter, part_len);
        part[part_len] = '\0'; // ensure null-termination since strncpy does not null-terminate on its own
        name = strstr(name, part);
        if (name == NULL) return 0;
        name += part_len;
        filter += part_len;
    }
}

void xtest_init(int argc, char **argv) {
    int c;
    int list_flag = 0;
    int help_flag = 0;
    const char *format = NULL;
#ifdef XTEST_PRNG
    const char *seed_arg = NULL;
#endif

    const char *progname = argv[0];

    struct option long_options[] =
        {
            {"list", no_argument, &list_flag, 1}, // list the tests instead of running them
            {"help", no_argument, &help_flag, 1}, // list the tests instead of running them
            {"match", required_argument, 0, 'm'}, // only run tests that match the provided filter
            {"output-format", required_argument, 0, 'f'}, // select output format
#ifdef XTEST_PRNG
            {"seed", required_argument, 0, 's'}, // provide a prng seed
#endif
            {0, 0, 0, 0}
        };
    while (1) {
        int option_index = 0;
        c = getopt_long(argc, argv, "?lm:f:", long_options, &option_index);
        if (c == -1) break;
        switch (c) {
            case 0:
                assert(long_options[option_index].flag != 0);
                break;
            case 'l':
                list_flag = 1;
                break;
            case 'm':
                name_filter = optarg;
                break;
            case 'f':
                format = optarg;
                break;
#ifdef XTEST_PRNG
            case 's':
                seed_arg = optarg;
                break;
#endif
            case '?':
                help_flag = 1;
                break;
            default:
                assert(0);
        }
    }

    if (help_flag) {
        usage(progname);
    }

    if (list_flag)
        list_tests = 1;

    if (format == NULL) {
        format = "default";
    }
    if (strcmp(format, "default") == 0) {
        output_format = format_default;
    } else if (strcmp(format, "subunit") == 0) {
        output_format = format_subunit;
    } else {
        printf("invalid output format '%s', allowed formats are 'default' and 'subunit'\n", format);
        exit(1);
    }

#ifdef XTEST_PRNG
    char seed[PRNG_SEED_LENGTH + 1] = {0};
    if (seed_arg != NULL) {
        snprintf(seed, PRNG_SEED_LENGTH + 1, "%s", seed_arg);
    }
    xtest_rand_init(seed);
    printf("[PRNG seed: %s]\n============\n", seed);
#endif
}


int xtest_complete() {
    if (!list_tests) {
        if (output_format == format_default) {
            printf("============\nTotal: %d, Failed: %d, Skipped: %d\n", num_tests, failed_tests, skipped_tests);
        }
        fflush(stdout);
        return failed_tests != 0 ? 1 : 0;
    } else {
        return 0;
    }
}

void subunit_message(const char *name, const char *message, const char *details) {
    printf("%s: %s", message, name);
    if (details == NULL) {
        printf("\n");
    } else {
        printf(" [\n%s%s]\n", details, details[strlen(details)] != '\n' ? "\n" : "");
    }
}

void xtest_internal_run(xtest_test_fn fn, const char *name, xtest_param *params, xtest_setup_fn setup,
                        xtest_teardown_fn teardown) {
    expecting_assertion = 0;

    char full_name[XTEST_TEST_NAME_MAX_LENGTH + 1] = "";
    size_t name_len = 0;
    for (int i = 0; i < group_nesting_pos; ++i) {
        name_len += snprintf(full_name + name_len, XTEST_TEST_NAME_MAX_LENGTH - name_len, "%s/", group_nesting[i]);
        assert(name_len <= XTEST_TEST_NAME_MAX_LENGTH);
    }
    name_len += snprintf(full_name + name_len, XTEST_TEST_NAME_MAX_LENGTH - name_len, "%s", name);
    assert(name_len <= XTEST_TEST_NAME_MAX_LENGTH);

    if (name_filter != NULL) {
        if (!matches(full_name, name_filter)) return;
    }

    if (list_tests) {
        printf("%s\n", full_name);
        return;
    }

    num_tests += 1;


    void *fixture = NULL;

    if (output_format == format_subunit) {
        subunit_message(full_name, "test", NULL);
    } else {
        printf("%*s%s: ", xtest_indent, "", name);
    }
    if (setup != NULL) {
        int jmpres = setjmp(xtest_jmp);
        if (jmpres == 0) {
            setup(&fixture);
        } else if (jmpres == 3) {
            goto skipped;
        } else {
            snprintf(assert_message, 1024, "assertion failure in setup function: %s in %s (%s:%d)", assertion_info.expr,
                     assertion_info.func, assertion_info.file, assertion_info.line);
            goto failed;
        }
    }

    int jmpres = setjmp(xtest_jmp);
    if (jmpres == 0) {
        if (params == NULL) {
            expecting_assertion = 0;
            fn(fixture, NULL);
        } else {
            int num_params = 0;
            while (params[num_params].name != NULL) num_params++;
            int validx[num_params];
            int valmax[num_params];
            void *current_vals[num_params];
            for (int i = 0; i < num_params; ++i) {
                validx[i] = 0;
                valmax[i] = 0;
                while ((*params[i].values)[valmax[i]] != NULL) valmax[i]++;
            }
            while (1) {
                for (int i = 0; i < num_params; ++i) {
                    current_vals[i] = (*params[i].values)[validx[i]];
                }
                expecting_assertion = 0;
                fn(fixture, current_vals);
                _Bool carry = 1;
                for (int i = num_params - 1; i >= 0; --i) {
                    validx[i] += 1;
                    if (validx[i] >= valmax[i]) {
                        validx[i] = 0;
                    } else {
                        carry = 0;
                        break;
                    }
                }
                if (carry) break;
            }
        }

        if (expecting_assertion) { // an assertion failure in the tested code was expected but didn't occur
            snprintf(assert_message, 1024, "expected an assertion failure in the tested code");
            goto failed;
        } else {
            goto success;
        }
    } else if (jmpres == 1) { // a testing assertion failed
        goto failed;
    } else if (jmpres == 2) { // an assertion in the tested code failed
        if (expecting_assertion) {
            goto success;
        } else {
            snprintf(assert_message, 1024, "unexpected assertion failure in tested code: %s in %s (%s:%d)",
                     assertion_info.expr, assertion_info.func, assertion_info.file, assertion_info.line);
            goto failed;
        }
    } else if (jmpres == 3) {
        goto skipped;
    }

  success:

    if (output_format == format_subunit) {
        subunit_message(full_name, "success", NULL);
    } else {
        printf("SUCCESS\n");
    }
    successful_tests += 1;
    goto cleanup;
  failed:

    if (output_format == format_subunit) {
        subunit_message(full_name, "failure", assert_message);
    } else {
        printf("FAILED:\n%*s  %s\n", xtest_indent, "", assert_message);
    }
    failed_tests += 1;
    goto cleanup;
  skipped:
    skipped_tests += 1;

    if (output_format == format_subunit) {
        subunit_message(full_name, "skipped", skip_reason);
    } else {
        if (skip_reason != NULL) {
            printf("SKIPPED: %s\n", skip_reason);
        } else {
            printf("SKIPPED\n");
        }
    }
    goto cleanup;

  cleanup:
    if (teardown != NULL && fixture != NULL) {
        teardown(fixture);
    }
}


/// this function replaces the standard library assert and jumps back to the xtest_internal_run invocation that is running the test
/// \param file file in which the assertion happened
/// \param line line on which the assertion happened
/// \param func function in which the assertion happened
/// \param expr stringified expression of the assertion
void xtest_internal_assert(const char *file, int line, const char *func, const char *expr) {
    assertion_info = (struct xtest_assert_info) {
        .file = file, .line = line, .func = func, .expr = expr
    };
    longjmp(xtest_jmp, 2);
}


#define FAIL_ASSERT_FORMAT_1 "assertion failed: %s, expected %s"
#define FAIL_ASSERT_FORMAT_2 " but got "
#define FAIL_ASSERT_FORMAT_3 " (%s:%d)"
#define FAIL_ASSERT_FORMAT(fmt) FAIL_ASSERT_FORMAT_1 fmt FAIL_ASSERT_FORMAT_2 fmt FAIL_ASSERT_FORMAT_3
#define CAST_PTR(type, ptr) *(type *)(ptr)
#define PRINT_ASSERT_MESSAGE(fmt, type) snprintf(assert_message, XTEST_ASSERT_MESSAGE_MAX_LEN, FAIL_ASSERT_FORMAT(fmt), \
expression, invert ? "not " : "", CAST_PTR(type, expected), CAST_PTR(type, actual), file, line)

void xtest_fail_assert(const char *expression, const char *file, int line, void *expected,
                       void *actual, _Bool invert, enum xtest_type type) {
    switch (type) {
        case xtest_type_bool:
            snprintf(assert_message, XTEST_ASSERT_MESSAGE_MAX_LEN, FAIL_ASSERT_FORMAT("%s"), expression,
                     invert ? "not " : "",
                     CAST_PTR(_Bool, expected) ? "true" : "false", CAST_PTR(_Bool, actual) ? "true" : "false", file,
                     line);
            break;
        case xtest_type_char:
            PRINT_ASSERT_MESSAGE("'%c'", char);
            break;
        case xtest_type_short:
            PRINT_ASSERT_MESSAGE("%i", short);
            break;
        case xtest_type_unsigned_short:
            PRINT_ASSERT_MESSAGE("%u", unsigned short);
            break;
        case xtest_type_int:
            PRINT_ASSERT_MESSAGE("%i", int);
            break;
        case xtest_type_unsigned_int:
            PRINT_ASSERT_MESSAGE("%u", unsigned int);
            break;
        case xtest_type_long:
            PRINT_ASSERT_MESSAGE("%li", long);
            break;
        case xtest_type_unsigned_long:
            PRINT_ASSERT_MESSAGE("%lu", unsigned long);
            break;
        case xtest_type_long_long:
            PRINT_ASSERT_MESSAGE("%lli", long long);
            break;
        case xtest_type_unsigned_long_long:
            PRINT_ASSERT_MESSAGE("%llu", unsigned long long);
            break;
        case xtest_type_float:
            PRINT_ASSERT_MESSAGE("%f", float);
            break;
        case xtest_type_double:
            PRINT_ASSERT_MESSAGE("%f", double);
            break;
        case xtest_type_long_double:
            PRINT_ASSERT_MESSAGE("%Lf", long double);
            break;
        case xtest_type_string:
            PRINT_ASSERT_MESSAGE("\"%s\"", char*);
            break;
        case xtest_type_void_pointer:
            PRINT_ASSERT_MESSAGE("%p", void*);
            break;
        case xtest_type_other:
        default:
            snprintf(assert_message, XTEST_ASSERT_MESSAGE_MAX_LEN, FAIL_ASSERT_FORMAT("%s"), expression,
                     invert ? "not " : "",
                     "<value of unknown type>", "<value of unknown type>", file, line);
            break;
    }
    longjmp(xtest_jmp, 1);
}

void
xtest_assert_float(double expected, double actual, int precision, _Bool invert, const char *expression,
                   const char *file,
                   int line) {

    double epsilon = 1.0;
    for (int i = 0; i < precision; ++i) epsilon /= 10.0;
    double diff = expected - actual;
    if (diff < 0) diff = -diff;
    _Bool equals = diff < epsilon;
    _Bool failed = invert == equals; // invert = true: fail if equals, invert = false: fail if not equals
    if (!failed) return;
    snprintf(assert_message, 1024,
             "assertion failed: %s, expected %s%g but got %g (difference %g, precision %d) (%s:%d)",
             expression, invert ? "not " : "", expected, actual, diff, precision, file, line);
    longjmp(xtest_jmp, 1);
}

size_t find_diff_offset(const char *a, const char *b, size_t length) {
    size_t offset = 0;
    while (offset < length && *(a + offset) == *(b + offset)) offset++;
    return offset;
}

void xtest_assert_mem(const char *expected, const char *actual, size_t length, _Bool invert, const char *expression,
                      const char *file,
                      int line) {
    size_t offset = find_diff_offset(expected, actual, length);
    _Bool equals = offset == length;
    _Bool failed = invert == equals; // invert = true: fail if equals, invert = false: fail if not equals
    if (!failed) return;

    if (invert) {
        snprintf(assert_message, 1024,
                 "assertion failed: %s, expected memory to not match but did (%s:%d)",
                 expression, file, line);
    } else {
        snprintf(assert_message, 1024,
                 "assertion failed: %s, expected memory to match but did not (first different byte at offset %zu, expected 0x%2x but got 0x%2x) (%s:%d)",
                 expression, offset, (unsigned int) expected[offset], (unsigned int) actual[offset], file, line);
    }

    longjmp(xtest_jmp, 1);
}

void xtest_skip(const char *reason) {
    skip_reason = reason;
    longjmp(xtest_jmp, 3);
}

void xtest_internal_start_group(const char *name) {
    if (!list_tests) {
        if (output_format == format_default) {
            printf("%*s%s:\n", xtest_indent, "", name);
        }
    }
    group_nesting[group_nesting_pos] = name;
    group_nesting_pos += 1;
    xtest_indent += XTEST_INDENT;
}

void xtest_internal_end_group() {
    assert(group_nesting_pos > 0);
    group_nesting_pos -= 1;
    xtest_indent -= XTEST_INDENT;
}

void xtest_expect_assertion_failure() {
    expecting_assertion = 1;
}


#ifdef XTEST_PRNG


#define SEED_ENC_0 2
#define SEED_ENC_A 12
#define SEED_ENC_a 38


union seed_encode_box {
    struct {
        uint64_t state;
        uint64_t seq;
    };
    uint8_t bytes[PRNG_SEED_LENGTH * 6 / 8 + 1];
};

void encode_seed(prng_seed seed, uint64_t state, uint64_t seq) {
    union seed_encode_box box;
    memset(&box, 0, sizeof(box));
    box.state = state;
    box.seq = seq;
    uint8_t cur;
    for (int i = 0; i < PRNG_SEED_LENGTH; i += 1) {
        switch (i % 4) {
            case 0:
                cur = box.bytes[i * 6 / 8] & 0x3F;
                break;
            case 1:
                cur = (box.bytes[i * 6 / 8] >> 6) | ((box.bytes[i * 6 / 8 + 1] & 0xF) << 2);
                break;
            case 2:
                cur = (box.bytes[i * 6 / 8] >> 4) | ((box.bytes[i * 6 / 8 + 1] & 0x3) << 4);
                break;
            case 3:
                cur = (box.bytes[i * 6 / 8] >> 2);
                break;
        }
        char c;
        if (cur == 0) c = '_';
        else if (cur == 1) c = '.';
        else if (cur >= SEED_ENC_0 && cur <= SEED_ENC_0 + '9' - '0') c = (char) (cur - SEED_ENC_0 + '0');
        else if (cur >= SEED_ENC_A && cur <= SEED_ENC_A + 'Z' - 'A') c = (char) (cur - SEED_ENC_A + 'A');
        else if (cur >= SEED_ENC_a && cur <= SEED_ENC_a + 'z' - 'a') c = (char) (cur - SEED_ENC_a + 'a');
        else
            assert(0);
        seed[i] = c;
    }
    seed[PRNG_SEED_LENGTH] = '\0';
}

void decode_seed(const prng_seed seed, uint64_t *state, uint64_t *seq) {
    union seed_encode_box box;
    memset(&box, 0, sizeof(box));
    uint8_t cur;
    for (int i = 0; i < PRNG_SEED_LENGTH; i += 1) {
        char c = seed[i];
        if (c == '_') cur = 0;
        else if (c == '.') cur = 1;
        else if (c >= '0' && c <= '9') cur = SEED_ENC_0 + c - '0';
        else if (c >= 'A' && c <= 'Z') cur = SEED_ENC_A + c - 'A';
        else if (c >= 'a' && c <= 'z') cur = SEED_ENC_a + c - 'a';
        else cur = -1; // invalid seed character is just ignored and set to 0xFF so no error handling is necessary
        switch (i % 4) {
            case 0:
                box.bytes[i * 6 / 8] |= cur;
                break;
            case 1:
                box.bytes[i * 6 / 8] |= cur << 6;
                box.bytes[i * 6 / 8 + 1] |= cur >> 2;
                break;
            case 2:
                box.bytes[i * 6 / 8] |= cur << 4;
                box.bytes[i * 6 / 8 + 1] |= cur >> 4;
                break;
            case 3:
                box.bytes[i * 6 / 8] |= cur << 2;
                break;
        }
    }
    *state = box.state;
    *seq = box.seq;
}


void xtest_rand_init(prng_seed seed) {
    uint64_t state;
    uint64_t seq;
    if (seed[0] == '\0') {
        int64_t t = time(NULL);
        int shift = (int) (t & 0x3F);
        if (t >> 6 & 0x1) shift = -shift;
        t = t >> 7;

        t = t << shift | t >> (64 - shift);

        state = (intptr_t) &printf ^ (intptr_t) seed << 32 ^ t;
        seq = (intptr_t) &longjmp << 32 ^ (intptr_t) &xtest_rand_init ^ (t << 32 | t >> 32);
//        state = time(NULL) ^ (intptr_t) &printf << 32 ^ (intptr_t) &time;
//        seq = time(NULL) << 32 ^ (intptr_t) seed ^ (intptr_t) &matches << 32;
        encode_seed(seed, state, seq);
    } else {
        decode_seed(seed, &state, &seq);
    }
    pcg32_srandom_r(&rng, state, seq);
}


void xtest_rand_fill(char *buf, size_t len) {
    uint32_t rand = 0;
    for (size_t i = 0; i < len; ++i) {
        if (i % 4 == 0) {
            rand = pcg32_random_r(&rng);
        }
        buf[i] = (char) (rand >> (8 * (i % 4)));
    }
}

double xtest_rand_double() {
    return ldexp((double) (xtest_rand_64() & 0xFFFFFFFFFFFFF), -52);
}

double xtest_rand_double_range(double min, double max) {
    assert(min < max);
    return (max - min) * xtest_rand_double() + min;
}

int xtest_rand_int() {
    uint32_t uval = pcg32_random_r(&rng);
    int32_t *ival;
    ival = (int32_t *) &uval;
    return (int) *ival;
}

int xtest_rand_int_range(int min, int max) {
    int bound = max - min;
    uint32_t uval = pcg32_boundedrand_r(&rng, bound);
    int32_t *ival;
    ival = (int32_t *) &uval;
    return (int) *ival + min;
}

uint8_t xtest_rand_8() {
    return (uint8_t) pcg32_random_r(&rng);
}

uint16_t xtest_rand_16() {
    return (uint16_t) pcg32_random_r(&rng);
}

uint32_t xtest_rand_32() {
    return pcg32_random_r(&rng);

}

uint64_t xtest_rand_64() {
    uint32_t l = pcg32_random_r(&rng);
    uint32_t h = pcg32_random_r(&rng);
    return ((uint64_t) h) << 32 | l;
}

#endif