package p.studio.compiler.backend.irvm;

import p.studio.compiler.backend.bytecode.BytecodeEmitter;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayList;

public record IRVMProgram(
        IRVMModule module,
        BytecodeEmitter.EmissionPlan emissionPlan) {

    public IRVMProgram {
        module = module == null ? IRVMModule.empty() : module;
        emissionPlan = emissionPlan == null ? BytecodeEmitter.EmissionPlan.empty() : emissionPlan;
        if (!emissionPlan.functions().isEmpty() && !containsRawSyscall(emissionPlan)) {
            validateCoherence(module, emissionPlan);
        }
    }

    public IRVMProgram(final IRVMModule module) {
        this(module, BytecodeEmitter.EmissionPlan.empty());
    }

    public IRVMProgram(
            final boolean hasInternalOpcodes,
            final BytecodeEmitter.EmissionPlan emissionPlan) {
        this(moduleFromPlan(hasInternalOpcodes, emissionPlan), emissionPlan);
    }

    public static IRVMProgram empty() {
        return new IRVMProgram(IRVMModule.empty(), BytecodeEmitter.EmissionPlan.empty());
    }

    public BytecodeEmitter.EmissionPlan coherentEmissionPlan() {
        if (emissionPlan.functions().isEmpty()) {
            return deriveEmissionPlan(module);
        }
        validateCanonicalPlan(emissionPlan);
        validateCoherence(module, emissionPlan);
        return emissionPlan;
    }

    public boolean hasInternalOpcodes() {
        for (final var function : module.functions()) {
            for (final var instruction : function.instructions()) {
                if (instruction.op().internal()) {
                    return true;
                }
            }
        }
        return false;
    }

    private void validateCoherence(
            final IRVMModule module,
            final BytecodeEmitter.EmissionPlan emissionPlan) {
        if (module.functions().size() != emissionPlan.functions().size()) {
            throw new IllegalArgumentException("irvm module/function-plan size mismatch");
        }
        for (var fnIndex = 0; fnIndex < module.functions().size(); fnIndex++) {
            final var irvmFunction = module.functions().get(fnIndex);
            final var planFunction = emissionPlan.functions().get(fnIndex);
            if (!irvmFunction.name().equals(planFunction.name())
                    || irvmFunction.paramSlots() != planFunction.paramSlots()
                    || irvmFunction.localSlots() != planFunction.localSlots()
                    || irvmFunction.returnSlots() != planFunction.returnSlots()
                    || irvmFunction.maxStackSlots() != planFunction.maxStackSlots()) {
                throw new IllegalArgumentException("irvm/emission plan function header mismatch at index " + fnIndex);
            }
            if (irvmFunction.instructions().size() != planFunction.operations().size()) {
                throw new IllegalArgumentException("irvm/emission plan operation count mismatch at index " + fnIndex);
            }
            for (var opIndex = 0; opIndex < irvmFunction.instructions().size(); opIndex++) {
                final var instruction = irvmFunction.instructions().get(opIndex);
                final var operation = planFunction.operations().get(opIndex);
                if (!coherentPair(instruction, operation)) {
                    throw new IllegalArgumentException(
                            "irvm/emission plan opcode mismatch at function=%d op=%d".formatted(fnIndex, opIndex));
                }
            }
        }
    }

    private boolean coherentPair(
            final IRVMInstruction instruction,
            final BytecodeEmitter.Operation operation) {
        if (instruction.op().internal()) {
            return true;
        }
        final var immediate = instruction.immediate() == null ? 0 : instruction.immediate();
        return switch (instruction.op().opcode()) {
            case 0x01 -> operation.kind() == BytecodeEmitter.OperationKind.HALT;
            case 0x51 -> operation.kind() == BytecodeEmitter.OperationKind.RET;
            case 0x10 -> operation.kind() == BytecodeEmitter.OperationKind.PUSH_CONST && operation.immediate() == immediate;
            case 0x11 -> operation.kind() == BytecodeEmitter.OperationKind.POP;
            case 0x16 -> operation.kind() == BytecodeEmitter.OperationKind.PUSH_BOOL;
            case 0x40 -> operation.kind() == BytecodeEmitter.OperationKind.GET_GLOBAL && operation.immediate() == immediate;
            case 0x42 -> operation.kind() == BytecodeEmitter.OperationKind.GET_LOCAL && operation.immediate() == immediate;
            case 0x41 -> operation.kind() == BytecodeEmitter.OperationKind.SET_GLOBAL && operation.immediate() == immediate;
            case 0x43 -> operation.kind() == BytecodeEmitter.OperationKind.SET_LOCAL && operation.immediate() == immediate;
            case 0x17 -> operation.kind() == BytecodeEmitter.OperationKind.PUSH_I32 && operation.immediate() == immediate;
            case 0x20 -> operation.kind() == BytecodeEmitter.OperationKind.ADD;
            case 0x21 -> operation.kind() == BytecodeEmitter.OperationKind.SUB;
            case 0x22 -> operation.kind() == BytecodeEmitter.OperationKind.MUL;
            case 0x23 -> operation.kind() == BytecodeEmitter.OperationKind.DIV;
            case 0x24 -> operation.kind() == BytecodeEmitter.OperationKind.MOD;
            case 0x30 -> operation.kind() == BytecodeEmitter.OperationKind.EQ;
            case 0x31 -> operation.kind() == BytecodeEmitter.OperationKind.NEQ;
            case 0x32 -> operation.kind() == BytecodeEmitter.OperationKind.LT;
            case 0x33 -> operation.kind() == BytecodeEmitter.OperationKind.GT;
            case 0x34 -> operation.kind() == BytecodeEmitter.OperationKind.AND;
            case 0x35 -> operation.kind() == BytecodeEmitter.OperationKind.OR;
            case 0x36 -> operation.kind() == BytecodeEmitter.OperationKind.NOT;
            case 0x3C -> operation.kind() == BytecodeEmitter.OperationKind.LTE;
            case 0x3D -> operation.kind() == BytecodeEmitter.OperationKind.GTE;
            case 0x3E -> operation.kind() == BytecodeEmitter.OperationKind.NEG;
            case 0x50 -> operation.kind() == BytecodeEmitter.OperationKind.CALL_FUNC && operation.immediate() == immediate;
            case 0x02 -> operation.kind() == BytecodeEmitter.OperationKind.JMP && operation.immediate() == immediate;
            case 0x04 -> operation.kind() == BytecodeEmitter.OperationKind.JMP_IF_TRUE && operation.immediate() == immediate;
            case 0x03 -> operation.kind() == BytecodeEmitter.OperationKind.JMP_IF_FALSE && operation.immediate() == immediate;
            case 0x71 -> operation.kind() == BytecodeEmitter.OperationKind.HOSTCALL;
            case 0x72 -> operation.kind() == BytecodeEmitter.OperationKind.INTRINSIC && operation.immediate() == immediate;
            default -> false;
        };
    }

    private BytecodeEmitter.EmissionPlan deriveEmissionPlan(final IRVMModule module) {
        final var functions = new ArrayList<BytecodeEmitter.FunctionPlan>(module.functions().size());
        for (final var function : module.functions()) {
            final var operations = new ArrayList<BytecodeEmitter.Operation>(function.instructions().size());
            for (final var instruction : function.instructions()) {
                operations.add(deriveOperation(instruction));
            }
            functions.add(new BytecodeEmitter.FunctionPlan(
                    function.name(),
                    function.paramSlots(),
                    function.localSlots(),
                    function.returnSlots(),
                    function.maxStackSlots(),
                    ReadOnlyList.wrap(operations)));
        }
        return new BytecodeEmitter.EmissionPlan(
                0,
                ReadOnlyList.empty(),
                ReadOnlyList.empty(),
                ReadOnlyList.wrap(functions));
    }

    private boolean containsRawSyscall(final BytecodeEmitter.EmissionPlan plan) {
        for (final var function : plan.functions()) {
            for (final var operation : function.operations()) {
                if (operation.kind() == BytecodeEmitter.OperationKind.RAW_SYSCALL) {
                    return true;
                }
            }
        }
        return false;
    }

    private void validateCanonicalPlan(final BytecodeEmitter.EmissionPlan plan) {
        for (final var function : plan.functions()) {
            for (final var operation : function.operations()) {
                if (operation.kind() == BytecodeEmitter.OperationKind.RAW_SYSCALL) {
                    throw new IllegalArgumentException("raw syscall operation is forbidden in canonical IRVM emission plan");
                }
                if (operation.kind() == BytecodeEmitter.OperationKind.HOSTCALL && operation.syscallDecl() == null) {
                    throw new IllegalArgumentException("hostcall operation requires syscall declaration metadata");
                }
            }
        }
    }

    private BytecodeEmitter.Operation deriveOperation(final IRVMInstruction instruction) {
        final var immediate = instruction.immediate() == null ? 0 : instruction.immediate();
        return switch (instruction.op().opcode()) {
            case 0x01 -> BytecodeEmitter.Operation.halt();
            case 0x51 -> BytecodeEmitter.Operation.ret();
            case 0x10 -> BytecodeEmitter.Operation.pushConst(immediate, null);
            case 0x11 -> BytecodeEmitter.Operation.pop(null);
            case 0x16 -> BytecodeEmitter.Operation.pushBool(immediate != 0, null);
            case 0x40 -> BytecodeEmitter.Operation.getGlobal(immediate, null);
            case 0x42 -> BytecodeEmitter.Operation.getLocal(immediate, null);
            case 0x41 -> BytecodeEmitter.Operation.setGlobal(immediate, null);
            case 0x43 -> BytecodeEmitter.Operation.setLocal(immediate, null);
            case 0x17 -> BytecodeEmitter.Operation.pushI32(immediate, null);
            case 0x20 -> BytecodeEmitter.Operation.add(null);
            case 0x21 -> BytecodeEmitter.Operation.sub(null);
            case 0x22 -> BytecodeEmitter.Operation.mul(null);
            case 0x23 -> BytecodeEmitter.Operation.div(null);
            case 0x24 -> BytecodeEmitter.Operation.mod(null);
            case 0x30 -> BytecodeEmitter.Operation.eq(null);
            case 0x31 -> BytecodeEmitter.Operation.neq(null);
            case 0x32 -> BytecodeEmitter.Operation.lt(null);
            case 0x33 -> BytecodeEmitter.Operation.gt(null);
            case 0x34 -> BytecodeEmitter.Operation.and(null);
            case 0x35 -> BytecodeEmitter.Operation.or(null);
            case 0x36 -> BytecodeEmitter.Operation.not(null);
            case 0x3C -> BytecodeEmitter.Operation.lte(null);
            case 0x3D -> BytecodeEmitter.Operation.gte(null);
            case 0x3E -> BytecodeEmitter.Operation.neg(null);
            case 0x50 -> BytecodeEmitter.Operation.callFunc(immediate);
            case 0x02 -> BytecodeEmitter.Operation.jmp(immediate, null);
            case 0x04 -> BytecodeEmitter.Operation.jmpIfTrue(immediate, null);
            case 0x03 -> BytecodeEmitter.Operation.jmpIfFalse(immediate, null);
            case 0x71 -> throw new IllegalArgumentException(
                    "cannot derive HOSTCALL operation without canonical syscall metadata");
            case 0x72 -> BytecodeEmitter.Operation.intrinsic(immediate);
            default -> throw new IllegalArgumentException("cannot derive emission op for irvm opcode: " + instruction.op().opcode());
        };
    }

    private static IRVMModule moduleFromPlan(
            final boolean hasInternalOpcodes,
            final BytecodeEmitter.EmissionPlan emissionPlan) {
        final var inputPlan = emissionPlan == null ? BytecodeEmitter.EmissionPlan.empty() : emissionPlan;
        if (inputPlan.functions().isEmpty()) {
            return new IRVMModule(
                    "core-v1",
                    ReadOnlyList.from(new IRVMFunction(
                            "__synthetic__",
                            0,
                            0,
                            0,
                            1,
                            ReadOnlyList.from(
                                    new IRVMInstruction(hasInternalOpcodes ? IRVMOp.INTERNAL_EXT : IRVMOp.HALT, null)))));
        }
        final var functions = new ArrayList<IRVMFunction>(inputPlan.functions().size());
        for (final var functionPlan : inputPlan.functions()) {
            final var instructions = new ArrayList<IRVMInstruction>(functionPlan.operations().size());
            for (final var operation : functionPlan.operations()) {
                instructions.add(operationToInstruction(operation));
            }
            if (hasInternalOpcodes) {
                instructions.add(0, new IRVMInstruction(IRVMOp.INTERNAL_EXT, null));
            }
            functions.add(new IRVMFunction(
                    functionPlan.name(),
                    functionPlan.paramSlots(),
                    functionPlan.localSlots(),
                    functionPlan.returnSlots(),
                    functionPlan.maxStackSlots(),
                    ReadOnlyList.wrap(instructions)));
        }
        return new IRVMModule("core-v1", ReadOnlyList.wrap(functions));
    }

    private static IRVMInstruction operationToInstruction(final BytecodeEmitter.Operation operation) {
        return switch (operation.kind()) {
            case HALT -> new IRVMInstruction(IRVMOp.HALT, null);
            case RET -> new IRVMInstruction(IRVMOp.RET, null);
            case PUSH_CONST -> new IRVMInstruction(IRVMOp.PUSH_CONST, operation.immediate());
            case POP -> new IRVMInstruction(IRVMOp.POP, null);
            case PUSH_BOOL -> new IRVMInstruction(IRVMOp.PUSH_BOOL, operation.immediate());
            case GET_GLOBAL -> new IRVMInstruction(IRVMOp.GET_GLOBAL, operation.immediate());
            case GET_LOCAL -> new IRVMInstruction(IRVMOp.GET_LOCAL, operation.immediate());
            case SET_GLOBAL -> new IRVMInstruction(IRVMOp.SET_GLOBAL, operation.immediate());
            case SET_LOCAL -> new IRVMInstruction(IRVMOp.SET_LOCAL, operation.immediate());
            case PUSH_I32 -> new IRVMInstruction(IRVMOp.PUSH_I32, operation.immediate());
            case ADD -> new IRVMInstruction(IRVMOp.ADD, null);
            case SUB -> new IRVMInstruction(IRVMOp.SUB, null);
            case MUL -> new IRVMInstruction(IRVMOp.MUL, null);
            case DIV -> new IRVMInstruction(IRVMOp.DIV, null);
            case MOD -> new IRVMInstruction(IRVMOp.MOD, null);
            case EQ -> new IRVMInstruction(IRVMOp.EQ, null);
            case NEQ -> new IRVMInstruction(IRVMOp.NEQ, null);
            case LT -> new IRVMInstruction(IRVMOp.LT, null);
            case LTE -> new IRVMInstruction(IRVMOp.LTE, null);
            case GT -> new IRVMInstruction(IRVMOp.GT, null);
            case GTE -> new IRVMInstruction(IRVMOp.GTE, null);
            case AND -> new IRVMInstruction(IRVMOp.AND, null);
            case OR -> new IRVMInstruction(IRVMOp.OR, null);
            case NOT -> new IRVMInstruction(IRVMOp.NOT, null);
            case NEG -> new IRVMInstruction(IRVMOp.NEG, null);
            case CALL_FUNC -> new IRVMInstruction(IRVMOp.CALL, operation.immediate());
            case JMP -> new IRVMInstruction(IRVMOp.JMP, operation.immediate());
            case JMP_IF_TRUE -> new IRVMInstruction(IRVMOp.JMP_IF_TRUE, operation.immediate());
            case JMP_IF_FALSE -> new IRVMInstruction(IRVMOp.JMP_IF_FALSE, operation.immediate());
            case HOSTCALL -> new IRVMInstruction(IRVMOp.HOSTCALL, 0);
            case INTRINSIC -> new IRVMInstruction(IRVMOp.INTRINSIC, operation.immediate());
            case RAW_SYSCALL -> new IRVMInstruction(IRVMOp.HALT, null);
        };
    }
}
