#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "ast.h"
#include "optimizer.h"

typedef struct {
    char *name;
    long long value;
    int active;
} ConstEntry;

static ConstEntry g_consts[256];
static int g_const_count = 0;

static void const_clear(void) {
    for (int i = 0; i < g_const_count; i++) free(g_consts[i].name);
    g_const_count = 0;
}

static void const_set(const char *name, long long val) {
    for (int i = 0; i < g_const_count; i++) {
        if (strcmp(g_consts[i].name, name) == 0) {
            g_consts[i].value = val;
            g_consts[i].active = 1;
            return;
        }
    }
    if (g_const_count < 256) {
        g_consts[g_const_count].name = strdup(name);
        g_consts[g_const_count].value = val;
        g_consts[g_const_count].active = 1;
        g_const_count++;
    }
}

static void const_invalidate(const char *name) {
    for (int i = 0; i < g_const_count; i++) {
        if (strcmp(g_consts[i].name, name) == 0) {
            g_consts[i].active = 0;
            return;
        }
    }
}

static int const_get(const char *name, long long *out) {
    for (int i = 0; i < g_const_count; i++) {
        if (g_consts[i].active && strcmp(g_consts[i].name, name) == 0) {
            *out = g_consts[i].value;
            return 1;
        }
    }
    return 0;
}

static int expr_const_eval(Expr *e, long long *out) {
    if (!e) return 0;
    if (e->type == EXPR_INT) { *out = e->int_val; return 1; }
    if (e->type == EXPR_VAR) return 0;
    if (e->type == EXPR_BINOP) {
        long long l, r;
        if (!expr_const_eval(e->left, &l)) return 0;
        if (!expr_const_eval(e->right, &r)) return 0;
        switch (e->op) {
            case '+': *out = l + r; return 1;
            case '-': *out = l - r; return 1;
            case '*': *out = l * r; return 1;
            case '/': *out = r ? l / r : 0; return 1;
            case '%': *out = r ? l % r : 0; return 1;
        }
    }
    return 0;
}

static Expr* expr_fold(Expr *e) {
    if (!e) return NULL;
    if (e->type == EXPR_BINOP) {
        e->left = expr_fold(e->left);
        e->right = expr_fold(e->right);
        long long val;
        if (expr_const_eval(e, &val)) {
            expr_free(e);
            return expr_new_int(val);
        }
    }
    return e;
}

static Expr* expr_propagate(Expr *e) {
    if (!e) return NULL;
    if (e->type == EXPR_VAR) {
        long long val;
        if (const_get(e->var_name, &val)) {
            expr_free(e);
            return expr_new_int(val);
        }
    }
    if (e->type == EXPR_BINOP) {
        e->left = expr_propagate(e->left);
        e->right = expr_propagate(e->right);
        long long val;
        if (expr_const_eval(e, &val)) {
            expr_free(e);
            return expr_new_int(val);
        }
    }
    return e;
}

static Test* test_propagate(Test *t) {
    if (!t) return NULL;
    if (t->type == TEST_NUM) {
        t->num_left = expr_propagate(t->num_left);
        t->num_right = expr_propagate(t->num_right);
    }
    return t;
}

static int test_const_eval(Test *t, int *out) {
    if (!t) return 0;
    if (t->type == TEST_NUM) {
        long long l, r;
        if (!expr_const_eval(t->num_left, &l)) return 0;
        if (!expr_const_eval(t->num_right, &r)) return 0;
        switch (t->op) {
            case 'e': *out = (l == r); return 1;
            case 'n': *out = (l != r); return 1;
            case 'l': *out = (l < r); return 1;
            case 'L': *out = (l <= r); return 1;
            case 'g': *out = (l > r); return 1;
            case 'G': *out = (l >= r); return 1;
        }
    }
    return 0;
}

static Stmt* optimize_stmt(Stmt *s, int level);

static Stmt* optimize_list(Stmt *head, int level) {
    Stmt *result = NULL, **tail = &result;
    for (Stmt *s = head; s; ) {
        Stmt *next = s->next;
        s->next = NULL;
        Stmt *opt = optimize_stmt(s, level);
        if (opt) {
            *tail = opt;
            while ((*tail)->next) tail = &(*tail)->next;
            tail = &(*tail)->next;
        }
        s = next;
    }
    return result;
}

static Stmt* optimize_stmt(Stmt *s, int level) {
    switch (s->type) {
        case STMT_ASSIGN: {
            if (s->arith_val) {
                s->arith_val = expr_fold(s->arith_val);
                s->arith_val = expr_propagate(s->arith_val);
                long long val;
                if (expr_const_eval(s->arith_val, &val)) {
                    const_set(s->var_name, val);
                } else {
                    const_invalidate(s->var_name);
                }
            } else {
                const_invalidate(s->var_name);
            }
            return s;
        }
        case STMT_ECHO: {
            for (int i = 0; i < s->echo_argc; i++) {
                Word *w = s->echo_args[i];
                if (w->type == WORD_VAR) {
                    long long val;
                    if (const_get(w->text, &val)) {
                        free(w->text);
                        char buf[32];
                        snprintf(buf, sizeof(buf), "%lld", val);
                        w->text = strdup(buf);
                        w->type = WORD_LITERAL;
                    }
                }
            }
            return s;
        }
        case STMT_IF: {
            s->cond = test_propagate(s->cond);
            int cval;
            if (test_const_eval(s->cond, &cval)) {
                Stmt *alive = cval ? s->then_body : s->else_body;
                Stmt *dead = cval ? s->else_body : s->then_body;
                ast_free(dead);
                test_free(s->cond);
                s->cond = NULL;
                s->then_body = NULL;
                s->else_body = NULL;
                Stmt *body = optimize_list(alive, level);
                free(s);
                return body;
            }
            ConstEntry saved[256];
            int saved_count = g_const_count;
            for (int i = 0; i < g_const_count; i++) {
                saved[i].name = g_consts[i].name;
                saved[i].value = g_consts[i].value;
                saved[i].active = g_consts[i].active;
            }
            s->then_body = optimize_list(s->then_body, level);
            for (int i = 0; i < saved_count; i++) {
                g_consts[i].name = saved[i].name;
                g_consts[i].value = saved[i].value;
                g_consts[i].active = saved[i].active;
            }
            g_const_count = saved_count;
            s->else_body = optimize_list(s->else_body, level);
            for (int i = 0; i < g_const_count; i++) g_consts[i].active = 0;
            return s;
        }
        case STMT_FOR: {
            if (level >= 2 && s->for_list_count >= 1 && s->for_list_count <= 8) {
                int can_unroll = 1;
                for (Stmt *b = s->body; b; b = b->next) {
                    if (b->type != STMT_ECHO) { can_unroll = 0; break; }
                }
                if (can_unroll) {
                    Stmt *expanded = NULL, **etail = &expanded;
                    for (int i = 0; i < s->for_list_count; i++) {
                        for (Stmt *b = s->body; b; b = b->next) {
                            Stmt *copy = calloc(1, sizeof(Stmt));
                            copy->type = STMT_ECHO;
                            copy->echo_argc = b->echo_argc;
                            copy->echo_args = malloc(sizeof(Word*) * b->echo_argc);
                            for (int j = 0; j < b->echo_argc; j++) {
                                copy->echo_args[j] = malloc(sizeof(Word));
                                if (b->echo_args[j]->type == WORD_VAR &&
                                    strcmp(b->echo_args[j]->text, s->for_var) == 0) {
                                    copy->echo_args[j]->type = WORD_LITERAL;
                                    char buf[32];
                                    snprintf(buf, sizeof(buf), "%s", s->for_list[i]);
                                    copy->echo_args[j]->text = strdup(buf);
                                } else {
                                    copy->echo_args[j]->type = b->echo_args[j]->type;
                                    copy->echo_args[j]->text = strdup(b->echo_args[j]->text);
                                }
                            }
                            *etail = copy;
                            etail = &copy->next;
                        }
                    }
                    s->body = NULL;
                    free(s);
                    return expanded;
                }
            }
            s->body = optimize_list(s->body, level);
            return s;
        }
        case STMT_WHILE: {
            for (int i = 0; i < g_const_count; i++) g_consts[i].active = 0;
            s->cond = test_propagate(s->cond);
            s->body = optimize_list(s->body, level);
            for (int i = 0; i < g_const_count; i++) g_consts[i].active = 0;
            return s;
        }
        case STMT_EXIT: {
            if (s->exit_expr) {
                s->exit_expr = expr_fold(s->exit_expr);
                s->exit_expr = expr_propagate(s->exit_expr);
            }
            return s;
        }
        default:
            return s;
    }
}

Stmt* optimizer_run(Stmt *ast, int level) {
    const_clear();
    Stmt *result = optimize_list(ast, level);
    const_clear();
    return result;
}
