#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#pragma clang diagnostic ignored "-Wparentheses"
#pragma clang diagnostic ignored "-Wreturn-type"
#pragma clang diagnostic ignored "-Wempty-body"

// pu2!
// maybe compacter? idk
// last cpu: 167 chars

// NOTE: if you want the cpu to
// crash or do something when an invalid opcode is met,
// replace the `8` in the code with whatever behavior you desire.
// you can use ({...}) to insert normal code, or use __builtin_trap() for a SIGTRAP

// clang-format off
#define S 1<<20
#define H !--c
#define J while
m[S],c,v,a,b,u;pu2(){J(c=m[m[7]++],v=m[m[7]++],c){H?a=v:H?a=m[b]:H?u?m[a]=b:0:H?v=a,a=b,b=v:H?a+=b:H?a-=b:H?a=a&&b:H?a=!a:H?u=1:8;u=v;}}
// clang-format on

// there's just memory. no, really
// ok actually there's two registers called a and b
// 0 HLT: halt
// 1 IMM: a = imm
// 2 LOD: a = m[b]
// 3 STR: m[a] = b (only happens if acc != 0)
// 4 SWP: swap a and b
// 5 ADD: a += b
// 6 SUB: a -= b
// 7 AND: a &&= b
// 8 NOT: a = !a
// 9 ACC: u = imm

// NOTE: Memory indexes less than 1024 should be considered as temporary/compute memory, and
// programs should be written at index 1024 and upwards.
//
// NOTE: Memory index 7 is the program
// counter, in which the current instruction is read from m[m[7]] at each cycle

// -------------------------

// yeehaw baby it's time for as2
// E is the error code loc. zero it before calling as2
// Y is macro storage. extend as needed
// set O for output before running. keep in mind, it gets modified whilst as2 is running.
//    therefore, i recommend you store a copy somewhere!

// clang-format off
char*C="hltimmlodstrswpaddsubandnotacc",*Y[1<<10],*E,*L,i;*O;as2(char*I){J(c=*I++){if(c==35){J(*++I-10);c=*I++;}L=C;(!isspace(c)?c-39?c-125?c-123?c-64?({i=0,I--;J(i++,*L)strncmp(L,I,3)?L+=3:({I+=3;*O++=i-1;*O++=i-2?0:strtol(I,&I,10);goto Q;});E=I;return 1;Q:;}):as2(Y[strtol(I,&I,10)]):({Y[strtol(I,&I,10)]=I+1;J(*I++-125);}):({break;}):({*O++=1;*O++=*I++;}):0);}}
// clang-format on

#undef S
#undef H
#undef J

#ifndef __WRAPPER

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++;
        }
    }

    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);

    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++; }

    fprintf(stderr, "  \x1b[1m");
    for (i = line_start; i < line_end; i++) { fputc(code[ i ], stderr); }
    fprintf(stderr, "\x1b[0m\n");

    fprintf(stderr, "  ");
    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");
}

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;

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

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

        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; int imm; int lod; int str; int swp; int add; int sub; int and; int not; int acc;
} opcode = {
    0, 1, 2, 3, 4, 5, 6, 7, 8,9
};

#define emit(OP, VAL) ({ *(*O)++ = (OP); *(*O)++ = (VAL); })
#define store() ({ emit(opcode.acc, 1); emit(opcode.str, 0); })
#define run_err(result)                                                                        \
    ({                                                                                         \
        printf("\n");                                                                          \
        error_debug((result), ERR_code, strlen(ERR_code));                                     \
        printf("\nir2 error message: %s\n", ERR);                                              \
        printf("\n");                                                                          \
        exit(1);                                                                               \
    })

#define checkskip(char, message)                                                               \
    ({                                                                                         \
        if (!*I) {                                                                             \
            ERR = (message); run_err(I);                                                       \
        }                                                                                      \
        if (*I++ != (char) ) {                                                                 \
            ERR = (message); ERR_code=init_I; run_err(I-1);                                    \
        }                                                                                      \
    })
#define scanword() ({ slice VAR = { 0, I}; while (isalnum(*I)) VAR.len++, I++; VAR; })
#define memslice slice __attribute__((cleanup(clean_slice_ref)))
#define readd() ({ if(used == new.len) new.ref = realloc(new.ref, new.len *= 2); })
// 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;

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

T;
string ERR;
string ERR_code;
_ir2(string I, int **O, entry 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++;
        }

        // handle assignments!
        if (T == '@') {
            I++;
            int mem_index = strtol(I, &I, 10);
            skip_space(&I);
            checkskip('=', "expected assigment after memory blob");
            skip_space(&I);

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

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

                emit(opcode.imm, mem2_index);
                emit(opcode.swp, 0);
                emit(opcode.lod, 0);
                emit(opcode.swp, 0);
                emit(opcode.imm, mem_index);
                store();
            } else if (isdigit(*I) || *I == '\'') {
                int imm;
                if (*I == '\'') {
                    I++;
                    imm = *I++;
                    checkskip('\'', "couldn't find closing quote on character");
                } else {
                    imm = strtol(I, &I, 10);
                }

                emit(opcode.imm, imm);
                emit(opcode.swp, 0);
                emit(opcode.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)) {
                ERR = "expected letter at macro name";
                run_err(I - 1);
            }
            slice name = scanword();

            skip_space(&I);

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

            skip_space(&I);

            checkskip(
                ')',
                "expected close paren in macro definition (hint: macros do not take named "
                "arguments)");

            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++) {
                entry g = macros[ x ];
                if (g.key.ref == NULL) { goto DONE; }
            }

            ERR = "macro count exceeded 1<<9";
            run_err(I - 1);

        DONE:;

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

            int dep = 0;
            while (dep != 0 || *I != '}') {
                if (*I == '{') dep++;
                if (*I == '}') dep--;
                g->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 != ')') {
                ERR = "invalid input to exit blob, expected number";
                run_err(I - 1);
            }

            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");

            _ir2("@0 = 2; exit();", O, macros);

            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;

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

                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 != ')') {
                    ERR = "expected comma or closing paren in blob";
                    run_err(I);
                }

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

                arg_count++;
            }

            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++) {
                entry cur = macros[ i ];
                if (cur.key.len != name.len) continue;
                if (strncmp(name.ref, cur.key.ref, name.len) == 0) goto FOUND;
            }

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

        FOUND:;
            slice  target     = macros[ i ].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, O, macros);

            continue;
        }
    }
}

ir2(string I, int **O) {
    entry macros[ 1 << 9 ];
    _ir2(I, O, macros);
}

main() {
    int   *out = &m[ 1024 ];
    size_t o_size;
    ir2(read_file("code.i2", &o_size), &out);
    m[ 7 ] = 1024;

    // size_t res_size = out - &m[ 1024 ];
    // for(i = 0; i < res_size; i++) {
    //     printf("%i ", m[1024+i]);
    // }
    // fflush(stdout);
    // printf("\n\n");

    emu();
}

main2(int argc, char **argv) {
    if (argc == 1) {
        return printf(
            "\nusage:\n"
            "    %s <file_name>\n\n",
            argv[ 0 ]);
    }

    char  *filename = argv[ 1 ];
    size_t file_size;
    char  *contents = read_file(filename, &file_size);

    E = 0;
    O = &m[ 1024 ];
    as2(contents);

    if (E != 0) {
        error_debug(E, contents, file_size);
        exit(1);
    }

    m[ 7 ] = 1024;
    emu();
}

#endif