#include "superinst.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;
        }
    }
}

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