#define _POSIX_C_SOURCE 200809L
#include "pu2.out.c"
#include <stdio.h>
#include <time.h>

typedef char *string;

string read_file(string path, size_t *out_size) {
    FILE  *f = fopen(path, "rb");
    string buf;
    size_t size;
    if (!f) return NULL;
    fseek(f, 0, SEEK_END);
    size = ftell(f);
    fseek(f, 0, SEEK_SET);
    buf = malloc(size + 1);
    if (!buf) {
        fclose(f);
        return NULL;
    }
    if (fread(buf, 1, size, f) != size) {
        free(buf);
        fclose(f);
        return NULL;
    }
    buf[ size ] = '\0';
    fclose(f);
    if (out_size) *out_size = size;
    return buf;
}

error_debug(string fail_point, string code, int sz) {
    int error_pos = (int) (fail_point - code);
    int i;

    int line = 1, col = 1;
    for (i = 0; i < error_pos && i < sz; i++) {
        if (code[ i ] == '\n') {
            line++;
            col = 1;
        } else {
            col++;
        }
    }

    int line_start = error_pos;
    while (line_start > 0 && code[ line_start - 1 ] != '\n') { line_start--; }

    int line_end = error_pos;
    while (line_end < sz && code[ line_end ] != '\n' && code[ line_end ] != '\0') { line_end++; }

    int tmp          = line;
    int lineno_width = 1;
    while (tmp >= 10) {
        lineno_width++;
        tmp /= 10;
    }

    fprintf(
        stderr,
        "\x1b[1m\x1b[31merror:\x1b[0m \x1b[1massembly failed\x1b[0m "
        "at \x1b[36mline %d\x1b[0m, \x1b[36mcolumn %d\x1b[0m "
        "(\x1b[33mposition %d\x1b[0m):\n",
        line,
        col,
        error_pos);

    fprintf(
        stderr,
        " %*d | \x1b[1m%.*s\x1b[0m\n",
        lineno_width,
        line,
        line_end - line_start,
        code + line_start);

    /* print caret line aligned under the column */
    fprintf(stderr, " %*s | ", lineno_width, "");
    for (i = line_start; i < error_pos; i++) { fputc((code[ i ] == '\t') ? '\t' : ' ', stderr); }
    fprintf(stderr, "\x1b[1m\x1b[31m^\x1b[0m \x1b[90mparsing stopped here\x1b[0m\n");
}

long powi(long x, unsigned n) {
    long p = x, r = 1;

    while (n > 0) {
        if (n & 1) r *= p;
        p *= p, n /= 2;
    }

    return r;
}

emu() {
    while (1) {
        pu2();
        int e_code = m[ 0 ];

        if (e_code != 2) return e_code;

        // handle hardware interrupts
        int int_code = m[ 1 ];

        // 0: output string
        if (int_code == 0) {
            int ptr  = m[ 2 ];
            int size = m[ 3 ], i = 0;

            while (size--) putchar(m[ ptr++ ]);
        }

        // 1: input character
        else if (int_code == 1) {
            m[ 2 ] = getchar();
        }

        // 2: arithmetics unit
        else if (int_code == 2) {
            int op   = m[ 2 ];
            int val1 = m[ 3 ];
            int val2 = m[ 4 ];

            m[ 2 ] = 0;

            if (op == 0) {
                // multiply
                m[ 3 ] = val1 * val2;
            } else if (op == 1) {
                // divide
                m[ 3 ] = val1 / val2;
            } else if (op == 2) {
                // modulo
                m[ 3 ] = val1 % val2;
            } else if (op == 3) {
                // left shift
                m[ 3 ] = val1 << val2;
            } else if (op == 4) {
                // right shift
                m[ 3 ] = val1 >> val2;
            } else if (op == 5) {
                // pow
                if (val1 == 2) m[ 3 ] = val1 << val2;
                else { m[ 3 ] = powi(val1, val2); }
            }

            else {
                // unknown op, set m[2]
                m[ 2 ] = 1;
            }
        }

        else {
            // unknown interrupt. what to do?
        }

        m[ 0 ] = 0;
    }
}

skip_space(string *I) {
    while ((**I == '\n' || isspace(**I)) && (**I != 0)) (*I)++;
}

// IR parser for pu2's as2
// clang-format off
struct {
    int hlt, imm, lod, str, swp, add, sub, and, not, acc, a2u, u2a, b2u, u2b;
} opcode = {
    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13
};
// clang-format on

typedef unsigned char byte;
typedef byte bool;
const bool true  = 1;
const bool false = 0;

// use this instead of strdup'ing and allocating extra storage!
typedef struct {
    size_t len;
    string ref;
} slice;

typedef struct {
    slice key;
    slice value;
} entry;

typedef struct {
    entry data;
    byte  arg_count;
} macro;

clean_slice_ref(slice *val) { free(val->ref); }

#define static_len(G) (sizeof(G) / sizeof((G)[ 0 ]))

#define emit(OP, VAL)                                                                              \
    ({                                                                                             \
        *Code++ = (OP);                                                                            \
        *Code++ = (VAL);                                                                           \
    })
#define store()                                                                                    \
    ({                                                                                             \
        emit(opcode.acc, 1);                                                                       \
        emit(opcode.str, 0);                                                                       \
    })
#define run_err(result, ...)                                                                       \
    ({                                                                                             \
        ERR_code = init_I;                                                                         \
        fprintf(stderr, "\n");                                                                     \
        error_debug((result), ERR_code, strlen(ERR_code));                                         \
        fprintf(stderr, "\n\x1b[1m\x1b[31mir2 error:\x1b[0m ");                                    \
        fprintf(stderr, __VA_ARGS__);                                                              \
        fprintf(stderr, "\n\n");                                                                   \
        exit(1);                                                                                   \
    })

#define checkskip(char, ...)                                                                       \
    ({                                                                                             \
        if (!*I) { run_err(I, __VA_ARGS__); }                                                      \
        if (*I++ != (char) ) { run_err(I - 1, __VA_ARGS__); }                                      \
    })
#define scanword()                                                                                 \
    ({                                                                                             \
        slice VAR = { 0, I };                                                                      \
        while (isalnum(*I) || *I == '_') VAR.len++, I++;                                           \
        VAR;                                                                                       \
    })
#define memslice slice __attribute__((cleanup(clean_slice_ref)))
#define readd()  ((used == new.len) ? (new.ref = realloc(new.ref, new.len *= 2)) : 0)
#define scanvalue()                                                                                \
    ({ /* value scanner! numbers, characters, strings. */                                          \
       int value;                                                                                  \
       if (isdigit(*I)) {                                                                          \
           value = strtol(I, &I, 10);                                                              \
       } else if (*I == '\'') {                                                                    \
           I++;                                                                                    \
           value = *I++;                                                                           \
           checkskip('\'', "couldn't find closing quote on character");                            \
       } else if (*I == '"') {                                                                     \
           string str_start = ++I;                                                                 \
           bool   T         = false;                                                               \
           size_t str_len   = 0;                                                                   \
           while (*I != '"' || T) {                                                                \
               if (*I == '\\') {                                                                   \
                   T = !T;                                                                         \
                   if (!T && I[ 1 ] == '"') str_len--;                                             \
               }                                                                                   \
               str_len++;                                                                          \
               I++;                                                                                \
           }                                                                                       \
           I++;                                                                                    \
           str_len++;                                                                              \
           int *dest_start = Data -= str_len;                                                      \
           int  y;                                                                                 \
           for (y = 0; y < str_len - 1; y++) { dest_start[ y ] = str_start[ y ]; }                 \
           Data[ -1 ] = 0;                                                                         \
           value      = Data - m;                                                                  \
       } else {                                                                                    \
           run_err(I, "expected value (number, character, string)");                               \
       }                                                                                           \
       value;                                                                                      \
    })

#define op(mode)   emit(opcode.mode, 0)
#define swap()     emit(opcode.swp, 0)
#define imm(value) emit(opcode.imm, value)
#define load()     emit(opcode.lod, 0)
#define imm_in_mem(__mem, __imm)                                                                   \
    ({                                                                                             \
        imm(__imm);                                                                                \
        swap();                                                                                    \
        imm(__mem);                                                                                \
        store();                                                                                   \
    })
#define mov_mem_mem(__mem1, __mem2)                                                                \
    ({                                                                                             \
        imm(__mem2);                                                                               \
        swap();                                                                                    \
        load();                                                                                    \
        swap();                                                                                    \
        imm(__mem1);                                                                               \
        store();                                                                                   \
    })
#define mem_in_a(__mem)                                                                            \
    ({                                                                                             \
        op(b2u);                                                                                   \
        imm(__mem);                                                                                \
        swap();                                                                                    \
        load();                                                                                    \
        op(u2b);                                                                                   \
    })
#define perform_op(operation)                                                                      \
    ({                                                                                             \
        mem_in_a(mem2_index);                                                                      \
        swap();                                                                                    \
        mem_in_a(mem_index);                                                                       \
        emit(opcode.operation, 0);                                                                 \
        swap();                                                                                    \
        imm(mem_index);                                                                            \
        store();                                                                                   \
    })

char   T;
string ERR;
string ERR_code;
int   *Code = &m[ 4000 ];
int   *Data = &m[ static_len(m) - 1 ]; // grows downwards
_ir2(string I, macro macros[ 1 << 9 ]) {
    string init_I = I;
    ERR           = "all good";

    while (T = *I) {
        skip_space(&I);

        if (T == '/' && I[ 1 ] == '/') { // comment
            I += 2;
            while (*I != '\n' && *I != '\0') I++;
        }

        // handle assignments!
        if (T == '@') {
            I++;
            int mem_index = strtol(I, &I, 10);
            skip_space(&I);
            char extra_op = '=';
            if (*I == '+' || *I == '-' || *I == '&' || *I == '!') extra_op = *I++;
            checkskip('=', "expected assigment after memory blob (hint: only +-&! are supported)");
            skip_space(&I);

            if (*I == '@') {
                I++;
                if (!isdigit(*I)) { run_err(I, "expected a digit after @"); }

                int mem2_index = strtol(I, &I, 10);

                if (extra_op == '=') {
                    mov_mem_mem(mem_index, mem2_index);
                } else if (extra_op == '+') {
                    perform_op(add);
                } else if (extra_op == '-') {
                    perform_op(sub);
                } else if (extra_op == '&') {
                    perform_op(and);
                } else if (extra_op == '!') {
                    // @1 != @2;
                    mem_in_a(mem2_index);
                    emit(opcode.not, 0);
                    swap();
                    imm(mem_index);
                    store();
                }
            } else {
                int imm = scanvalue();

                if (extra_op == '=') {
                    imm_in_mem(mem_index, imm);
                } else if (extra_op == '+') {
                    imm(imm);
                    swap();
                    mem_in_a(mem_index);
                    emit(opcode.add, 0);
                    store();
                } else if (extra_op == '-') {
                    imm(imm);
                    swap();
                    mem_in_a(mem_index);
                    emit(opcode.sub, 0);
                    store();
                } else if (extra_op == '&') {
                    imm(imm);
                    swap();
                    mem_in_a(mem_index);
                    emit(opcode.and, 0);
                    store();
                } else if (extra_op == '!') {
                    // @1 != 5;
                    imm(imm);
                    emit(opcode.not, 0);
                    swap();
                    imm(mem_index);
                    store();
                }
            }

            skip_space(&I);
            checkskip(';', "expected semicolon");
            continue;
        }

        // register macros
        if (strncmp("macro", I, 5) == 0 && isspace(I[ 5 ])) {
            I += 6;
            skip_space(&I);
            // parse macro stuff

            if (!isalpha(*I) && *I != '_') {
                run_err(I - 1, "expected letter or underscore at macro name");
            }
            slice name = scanword();

            skip_space(&I);

            checkskip('(', "expected open paren in macro defintion");

            skip_space(&I);

            string before    = I;
            int    arg_count = strtol(I, &I, 10);
            if (I == before) {
                run_err(I, "expected number (argument count) in macro definition parens");
            }

            skip_space(&I);

            checkskip(
                ')',
                "expected close paren in macro definition (hint: macro defintions only take an "
                "argument count)");

            skip_space(&I);

            checkskip('{', "expected open curly bracket in macro body");

            // find open space in macros
            int x;
            for (x = 0; x < (1 << 9); x++) {
                macro g = macros[ x ];
                if (g.data.key.ref == NULL) { goto DONE; }
            }

            run_err(I - 1, "registered macro count exceeded 1<<9 (how?)");

        DONE:;

            skip_space(&I); // might aswell skip whitespace for funsies.
            macro *g          = &macros[ x ];
            g->data.key       = name;
            g->data.value.ref = I;
            g->data.value.len = 0;
            g->arg_count      = arg_count;

            int dep = 0;
            while (dep != 0 || *I != '}') {
                if (*I == '{') dep++;
                if (*I == '}') dep--;
                g->data.value.len++;
                I++;

                if (*I == 0) {
                    checkskip('}', "expected close curly bracket in macro body (met EOF instead)");
                }
            }

            I++;
            continue;
        }

        // emit exit calls
        if (strncmp("exit", I, 4) == 0) {
            I += 4;
            checkskip('(', "expected open paren on exit blob");

            string old   = I;
            int    value = strtol(I, &I, 10);
            if (old == I && *I != ')') {
                run_err(I - 1, "invalid input to exit blob, expected number");
            }

            skip_space(&I);
            checkskip(')', "expected open paren on exit blob");
            skip_space(&I);
            checkskip(';', "expected semicolon after call");

            // actually emit stuff
            if (old != I) emit(opcode.imm, value);
            emit(opcode.hlt, 0);

            continue;
        }

        // emit system
        if (strncmp("system", I, 6) == 0 && !isalnum(I[ 7 ])) {
            I += 6;
            skip_space(&I);
            checkskip('(', "expected open paren on system blob");
            skip_space(&I);
            checkskip(
                ')',
                "expected closing paren on system blob (hint: system does not take any arguments)");
            skip_space(&I);
            checkskip(';', "expected semicolon after system blob");

            emit(opcode.imm, 2);
            emit(opcode.swp, 0);
            emit(opcode.imm, 0);
            store();
            emit(opcode.hlt, 0);

            continue;
        }

        // macro call, maybe?
        if (isalpha(*I)) {
            slice name = scanword();
            skip_space(&I);

            checkskip('(', "expected open paren after alphanumeric word (assuming this is a blob)");
            skip_space(&I);

            slice  arguments[ 1 << 7 ];
            size_t arg_count = 0;

            string fail_I = I;

            while (*I != ')') {
                if (!*I) run_err(fail_I, "reached EOF whilst looking for closing paren on blob");

                checkskip('[', "expected bracket in blob argument (macro-type)");
                arguments[ arg_count ].len = 0;
                arguments[ arg_count ].ref = I;

                int dep = 0;
                while (dep != 0 || *I != ']') {
                    if (*I == '[') dep++;
                    if (*I == ']') dep--;
                    arguments[ arg_count ].len++;
                    I++;

                    if (*I == 0) {
                        checkskip(
                            ']', "expected closing bracket in macro argument (met EOF instead)");
                    }
                }

                skip_space(&I);
                checkskip(']', "expected closing bracket in macro arguments");

                skip_space(&I);
                if (*I != ',' && *I != ')') {
                    run_err(I, "expected comma or closing paren in blob");
                }

                if (*I == ',') {
                    I++;
                    skip_space(&I);
                }

                arg_count++;
            }

            fail_I = I;

            checkskip(')', "expected closing paren after blob arguments");
            skip_space(&I);
            checkskip(';', "expected semicolon after blob");

            // done reading blob! check it and parse macro :)
            // find index
            size_t i;
            for (i = 0; i < (1 << 9); i++) {
                macro cur = macros[ i ];
                if (cur.data.key.len != name.len) continue;
                if (strncmp(name.ref, cur.data.key.ref, name.len) == 0) goto FOUND;
            }

            run_err(I, "could not find a macro by that name");

        FOUND:;
            if (arg_count != macros[ i ].arg_count) {
                run_err(
                    fail_I,
                    "macro argument count mismatch (expected %i, got %li)",
                    macros[ i ].arg_count,
                    arg_count);
            }
            slice  target     = macros[ i ].data.value;
            string N          = target.ref;
            string oldN       = N;
            memslice new      = { target.len, malloc(target.len) };
            size_t used       = 0;
            bool   needs_scan = true;

            while (needs_scan) {
                needs_scan = false;

                while (*N) {
                    if (N - oldN >= target.len) break;

                    if (*N == '/' && isdigit(N[ 1 ])) {
                        N++;
                        string old   = N - 1;
                        size_t index = strtol(N, &N, 10);

                        if (index > arg_count) {
                            N--;
                            goto END;
                        }

                        if (*N++ != '/') {
                            N = old + 1;
                            while (*N != '/') {
                                new.ref[ used++ ] = *N++;
                                readd();
                            }
                            goto END;
                        }

                        needs_scan = true;
                        slice  rep = arguments[ index - 1 ];
                        string val = rep.ref;
                        for (i = 0; i < rep.len; i++) {
                            new.ref[ used++ ] = *val++;
                            readd();
                        }

                        continue;
                    }

                END:
                    new.ref[ used++ ] = *N++;
                    readd();
                }
            }

            new.ref[ used ] = 0;

            _ir2(new.ref, macros);

            continue;
        }
    }
}

ir2(string I) {
    macro macros[ 1 << 9 ] = {0};
    _ir2(I, macros);
}

string op_names[] = { "hlt", "imm", "lod", "str", "swp", "add", "sub",
                      "and", "not", "acc", "a2u", "u2a", "b2u", "u2b" };

unsigned long long rdtsc(void) {
    unsigned hi, lo;
    __asm__ __volatile__("rdtsc" : "=a"(lo), "=d"(hi));
    return ((unsigned long long) hi << 32) | lo;
}

double bench_cycle(int instr) {
    int i;
    for (i = 0; i < S - 2; i += 2) {
        m[ i ]     = instr;
        m[ i + 1 ] = 5;
    }
    m[ S - 2 ] = 0;
    m[ S - 1 ] = 0;

    m[ 7 ] = 0;
    c      = 0;

    clock_t start = clock();
    pu2();
    clock_t end = clock();

    double seconds = (double) (end - start) / CLOCKS_PER_SEC;
    printf("time: %.6f s\n", seconds);
    return (double) m[ 7 ] / seconds / 1e6;
}

double bench2_cycle(int instr) {
    unsigned long long start, end;
    double             vminst_count, cpu_cycles;
    int                i;
    for (i = 0; i < S - 2; i += 2) {
        m[ i ]     = instr;
        m[ i + 1 ] = 5;
    }
    m[ S - 2 ] = 0;
    m[ S - 1 ] = 0;
    m[ 7 ]     = 0;
    c          = 0;
    start      = rdtsc();
    pu2();
    end          = rdtsc();
    cpu_cycles   = (double) (end - start);
    vminst_count = (double) m[ 7 ];
    printf("pu2/real: %f\n", cpu_cycles / vminst_count);
    return cpu_cycles / vminst_count;
}

benchmark(double (*target)(int)) {
    double results[ 13 ];
    double sum   = 0;
    int    count = 0;
    int    i;
    for (i = 1; i < 14; i++) {
        printf("running: %s\n", op_names[ i ]);
        results[ i ] = target(i);
        sum += results[ i ];
        count++;
    }
    sum /= count;

    unsigned long long start, end;
    start = rdtsc();
    sleep(1);
    end = rdtsc();

    float clock_speed = (end - start) / 1e6;
    printf("results sum: %f\n", sum);
    if (target == bench_cycle) {
        printf("clock speed ~ %.6f MHz\n", clock_speed);
        printf("instructions per cycle: %f\n", clock_speed / sum);
    }
}

long nstime() {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_nsec;
}

main(int argc, char **argv) {
    if (argc == 1) {
        return printf(
            "\nusage:\n"
            "    %s <file_name>              open and compile ir2-compliant code\n"
            "    %s --bench                  benchmark pu2\n\n",
            argv[ 0 ],
            argv[ 0 ]);
    }

    if (strcmp(argv[ 1 ], "--bench") == 0) {
        benchmark(bench_cycle);
        benchmark(bench2_cycle);
        exit(0);
    }

    char  *filename = argv[ 1 ];
    size_t file_size;
    string contents = read_file(filename, &file_size);
    if (contents == NULL) {
        perror("unable to open file");
        exit(1);
    }

    size_t o_size;
    long before = nstime();
    long cycles_bef = rdtsc();
    ir2(contents);
    long cycles_aft = rdtsc();
    long after = nstime();
    m[ 7 ] = 4000;

    // fprintf(stderr, "assembly took %li nanoseconds (%li cycles)\n", (after - before), (cycles_aft - cycles_bef));
    free(contents);

    size_t res_size = Code - &m[ 4000 ];
    // --- BYTES PRINTING ---

    // for(i = 0; i < res_size; i++) {
    //     printf("%X%X ", m[4000+i]>>4, m[4000+i]&0b1111);
    //     if(i % 8 == 7) printf("\n");
    // }
    // fflush(stdout);
    // printf("\n\n");

    // --- OPCODES PRINTING ---

    // for (i = 0; i < res_size; i++) {
    //     int op   = m[ 4000 + i++ ];
    //     int data = m[ 4000 + i ];

    //     printf("%s(%i);\n", op_names[ op ], data);
    // }
    // printf("\n\n");

    emu();
}