#include <stdint.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdarg.h>
#include <string.h>
#include <assert.h>


// Global defines

#define STR_SIZE 128
#define MEMORY_SIZE 1024*1024
#define EDIT_QUEUE_SIZE 10

// Memory

static char g_memory[1024*1024];
static int g_memory_index = 0;

void * alloc(int num, int size) {
    assert(g_memory_index + num*size < MEMORY_SIZE);

    void * result = g_memory + g_memory_index;
    for (int i = 0; i < num*size; i++)
        g_memory[g_memory_index+i] = 0;
    g_memory_index += num*size;
    return result;
}

#define NEW(TYPE) ((TYPE *)alloc(1, sizeof(TYPE)))
#define NEWARR(TYPE, NUM) ((TYPE *)alloc(NUM, sizeof(TYPE)))
#define NEWSTR NEWARR(char, STR_SIZE)

// getch

#ifdef _WIN32
#include <windows.h>
#include <conio.h>
#else
#include <sys/ioctl.h>
#include <termios.h>
#include <unistd.h>
#include <stdio.h>

/* reads from keypress, doesn't echo */
int getch(void)
{
    struct termios oldattr, newattr;
    int ch;
    tcgetattr( STDIN_FILENO, &oldattr );
    newattr = oldattr;
    newattr.c_lflag &= ~( ICANON | ECHO ); // no ECHO for echo(?)
    tcsetattr( STDIN_FILENO, TCSANOW, &newattr );
    ch = getchar();
    tcsetattr( STDIN_FILENO, TCSANOW, &oldattr );
    return ch;
}

/* ungets keypress */
void ungetch(int ch)
{
    struct termios oldattr, newattr;
    tcgetattr( STDIN_FILENO, &oldattr );
    newattr = oldattr;
    newattr.c_lflag &= ~( ICANON | ECHO );
    tcsetattr( STDIN_FILENO, TCSANOW, &newattr );
    ungetc(ch, stdin);
    tcsetattr( STDIN_FILENO, TCSANOW, &oldattr );
}
#endif

int
peekch() {
    int c = getch();
    ungetch(c);
    return c;
}

// VT100

#define ASCII_ESC 27

void vt100Escape(const char * str, ...) {
    va_list args;
    va_start(args, str);

    printf("%c", ASCII_ESC);
    vprintf(str, args);
}

void vt100ClearScreen() { vt100Escape("[2J"); }
void vt100CursorHome() { vt100Escape("[H"); }
void vt100CursorPos(int v, int h) { vt100Escape("[%d;%df", v, h); }
void vt100SaveCursor() { vt100Escape("7"); }
void vt100RestoreCursor() { vt100Escape("8"); }
void vt100EnableAlternateBuffer() { vt100Escape("[?1049h"); }
void vt100DisableAlternateBuffer() { vt100Escape("[?1049l"); }
void vt100EnableNegative() { vt100Escape("[7m"); }
void vt100DisableNegative() { vt100Escape("[27m"); }
void vt100ShowCursor() { vt100Escape("[?25h"); }
void vt100HideCursor() { vt100Escape("[?25l"); }
void vt100GetScreenSize(int * v, int * h) {
#ifdef _WIN32
    CONSOLE_SCREEN_BUFFER_INFO csbi;
    GetConsoleScreenBufferInfo(GetStdHandle(STD_OUTPUT_HANDLE), &csbi);
    *h = csbi.srWindow.Right - csbi.srWindow.Left + 1;
    *v = csbi.srWindow.Bottom - csbi.srWindow.Top + 1;
#else
    struct winsize w;
    ioctl(STDOUT_FILENO, TIOCGWINSZ, &w);
    *h = w.ws_col;
    *v = w.ws_row;
#endif
}

// Node

typedef enum NodeKind {
    NK_Namespace,
    NK_Struct,
    NK_Func,
    NK_VarList,
    NK_ExprList,
    NK_Var,
    NK_VarDecl,
    NK_VarType,
    NK_Type,
    NK_Body,
    NK_If,
    NK_While,
    NK_Num,
    NK_Str,
    NK_Call,
    NK_Op,

    NK_COUNT
} NodeKind;

const char *NK_STRINGS[NK_COUNT] = {
    "NK_Namespace",
    "NK_Struct",
    "NK_Func",
    "NK_VarList",
    "NK_ExprList",
    "NK_Var",
    "NK_VarDecl",
    "NK_VarType",
    "NK_Type",
    "NK_Body",
    "NK_If",
    "NK_While",
    "NK_Num",
    "NK_Str",
    "NK_Call",
    "NK_Op",
};

typedef struct Node Node;
struct Node {
    NodeKind kind;
    void * data;

    int childCount;
    Node *next, *prev, *prnt, *chld;
};

Node *NodeGetChild(Node *n, unsigned int index) {
    if (index >= n->childCount)
        return NULL;
    
    Node *result = n->chld;
    for (int i = 0; i < index; i++)
        result = result->next;
    return result;
}

void NodeAppend(Node *n1, Node *n2) {
    if (n1->childCount == 0) {
        n1->chld = n2;
    }
    else {
        Node *lastChild = NodeGetChild(n1, n1->childCount-1);

        lastChild->next = n2;
        n2->prev = lastChild;
        n2->next = NULL;
    }
    n2->prnt = n1;
    n1->childCount += 1;
}

void NodeInsert(Node *n1, Node *n2, unsigned int index) {
    if (index >= n1->childCount-1) {
        NodeAppend(n1, n2);
        return;
    }

    Node *insertAfter = NodeGetChild(n1, index);
    Node *insertBefore = NodeGetChild(n1, index+1);

    insertAfter->next = n2;
    insertBefore->prev = n2;
    n2->prev = insertAfter;
    n2->next = insertBefore;

    n1->childCount += 1;
}

Node *NodeRemove(Node *n1, Node *n2) {
    if (n1 == NULL)
        return n2;
    
    Node *prev = n2->prev;
    Node *next = n2->next;

    if (prev != NULL)
        prev->next = next;
    if (next != NULL)
        next->prev = prev;
    
    n2->prev = n2->next = n2->prnt = NULL;

    if (n2 == n1->chld)
        n1->chld = next;

    n1->childCount -= 1;
    
    if (prev == NULL)
        return n1;
    else
        return prev;
}

void NodeDraw(Node *n);
Node *g_NodeDrawSelected;

void Printf(Node *n, char* format, ...)
{
  va_list argp;
  va_start(argp, format);
  while (*format != '\0') {
    if (*format == '%') {
      format++;
      if (*format == '%') {
        putchar('%');
      } else if (*format == 'n') {
        Node *n = va_arg(argp, Node*);
        NodeDraw(n);
      } else if (*format == 's') {
        char *s = va_arg(argp, char*);
        printf("%s", s);
        if (g_NodeDrawSelected == n && s == n->data)
            printf("%c", ' ');
      } else if (*format == ',' && format[1] == 'n') {
        Node *n = va_arg(argp, Node*);
        for (int i = 0; i < n->childCount; i++) {
            if (i > 0) printf(", ");
            NodeDraw(NodeGetChild(n, i));
        }
        format++;
      } else if (*format == ';' && format[1] == 'n') {
        Node *n = va_arg(argp, Node*);
        for (int i = 0; i < n->childCount; i++) {
            if (i == 0) printf("\n");
            NodeDraw(NodeGetChild(n, i));
            printf("\n");
        }
        format++;
      } else {
        fputs("Not implemented", stderr);
      }
    } else if (*format == '_') {
        static bool highlighted = false;
        // Node *n = va_arg(argp, Node*);
        if (highlighted) {
            vt100DisableNegative();
            highlighted = false;
        }
        else if (g_NodeDrawSelected == n) {
            vt100EnableNegative();
            highlighted = true;
        }
    } else {
      putchar(*format);
    }
    format++;
  }
  va_end(argp);
}

void NodeDraw(Node *n) {
    static int indent = 0;
    #define INDENT printf("\n"); for (int i = 0; i < indent; i++) printf("  ");
    
    #define PRINTF(...) do { if (n == g_NodeDrawSelected) vt100EnableNegative(); printf(__VA_ARGS__); vt100DisableNegative(); } while(0);

    switch (n->kind) {
    case NK_Namespace: { Printf(n, "_namespace %s {_%;n_}_\n", n->data, n); break; }
    case NK_Struct:    { Printf(n, "_struct %s {_%;n_}_\n", n->data, n); break; }
    case NK_Func:      { Printf(n, "_fn %s_ %n %n", n->data, NodeGetChild(n, 0), NodeGetChild(n, 1)); break; }
    case NK_VarList:
    case NK_ExprList:  { Printf(n, "_(_%,n_)_", n); break; }
    case NK_Var:       { Printf(n, "_%s_", n->data); break; }
    case NK_VarDecl:   { Printf(n, "_%s:_ %n", n->data, NodeGetChild(n, 0)); break; }
    case NK_VarType:
    case NK_Type:      { Printf(n, "_%s%s%,n%s_", n->data, (n->childCount > 0 ? "<" : ""), n, (n->childCount > 0 ? ">" : "")); break; }
    case NK_Body:      { Printf(n, "_{_%;n_}_", n); break; }
    case NK_If:        { Printf(n, "_if_ %n %n", NodeGetChild(n, 0), NodeGetChild(n, 1)); break; }
    case NK_While:     { Printf(n, "_while_ %n %n", NodeGetChild(n, 0), NodeGetChild(n, 1)); break; }
    case NK_Num:       { Printf(n, "_%s_", n->data); break; }
    case NK_Str:       { Printf(n, "_'%s'_", n->data); break; }
    case NK_Call:      { Printf(n, "_%s_(%n)", n->data, NodeGetChild(n, 0)); break; }
    case NK_Op:        { for (int i = 0; i < n->childCount; i++) { if (n->childCount == 0 || i != 0) { Printf(n, " _%s_ ", n->data); } NodeDraw(NodeGetChild(n, i)); } break; }
    }
}

// Input

#define KEY_CTRL_C 3
#define KEY_BACKSPACE1 8
#define KEY_BACKSPACE2 127

typedef enum Input {
    IN_None,

    IN_H,
    IN_J,
    IN_K,
    IN_L,

    IN_SPACE,
    IN_D,
    IN_A,

    IN_F,
    IN_V,
    IN_T,
    IN_I,
    IN_W,
    IN_N,
    IN_S,
    IN_C,
    IN_O,

    IN_Quit,

    IN_COUNT
} Input;

const char *IN_STRINGS[IN_COUNT] = {
    "IN_None",

    "IN_H",
    "IN_J",
    "IN_K",
    "IN_L",

    "IN_SPACE",
    "IN_D",
    "IN_A",

    "IN_F",
    "IN_V",
    "IN_T",
    "IN_I",
    "IN_W",
    "IN_N",
    "IN_S",
    "IN_C",
    "IN_O",

    "IN_Quit",
};

Input InputGet() {
    int c;

    while (true) {
        c = getch();
        switch (c) {
            case 'h': return IN_H;
            case 'j': return IN_J;
            case 'k': return IN_K;
            case 'l': return IN_L;

            case ' ': return IN_SPACE;
            case 'd': return IN_D;
            
            case 'f': return IN_F;
            case 'v': return IN_V;
            case 't': return IN_T;
            case 'i': return IN_I;
            case 'w': return IN_W;
            case 'n': return IN_N;
            case 's': return IN_S;
            case 'c': return IN_C;
            case 'o': return IN_O;

            case 'q': return IN_Quit;
        }
    }
}

typedef enum InputAction {
    IA_None,

    IA_MoveLeft,
    IA_MoveUp,
    IA_MoveDown,
    IA_MoveRight,

    IA_StartEditing,
    IA_Delete,

    IA_AddNamespace,
    IA_AddStruct,
    IA_AddFunc,
    IA_AddVar,
    IA_AddVarDecl,
    IA_AddType,
    IA_AddIf,
    IA_AddWhile,
    IA_AddNum,
    IA_AddStr,
    IA_AddCall,
    IA_AddOp,

    IA_COUNT
} InputAction;

const char *IA_STRINGS[IA_COUNT] = {
    "IA_None",

    "IA_MoveLeft",
    "IA_MoveUp",
    "IA_MoveDown",
    "IA_MoveRight",

    "IA_StartEditing",
    "IA_Delete",

    "IA_AddNamespace",
    "IA_AddStruct",
    "IA_AddFunc",
    "IA_AddVar",
    "IA_AddVarDecl",
    "IA_AddType",
    "IA_AddIf",
    "IA_AddWhile",
    "IA_AddNum",
    "IA_AddStr",
    "IA_AddCall",
    "IA_AddOp",
};

typedef enum InputMode {
    IM_Normal,
    IM_Editing,
} InputMode;

void DrawInfo(InputAction actions[NK_COUNT][IN_COUNT], NodeKind nk, InputMode mode) {
    int v, h;
    vt100GetScreenSize(&v, &h);

    int line = 2;

    vt100CursorPos(line++, h-30);
    printf("%d:%d", v, h);
    vt100CursorPos(line++, h-30);

    printf("%s:%s", NK_STRINGS[nk], (mode == IM_Normal ? "" : " (editing)"));

    for (int i = IN_L+1; i < IN_COUNT; i++) {
        InputAction action = actions[nk][i];
        if (action != IA_None) {
            vt100CursorPos(line++, h-30);
            printf("%s %s", IN_STRINGS[i], IA_STRINGS[action]);
        }
    }
}

bool ValidChar(NodeKind nk, int c) {
    const char *allow = NULL;
    const char *block = NULL;
    
    switch (nk) {
    case NK_Namespace:
    case NK_Struct:
    case NK_Func:
    case NK_Var:
    case NK_VarDecl:
    case NK_VarType:
    case NK_Type:
    case NK_Str:
    case NK_Call:
    case NK_Op:
        block = "";
        break;
       
    case NK_Num:
        allow = "+-.0123456789";
        break;

    case NK_VarList:
    case NK_ExprList:
    case NK_Body:
    case NK_If:
    case NK_While:
        allow = "";
        break;
    }

    if (allow != NULL) {
        char a;
        for (int i = 0; (a = allow[i]) != '\0'; i++) {
            if (a == c) return true;
        }
        return false;
    }
    else if (block != NULL) {
        char b;
        for (int i = 0; (b = block[i]) != '\0'; i++) {
            if (b == c) return false;
        }
        return true;
    }
    else {
        return false;
    }
}

Node *PopEditQueue(Node ** q) {
    Node *result = q[0];
    for (int i = 0; i < EDIT_QUEUE_SIZE-1 && q[i] != NULL; i++) {
        q[i] = q[i+1];
    }
    return result;
}

Node *GetNode(InputAction actions[NK_COUNT][IN_COUNT]) {
    Node *result = NEW(Node); result->data = NEWSTR; result->kind = NK_Namespace;

    Node * n = result;
    Node * q[EDIT_QUEUE_SIZE] = {0};

    Input in;
    InputAction action;

    InputMode mode = IM_Normal;

    while (true) {
        vt100ClearScreen();
        vt100CursorHome();
        g_NodeDrawSelected = n;
        NodeDraw(result);
        g_NodeDrawSelected = NULL;
        DrawInfo(actions, n->kind, mode);

        if (mode == IM_Editing) {
            int c = getch();
            char *s = (char*)n->data;
            int slen = strlen(s);
            
            if (c == KEY_BACKSPACE1 || c == KEY_BACKSPACE2) {
                s[slen-1] = '\0';
            }
            else if (c == '\n' || c == '\r') {
                if (q[0] == NULL)
                    mode = IM_Normal;
                else
                    n = PopEditQueue(q);
            }
            else if (slen < STR_SIZE && ValidChar(n->kind, c)) {
                s[slen++] = (char)c;
            }
        }
        else if (mode == IM_Normal) {
            in = InputGet();

            if (in == IN_Quit)
                break;

            action = actions[n->kind][in];

            #define N(NAME, PARENT, KIND) Node *NAME = NEW(Node); NAME->kind = KIND; NodeAppend(PARENT, NAME);
            #define S(NAME, PARENT, KIND) Node *NAME = NEW(Node); NAME->kind = KIND; NAME->data = NEWSTR; NodeAppend(PARENT, NAME);

            switch (action) {
            case IA_MoveLeft:  { if (n->prev != NULL) n = n->prev; break; }
            case IA_MoveDown:  { if (n->chld != NULL) n = n->chld; break; }
            case IA_MoveUp:    { if (n->prnt != NULL) n = n->prnt; break; }
            case IA_MoveRight: { if (n->next != NULL) n = n->next; break; }
            
            case IA_StartEditing: { mode = IM_Editing; break; }
            case IA_Delete:       { n = NodeRemove(n->prnt, n); break; }

            case IA_AddNamespace: { S(n1, n, NK_Namespace) n=n1;    mode=IM_Editing; break; }
            case IA_AddStruct:    { S(n1, n, NK_Struct)    n=n1;    mode=IM_Editing; break; }
            case IA_AddFunc:      { S(n1, n, NK_Func)
                                    N(n2, n1, NK_VarList)
                                    N(n3, n1, NK_Body)     n=n1;    mode=IM_Editing; break; }
            case IA_AddVar:       { S(n1, n, NK_Var)       n=n1;    mode=IM_Editing; break; }
            case IA_AddVarDecl:   { S(n1, n, NK_VarDecl)   n=n1;
                                    S(n2, n1, NK_VarType)  q[0]=n2; mode=IM_Editing; break; }
            case IA_AddType:      { S(n1, n, NK_Type)      n=n1;    mode=IM_Editing; break; }
            case IA_AddIf:        { N(n1, n, NK_If)
                                    N(n2, n1, NK_ExprList)
                                    N(n3, n1, NK_Body)     n=n2;                     break; }
            case IA_AddWhile:     { N(n1, n, NK_While)
                                    N(n2, n1, NK_ExprList)
                                    N(n3, n1, NK_Body)     n=n2;                     break; }
            case IA_AddNum:       { S(n1, n, NK_Num)       n=n1;    mode=IM_Editing; break; }
            case IA_AddStr:       { S(n1, n, NK_Str)       n=n1;    mode=IM_Editing; break; }
            case IA_AddCall:      { S(n1, n, NK_Call)
                                    N(n2, n1, NK_ExprList) n=n1;    mode=IM_Editing; break; }
            case IA_AddOp:        { S(n1, n, NK_Op)        n=n1;    mode=IM_Editing; break; }
            }

            #undef N
            #undef S
        }
    }

    return result;
}


void SetupInputActions(InputAction actions[NK_COUNT][IN_COUNT]) {
    for (int i = 0; i < NK_COUNT; i++) {
        actions[i][IN_H] = IA_MoveLeft;
        actions[i][IN_J] = IA_MoveDown;
        actions[i][IN_K] = IA_MoveUp;
        actions[i][IN_L] = IA_MoveRight;
    }

    actions[NK_Namespace][IN_D] = IA_Delete;
    actions[NK_Namespace][IN_SPACE] = IA_StartEditing;
    actions[NK_Namespace][IN_N] = IA_AddNamespace;
    actions[NK_Namespace][IN_S] = IA_AddStruct;
    actions[NK_Namespace][IN_F] = IA_AddFunc;
    
    actions[NK_Struct][IN_D] = IA_Delete;
    actions[NK_Struct][IN_SPACE] = IA_StartEditing;
    actions[NK_Struct][IN_V] = IA_AddVarDecl;
    actions[NK_Struct][IN_F] = IA_AddFunc;

    actions[NK_Func][IN_D] = IA_Delete;
    actions[NK_Func][IN_SPACE] = IA_StartEditing;

    actions[NK_VarList][IN_V] = IA_AddVarDecl;
    
    actions[NK_ExprList][IN_V] = IA_AddVar;
    actions[NK_ExprList][IN_I] = IA_AddIf;
    actions[NK_ExprList][IN_W] = IA_AddWhile;
    actions[NK_ExprList][IN_N] = IA_AddNum;
    actions[NK_ExprList][IN_S] = IA_AddStr;
    actions[NK_ExprList][IN_C] = IA_AddCall;
    actions[NK_ExprList][IN_O] = IA_AddOp;

    actions[NK_Var][IN_D] = IA_Delete;
    actions[NK_Var][IN_SPACE] = IA_StartEditing;

    actions[NK_VarDecl][IN_D] = IA_Delete;
    actions[NK_VarDecl][IN_SPACE] = IA_StartEditing;

    actions[NK_VarType][IN_SPACE] = IA_StartEditing;
    actions[NK_VarType][IN_T] = IA_AddType;

    actions[NK_Type][IN_D] = IA_Delete;
    actions[NK_Type][IN_SPACE] = IA_StartEditing;
    actions[NK_Type][IN_T] = IA_AddType;

    actions[NK_Body][IN_F] = IA_AddFunc;
    actions[NK_Body][IN_V] = IA_AddVarDecl;
    actions[NK_Body][IN_I] = IA_AddIf;
    actions[NK_Body][IN_W] = IA_AddWhile;
    actions[NK_Body][IN_N] = IA_AddNum;
    actions[NK_Body][IN_S] = IA_AddStr;
    actions[NK_Body][IN_C] = IA_AddCall;
    actions[NK_Body][IN_O] = IA_AddOp;

    actions[NK_If][IN_D] = IA_Delete;

    actions[NK_While][IN_D] = IA_Delete;

    actions[NK_Num][IN_D] = IA_Delete;
    actions[NK_Num][IN_SPACE] = IA_StartEditing;

    actions[NK_Str][IN_D] = IA_Delete;
    actions[NK_Str][IN_SPACE] = IA_StartEditing;

    actions[NK_Call][IN_D] = IA_Delete;
    actions[NK_Call][IN_SPACE] = IA_StartEditing;

    actions[NK_Op][IN_D] = IA_Delete;
    actions[NK_Op][IN_SPACE] = IA_StartEditing;
    actions[NK_Op][IN_V] = IA_AddVar;
    actions[NK_Op][IN_I] = IA_AddIf;
    actions[NK_Op][IN_W] = IA_AddWhile;
    actions[NK_Op][IN_N] = IA_AddNum;
    actions[NK_Op][IN_S] = IA_AddStr;
    actions[NK_Op][IN_C] = IA_AddCall;
    actions[NK_Op][IN_O] = IA_AddOp;
}

int main() {
    // Setup
    static InputAction actions[NK_COUNT][IN_COUNT];
    SetupInputActions(actions);

    // Main
    vt100EnableAlternateBuffer();
    vt100HideCursor();

    Node *n = GetNode(actions);

    vt100DisableAlternateBuffer();
    vt100ShowCursor();

    NodeDraw(n);
    printf("\n");

    return 0;
}