#include "superinst.h"
#include "fastframe.h"
#include <string.h>

static void optimize_inc_dec_local(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 3 < count; i++) {
        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_CONST &&
            instrs[i+1].operand.int_value == 1 &&
            instrs[i+2].opcode == RAVA_OP_ADD &&
            instrs[i+3].opcode == RAVA_OP_STORE_LOCAL &&
            instrs[i].operand.var.index == instrs[i+3].operand.var.index) {

            instrs[i].opcode = RAVA_OP_INC_LOCAL;
            instrs[i+1].opcode = RAVA_OP_NOP;
            instrs[i+2].opcode = RAVA_OP_NOP;
            instrs[i+3].opcode = RAVA_OP_NOP;
        }

        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_CONST &&
            instrs[i+1].operand.int_value == 1 &&
            instrs[i+2].opcode == RAVA_OP_SUB &&
            instrs[i+3].opcode == RAVA_OP_STORE_LOCAL &&
            instrs[i].operand.var.index == instrs[i+3].operand.var.index) {

            instrs[i].opcode = RAVA_OP_DEC_LOCAL;
            instrs[i+1].opcode = RAVA_OP_NOP;
            instrs[i+2].opcode = RAVA_OP_NOP;
            instrs[i+3].opcode = RAVA_OP_NOP;
        }
    }
}

static void optimize_loop_comparison(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 3 < count; i++) {
        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_CONST &&
            instrs[i+2].opcode == RAVA_OP_LT &&
            instrs[i+3].opcode == RAVA_OP_JUMP_IF_FALSE) {

            int local_idx = instrs[i].operand.var.index;
            int64_t const_val = instrs[i+1].operand.int_value;
            int label = instrs[i+3].operand.label_id;

            instrs[i].opcode = RAVA_OP_LOAD_LOCAL_CONST_LT_JUMPFALSE;
            instrs[i].operand.super.local_index = local_idx;
            instrs[i].operand.super.const_value = const_val;
            instrs[i].operand.super.label_id = label;

            instrs[i+1].opcode = RAVA_OP_NOP;
            instrs[i+2].opcode = RAVA_OP_NOP;
            instrs[i+3].opcode = RAVA_OP_NOP;
        }

        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_CONST &&
            instrs[i+2].opcode == RAVA_OP_LE &&
            instrs[i+3].opcode == RAVA_OP_JUMP_IF_FALSE) {

            int local_idx = instrs[i].operand.var.index;
            int64_t const_val = instrs[i+1].operand.int_value;
            int label = instrs[i+3].operand.label_id;

            instrs[i].opcode = RAVA_OP_LOAD_LOCAL_CONST_LE_JUMPFALSE;
            instrs[i].operand.super.local_index = local_idx;
            instrs[i].operand.super.const_value = const_val;
            instrs[i].operand.super.label_id = label;

            instrs[i+1].opcode = RAVA_OP_NOP;
            instrs[i+2].opcode = RAVA_OP_NOP;
            instrs[i+3].opcode = RAVA_OP_NOP;
        }
    }
}

static void optimize_two_locals(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 1 < count; i++) {
        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_LOCAL) {

            int idx1 = instrs[i].operand.var.index;
            int idx2 = instrs[i+1].operand.var.index;

            instrs[i].opcode = RAVA_OP_LOAD_TWO_LOCALS;
            instrs[i].operand.two_locals.index1 = idx1;
            instrs[i].operand.two_locals.index2 = idx2;

            instrs[i+1].opcode = RAVA_OP_NOP;
        }
    }
}

static void optimize_add_local_to_local(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 3 < count; i++) {
        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+2].opcode == RAVA_OP_ADD &&
            instrs[i+3].opcode == RAVA_OP_STORE_LOCAL &&
            instrs[i].operand.var.index == instrs[i+3].operand.var.index) {

            int dest_idx = instrs[i].operand.var.index;
            int src_idx = instrs[i+1].operand.var.index;

            instrs[i].opcode = RAVA_OP_ADD_LOCAL_TO_LOCAL;
            instrs[i].operand.add_locals.dest_index = dest_idx;
            instrs[i].operand.add_locals.src_index = src_idx;

            instrs[i+1].opcode = RAVA_OP_NOP;
            instrs[i+2].opcode = RAVA_OP_NOP;
            instrs[i+3].opcode = RAVA_OP_NOP;
        }
    }
}

static void optimize_local_lt_local_jump(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 3 < count; i++) {
        if (instrs[i].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+1].opcode == RAVA_OP_LOAD_LOCAL &&
            instrs[i+2].opcode == RAVA_OP_LT &&
            instrs[i+3].opcode == RAVA_OP_JUMP_IF_FALSE) {

            int local1 = instrs[i].operand.var.index;
            int local2 = instrs[i+1].operand.var.index;
            int label = instrs[i+3].operand.label_id;

            instrs[i].opcode = RAVA_OP_LOAD_LOCAL_LT_LOCAL_JUMPFALSE;
            instrs[i].operand.cmp_locals.local1 = local1;
            instrs[i].operand.cmp_locals.local2 = local2;
            instrs[i].operand.cmp_locals.label_id = label;

            instrs[i+1].opcode = RAVA_OP_NOP;
            instrs[i+2].opcode = RAVA_OP_NOP;
            instrs[i+3].opcode = RAVA_OP_NOP;
        }
    }
}

static int is_power_of_two(int64_t n) {
    return n > 0 && (n & (n - 1)) == 0;
}

static int log2_int(int64_t n) {
    int result = 0;
    while (n > 1) {
        n >>= 1;
        result++;
    }
    return result;
}

static void optimize_strength_reduction(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 2 < count; i++) {
        if (instrs[i+1].opcode == RAVA_OP_LOAD_CONST) {
            int64_t const_val = instrs[i+1].operand.int_value;

            if (instrs[i+2].opcode == RAVA_OP_MUL && is_power_of_two(const_val)) {
                int shift_amount = log2_int(const_val);
                instrs[i+1].operand.int_value = shift_amount;
                instrs[i+2].opcode = RAVA_OP_SHL;
            }
            else if (instrs[i+2].opcode == RAVA_OP_DIV && is_power_of_two(const_val)) {
                int shift_amount = log2_int(const_val);
                instrs[i+1].operand.int_value = shift_amount;
                instrs[i+2].opcode = RAVA_OP_SHR;
            }
            else if (instrs[i+2].opcode == RAVA_OP_MOD && const_val == 2) {
                instrs[i+1].operand.int_value = 1;
                instrs[i+2].opcode = RAVA_OP_AND;
            }
        }
    }
}

static void optimize_constant_folding(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i + 2 < count; i++) {
        if (instrs[i].opcode == RAVA_OP_LOAD_CONST &&
            instrs[i+1].opcode == RAVA_OP_LOAD_CONST) {

            int64_t val1 = instrs[i].operand.int_value;
            int64_t val2 = instrs[i+1].operand.int_value;
            int64_t result = 0;
            int can_fold = 1;

            switch (instrs[i+2].opcode) {
                case RAVA_OP_ADD:
                    result = val1 + val2;
                    break;
                case RAVA_OP_SUB:
                    result = val1 - val2;
                    break;
                case RAVA_OP_MUL:
                    result = val1 * val2;
                    break;
                case RAVA_OP_DIV:
                    if (val2 != 0) {
                        result = val1 / val2;
                    } else {
                        can_fold = 0;
                    }
                    break;
                case RAVA_OP_MOD:
                    if (val2 != 0) {
                        result = val1 % val2;
                    } else {
                        can_fold = 0;
                    }
                    break;
                case RAVA_OP_AND:
                    result = val1 & val2;
                    break;
                case RAVA_OP_OR:
                    result = val1 | val2;
                    break;
                case RAVA_OP_XOR:
                    result = val1 ^ val2;
                    break;
                case RAVA_OP_SHL:
                    result = val1 << val2;
                    break;
                case RAVA_OP_SHR:
                    result = val1 >> val2;
                    break;
                case RAVA_OP_EQ:
                    result = val1 == val2;
                    break;
                case RAVA_OP_NE:
                    result = val1 != val2;
                    break;
                case RAVA_OP_LT:
                    result = val1 < val2;
                    break;
                case RAVA_OP_LE:
                    result = val1 <= val2;
                    break;
                case RAVA_OP_GT:
                    result = val1 > val2;
                    break;
                case RAVA_OP_GE:
                    result = val1 >= val2;
                    break;
                default:
                    can_fold = 0;
                    break;
            }

            if (can_fold) {
                instrs[i].operand.int_value = result;
                instrs[i+1].opcode = RAVA_OP_NOP;
                instrs[i+2].opcode = RAVA_OP_NOP;
            }
        }
    }
}

static void optimize_inline_calls(RavaInstruction_t *instrs, size_t count) {
    for (size_t i = 0; i < count; i++) {
        if (instrs[i].opcode == RAVA_OP_CALL_STATIC ||
            instrs[i].opcode == RAVA_OP_CALL_RECURSIVE) {

            RavaMethod_t *target = (RavaMethod_t*)instrs[i].operand.call.cached_method;
            if (target && target->is_simple && target->is_leaf &&
                target->local_count <= RAVA_INLINE_FRAME_LOCALS &&
                target->max_stack <= RAVA_INLINE_FRAME_STACK) {

                instrs[i].opcode = RAVA_OP_CALL_INLINE;
            }
        }
    }
}

void rava_optimize_superinstructions(RavaMethod_t* method) {
    if (!method || !method->instructions) return;

    RavaInstruction_t *instrs = method->instructions->instructions;
    size_t count = method->instructions->count;

    optimize_constant_folding(instrs, count);
    optimize_inline_calls(instrs, count);
    optimize_strength_reduction(instrs, count);
    optimize_add_local_to_local(instrs, count);
    optimize_local_lt_local_jump(instrs, count);
    optimize_inc_dec_local(instrs, count);
    optimize_loop_comparison(instrs, count);
    optimize_two_locals(instrs, count);
}