package org.jcnc.snow.compiler.backend;

import org.jcnc.snow.compiler.ir.core.IRFunction;
import org.jcnc.snow.compiler.ir.core.IRInstruction;
import org.jcnc.snow.compiler.ir.core.IRValue;
import org.jcnc.snow.compiler.ir.instruction.BinaryOperationInstruction;
import org.jcnc.snow.compiler.ir.instruction.CallInstruction;
import org.jcnc.snow.compiler.ir.instruction.LoadConstInstruction;
import org.jcnc.snow.compiler.ir.instruction.ReturnInstruction;
import org.jcnc.snow.compiler.ir.instruction.UnaryOperationInstruction;
import org.jcnc.snow.compiler.ir.value.IRVirtualRegister;
import org.jcnc.snow.compiler.ir.value.IRConstant;
import org.jcnc.snow.vm.engine.VMOpCode;

import java.lang.reflect.Field;
import java.util.Map;

/**
 * VMCodeGenerator为单个 {@link IRFunction} 生成对应的虚拟机指令序列。
 * <p>
 * 本类依据中间表示函数的指令列表，调用 {@link VMProgramBuilder} 输出相应的 VM 操作码，
 * 包括常量加载、算术运算、函数调用和返回等指令生成逻辑。
 * </p>
 */
public final class VMCodeGenerator {

    /**
     * 虚拟寄存器到槽编号的映射表，用于确定加载和存储的目标槽。
     */
    private final Map<IRVirtualRegister, Integer> slotMap;

    /**
     * VM 指令输出器，用于构建最终的指令流。
     */
    private final VMProgramBuilder out;

    /**
     * 当前正在生成代码的函数名称，用于在返回指令处区分主函数与其他函数。
     */
    private String currentFn;

    /**
     * 构造 VMCodeGenerator 实例。
     *
     * @param slotMap 虚拟寄存器与槽编号的映射表
     * @param out     用于输出 VM 指令的构建器
     */
    public VMCodeGenerator(Map<IRVirtualRegister, Integer> slotMap,
                           VMProgramBuilder out) {
        this.slotMap = slotMap;
        this.out = out;
    }

    /**
     * 为指定的 IR 函数生成完整的 VM 指令流。
     *
     * @param fn 目标 IRFunction 实例
     * @throws IllegalStateException 若遇到不支持的 IR 指令类型
     */
    public void generate(IRFunction fn) {
        currentFn = fn.name();
        out.beginFunction(currentFn);

        for (IRInstruction inst : fn.body()) {
            switch (inst) {
                case LoadConstInstruction c -> genLoadConst(c);
                case BinaryOperationInstruction b -> genBinOp(b);
                case UnaryOperationInstruction u -> genUnary(u);
                case CallInstruction c -> genCall(c);
                case ReturnInstruction r -> genRet(r);
                default -> throw new IllegalStateException("Unsupported IR: " + inst);
            }
        }
    }

    /**
     * 生成常量加载指令序列：将常量推入栈并存入目标寄存器槽。
     *
     * @param c 常量加载指令
     */
    private void genLoadConst(LoadConstInstruction c) {
        IRConstant constant = (IRConstant) c.operands().getFirst();
        emit(op("I_PUSH"), constant.value().toString());
        emit(op("I_STORE"), slot(c.dest()));
    }

    /**
     * 生成二元算术或比较运算指令序列：
     * 依次加载两个操作数，执行运算，并存回目标寄存器槽。
     *
     * @param b 二元操作指令
     */
    private void genBinOp(BinaryOperationInstruction b) {
        emit(op("I_LOAD"), slot((IRVirtualRegister) b.operands().get(0)));
        emit(op("I_LOAD"), slot((IRVirtualRegister) b.operands().get(1)));
        emit(op(IROpCodeMapper.toVMOp(b.op())));
        emit(op("I_STORE"), slot(b.dest()));
    }

    /**
     * 生成一元运算指令序列：加载操作数，执行运算，并存回目标寄存器槽。
     *
     * @param u 一元操作指令
     */
    private void genUnary(UnaryOperationInstruction u) {
        emit(op("I_LOAD"), slot((IRVirtualRegister) u.operands().getFirst()));
        emit(op(IROpCodeMapper.toVMOp(u.op())));
        emit(op("I_STORE"), slot(u.dest()));
    }

    /**
     * 生成函数调用指令序列：
     * 按参数顺序加载每个参数，执行调用，并将返回值存入目标寄存器槽。
     *
     * @param c 调用指令
     */
    private void genCall(CallInstruction c) {
        for (IRValue arg : c.getArguments()) {
            emit(op("I_LOAD"), slot((IRVirtualRegister) arg));
        }
        int nArgs = c.getArguments().size();
        out.emitCall(c.getFunctionName(), nArgs);
        emit(op("I_STORE"), slot(c.getDest()));
    }

    /**
     * 生成返回指令序列：可选地加载返回值，然后根据是否为主函数选择 HALT 或 RET。
     *
     * @param r 返回指令
     */
    private void genRet(ReturnInstruction r) {
        if (r.value() != null) {
            emit(op("I_LOAD"), slot(r.value()));
        }
        String opcode = "main".equals(currentFn) ? op("HALT") : op("RET");
        out.emit(opcode);
    }

    /**
     * 构造并输出一条 VM 指令。
     *
     * @param opcode 操作码字符串
     * @param args   可选的操作数字符串列表
     */
    private void emit(String opcode, String... args) {
        StringBuilder sb = new StringBuilder(opcode);
        for (String a : args) {
            sb.append(' ').append(a);
        }
        out.emit(sb.toString());
    }

    /**
     * 将虚拟寄存器映射为对应的槽编号字符串。
     *
     * @param r 虚拟寄存器
     * @return 目标槽编号的字符串形式
     */
    private String slot(IRVirtualRegister r) {
        return slotMap.get(r).toString();
    }

    /**
     * 根据 VMOpCode 常量名反射获取对应的操作码字符串。
     *
     * @param name VMOpCode 常量字段名
     * @return 对应的操作码字符串
     * @throws RuntimeException 若未找到指定操作码字段
     */
    private String op(String name) {
        try {
            Field field = VMOpCode.class.getField(name);
            return field.get(null).toString();
        } catch (Exception e) {
            throw new RuntimeException("Unknown opcode: " + name, e);
        }
    }
}
