package p.studio.compiler.backend.bytecode;

import p.studio.utilities.structures.ReadOnlyList;

import java.io.ByteArrayOutputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Objects;

public class BytecodeEmitter {
    private static final int OP_HALT = 0x01;
    private static final int OP_RET = 0x51;
    private static final int OP_CALL = 0x50;
    private static final int OP_JMP = 0x02;
    private static final int OP_JMP_IF_FALSE = 0x03;
    private static final int OP_JMP_IF_TRUE = 0x04;
    private static final int OP_PUSH_CONST = 0x10;
    private static final int OP_POP = 0x11;
    private static final int OP_PUSH_BOOL = 0x16;
    private static final int OP_ADD = 0x20;
    private static final int OP_SUB = 0x21;
    private static final int OP_MUL = 0x22;
    private static final int OP_DIV = 0x23;
    private static final int OP_MOD = 0x24;
    private static final int OP_EQ = 0x30;
    private static final int OP_NEQ = 0x31;
    private static final int OP_LT = 0x32;
    private static final int OP_GT = 0x33;
    private static final int OP_AND = 0x34;
    private static final int OP_OR = 0x35;
    private static final int OP_NOT = 0x36;
    private static final int OP_LTE = 0x3C;
    private static final int OP_GTE = 0x3D;
    private static final int OP_NEG = 0x3E;
    private static final int OP_GET_GLOBAL = 0x40;
    private static final int OP_SET_GLOBAL = 0x41;
    private static final int OP_GET_LOCAL = 0x42;
    private static final int OP_SET_LOCAL = 0x43;
    private static final int OP_PUSH_I32 = 0x17;
    private static final int OP_HOSTCALL = 0x71;
    private static final int OP_SYSCALL = 0x70;
    private static final int OP_INTRINSIC = 0x72;

    public BytecodeModule emit(
            final EmissionPlan plan) {
        final var inputPlan = plan == null ? EmissionPlan.empty() : plan;
        final var canonicalConstPool = dedupeConstPool(inputPlan.constPool());
        final var canonicalExports = dedupeExports(inputPlan.exports());
        final var functionFragments = new ArrayList<BytecodeFunctionLayoutBuilder.FunctionFragment>(inputPlan.functions().size());
        final var orderedSyscalls = new LinkedHashMap<SyscallIdentity, BytecodeModule.SyscallDecl>();
        final var syscallIndexByIdentity = new LinkedHashMap<SyscallIdentity, Integer>();

        for (final var function : inputPlan.functions()) {
            final var code = new ByteArrayOutputStream();
            final var spans = new ArrayList<BytecodeFunctionLayoutBuilder.InstructionSpan>();
            for (final var op : function.operations()) {
                final var pc = code.size();
                switch (op.kind()) {
                    case HALT -> {
                        writeOpNoImm(code, OP_HALT);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case RET -> {
                        writeOpNoImm(code, OP_RET);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case CALL_FUNC -> {
                        writeOpU32(code, OP_CALL, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case PUSH_CONST -> {
                        writeOpU32(code, OP_PUSH_CONST, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case POP -> {
                        writeOpNoImm(code, OP_POP);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case GET_LOCAL -> {
                        writeOpU32(code, OP_GET_LOCAL, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case GET_GLOBAL -> {
                        writeOpU32(code, OP_GET_GLOBAL, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case SET_LOCAL -> {
                        writeOpU32(code, OP_SET_LOCAL, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case SET_GLOBAL -> {
                        writeOpU32(code, OP_SET_GLOBAL, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case PUSH_BOOL -> {
                        writeOpU8(code, OP_PUSH_BOOL, op.immediate() == 0 ? 0 : 1);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case PUSH_I32 -> {
                        writeOpU32(code, OP_PUSH_I32, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case ADD -> {
                        writeOpNoImm(code, OP_ADD);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case SUB -> {
                        writeOpNoImm(code, OP_SUB);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case MUL -> {
                        writeOpNoImm(code, OP_MUL);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case DIV -> {
                        writeOpNoImm(code, OP_DIV);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case MOD -> {
                        writeOpNoImm(code, OP_MOD);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case EQ -> {
                        writeOpNoImm(code, OP_EQ);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case NEQ -> {
                        writeOpNoImm(code, OP_NEQ);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case LT -> {
                        writeOpNoImm(code, OP_LT);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case LTE -> {
                        writeOpNoImm(code, OP_LTE);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case GT -> {
                        writeOpNoImm(code, OP_GT);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case GTE -> {
                        writeOpNoImm(code, OP_GTE);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case AND -> {
                        writeOpNoImm(code, OP_AND);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case OR -> {
                        writeOpNoImm(code, OP_OR);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case NOT -> {
                        writeOpNoImm(code, OP_NOT);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case NEG -> {
                        writeOpNoImm(code, OP_NEG);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case JMP -> {
                        writeOpU32(code, OP_JMP, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case JMP_IF_TRUE -> {
                        writeOpU32(code, OP_JMP_IF_TRUE, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case JMP_IF_FALSE -> {
                        writeOpU32(code, OP_JMP_IF_FALSE, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case INTRINSIC -> {
                        writeOpU32(code, OP_INTRINSIC, op.immediate());
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case HOSTCALL -> {
                        final var decl = Objects.requireNonNull(op.syscallDecl(), "syscallDecl");
                        if (op.expectedArgSlots() != null && op.expectedArgSlots() != decl.argSlots()) {
                            throw new BytecodeMarshalingException(
                                    BytecodeMarshalingErrorCode.MARSHAL_LINKAGE_HOST_ABI_MISMATCH,
                                    "host arg_slots mismatch for " + decl.module() + "." + decl.name());
                        }
                        if (op.expectedRetSlots() != null && op.expectedRetSlots() != decl.retSlots()) {
                            throw new BytecodeMarshalingException(
                                    BytecodeMarshalingErrorCode.MARSHAL_LINKAGE_HOST_ABI_MISMATCH,
                                    "host ret_slots mismatch for " + decl.module() + "." + decl.name());
                        }

                        final var identity = new SyscallIdentity(decl.module(), decl.name(), decl.version());
                        if (!syscallIndexByIdentity.containsKey(identity)) {
                            syscallIndexByIdentity.put(identity, syscallIndexByIdentity.size());
                            orderedSyscalls.put(identity, decl);
                        }
                        final var index = syscallIndexByIdentity.get(identity);
                        writeOpU32(code, OP_HOSTCALL, index);
                        spans.add(new BytecodeFunctionLayoutBuilder.InstructionSpan(pc, spanOrUnknown(op.span())));
                    }
                    case RAW_SYSCALL -> throw new BytecodeMarshalingException(
                            BytecodeMarshalingErrorCode.MARSHAL_LINKAGE_RAW_SYSCALL_IN_PRELOAD,
                            "raw syscall is forbidden in pre-load artifact");
                }
            }
            functionFragments.add(new BytecodeFunctionLayoutBuilder.FunctionFragment(
                    function.name(),
                    code.toByteArray(),
                    function.paramSlots(),
                    function.localSlots(),
                    function.returnSlots(),
                    function.maxStackSlots(),
                    ReadOnlyList.wrap(spans)));
        }

        final var layout = BytecodeFunctionLayoutBuilder.build(ReadOnlyList.wrap(functionFragments));
        return new BytecodeModule(
                inputPlan.version(),
                canonicalConstPool,
                layout.functions(),
                layout.code(),
                layout.debugInfo(),
                canonicalExports,
                ReadOnlyList.wrap(orderedSyscalls.values()));
    }

    private static void writeOpNoImm(final ByteArrayOutputStream out, final int opcode) {
        out.writeBytes(ByteBuffer.allocate(2).order(ByteOrder.LITTLE_ENDIAN).putShort((short) opcode).array());
    }

    private static void writeOpU32(final ByteArrayOutputStream out, final int opcode, final int immediate) {
        out.writeBytes(ByteBuffer.allocate(2).order(ByteOrder.LITTLE_ENDIAN).putShort((short) opcode).array());
        out.writeBytes(ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(immediate).array());
    }

    private static void writeOpU8(final ByteArrayOutputStream out, final int opcode, final int immediate) {
        out.writeBytes(ByteBuffer.allocate(2).order(ByteOrder.LITTLE_ENDIAN).putShort((short) opcode).array());
        out.write(immediate & 0xFF);
    }

    private static BytecodeModule.SourceSpan spanOrUnknown(final BytecodeModule.SourceSpan span) {
        if (span == null) {
            return new BytecodeModule.SourceSpan(-1, 0, 0);
        }
        return span;
    }

    private ReadOnlyList<BytecodeModule.ConstantPoolEntry> dedupeConstPool(
            final ReadOnlyList<BytecodeModule.ConstantPoolEntry> input) {
        if (input == null || input.isEmpty()) {
            return ReadOnlyList.empty();
        }
        final var ordered = new LinkedHashSet<BytecodeModule.ConstantPoolEntry>(input.size());
        for (final var constant : input) {
            if (constant != null) {
                ordered.add(constant);
            }
        }
        return ReadOnlyList.wrap(ordered.stream().toList());
    }

    private ReadOnlyList<BytecodeModule.Export> dedupeExports(
            final ReadOnlyList<BytecodeModule.Export> input) {
        if (input == null || input.isEmpty()) {
            return ReadOnlyList.empty();
        }
        final var ordered = new LinkedHashSet<BytecodeModule.Export>(input.size());
        for (final var export : input) {
            if (export != null) {
                ordered.add(export);
            }
        }
        return ReadOnlyList.wrap(ordered.stream().toList());
    }

    private record SyscallIdentity(
            String module,
            String name,
            int version) {
    }

    public record EmissionPlan(
            int version,
            ReadOnlyList<BytecodeModule.ConstantPoolEntry> constPool,
            ReadOnlyList<BytecodeModule.Export> exports,
            ReadOnlyList<FunctionPlan> functions) {
        public EmissionPlan {
            constPool = constPool == null ? ReadOnlyList.empty() : constPool;
            exports = exports == null ? ReadOnlyList.empty() : exports;
            functions = functions == null ? ReadOnlyList.empty() : functions;
        }

        public static EmissionPlan empty() {
            return new EmissionPlan(0, ReadOnlyList.empty(), ReadOnlyList.empty(), ReadOnlyList.empty());
        }
    }

    public record FunctionPlan(
            String name,
            int paramSlots,
            int localSlots,
            int returnSlots,
            int maxStackSlots,
            ReadOnlyList<Operation> operations) {
        public FunctionPlan {
            Objects.requireNonNull(name, "name");
            operations = operations == null ? ReadOnlyList.empty() : operations;
        }
    }

    public enum OperationKind {
        HALT,
        RET,
        PUSH_CONST,
        POP,
        GET_LOCAL,
        GET_GLOBAL,
        SET_LOCAL,
        SET_GLOBAL,
        PUSH_BOOL,
        PUSH_I32,
        ADD,
        SUB,
        MUL,
        DIV,
        MOD,
        EQ,
        NEQ,
        LT,
        LTE,
        GT,
        GTE,
        AND,
        OR,
        NOT,
        NEG,
        CALL_FUNC,
        JMP,
        JMP_IF_TRUE,
        JMP_IF_FALSE,
        HOSTCALL,
        RAW_SYSCALL,
        INTRINSIC,
    }

    public record Operation(
            OperationKind kind,
            int immediate,
            BytecodeModule.SyscallDecl syscallDecl,
            Integer expectedArgSlots,
            Integer expectedRetSlots,
            BytecodeModule.SourceSpan span) {
        public Operation {
            Objects.requireNonNull(kind, "kind");
        }

        public static Operation halt() {
            return halt(null);
        }

        public static Operation halt(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.HALT, 0, null, null, null, span);
        }

        public static Operation ret() {
            return ret(null);
        }

        public static Operation ret(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.RET, 0, null, null, null, span);
        }

        public static Operation intrinsic(final int intrinsicId) {
            return intrinsic(intrinsicId, null);
        }

        public static Operation intrinsic(final int intrinsicId, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.INTRINSIC, intrinsicId, null, null, null, span);
        }

        public static Operation intrinsic(
                final int intrinsicId,
                final Integer expectedArgSlots,
                final Integer expectedRetSlots,
                final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.INTRINSIC, intrinsicId, null, expectedArgSlots, expectedRetSlots, span);
        }

        public static Operation callFunc(final int funcId) {
            return callFunc(funcId, null);
        }

        public static Operation callFunc(final int funcId, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.CALL_FUNC, funcId, null, null, null, span);
        }

        public static Operation pushConst(final int constPoolIndex, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.PUSH_CONST, constPoolIndex, null, null, null, span);
        }

        public static Operation pop(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.POP, 0, null, null, null, span);
        }

        public static Operation getLocal(final int slot, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.GET_LOCAL, slot, null, null, null, span);
        }

        public static Operation getGlobal(final int slot, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.GET_GLOBAL, slot, null, null, null, span);
        }

        public static Operation setLocal(final int slot, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.SET_LOCAL, slot, null, null, null, span);
        }

        public static Operation setGlobal(final int slot, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.SET_GLOBAL, slot, null, null, null, span);
        }

        public static Operation pushBool(final boolean value, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.PUSH_BOOL, value ? 1 : 0, null, null, null, span);
        }

        public static Operation pushI32(final int value, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.PUSH_I32, value, null, null, null, span);
        }

        public static Operation add(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.ADD, 0, null, null, null, span);
        }

        public static Operation sub(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.SUB, 0, null, null, null, span);
        }

        public static Operation mul(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.MUL, 0, null, null, null, span);
        }

        public static Operation div(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.DIV, 0, null, null, null, span);
        }

        public static Operation mod(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.MOD, 0, null, null, null, span);
        }

        public static Operation eq(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.EQ, 0, null, null, null, span);
        }

        public static Operation neq(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.NEQ, 0, null, null, null, span);
        }

        public static Operation lt(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.LT, 0, null, null, null, span);
        }

        public static Operation lte(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.LTE, 0, null, null, null, span);
        }

        public static Operation gt(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.GT, 0, null, null, null, span);
        }

        public static Operation gte(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.GTE, 0, null, null, null, span);
        }

        public static Operation and(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.AND, 0, null, null, null, span);
        }

        public static Operation or(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.OR, 0, null, null, null, span);
        }

        public static Operation not(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.NOT, 0, null, null, null, span);
        }

        public static Operation neg(final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.NEG, 0, null, null, null, span);
        }

        public static Operation hostcall(
                final BytecodeModule.SyscallDecl decl,
                final Integer expectedArgSlots,
                final Integer expectedRetSlots) {
            return hostcall(decl, expectedArgSlots, expectedRetSlots, null);
        }

        public static Operation hostcall(
                final BytecodeModule.SyscallDecl decl,
                final Integer expectedArgSlots,
                final Integer expectedRetSlots,
                final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.HOSTCALL, 0, decl, expectedArgSlots, expectedRetSlots, span);
        }

        public static Operation jmp(final int targetPc, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.JMP, targetPc, null, null, null, span);
        }

        public static Operation jmpIfTrue(final int targetPc, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.JMP_IF_TRUE, targetPc, null, null, null, span);
        }

        public static Operation jmpIfFalse(final int targetPc, final BytecodeModule.SourceSpan span) {
            return new Operation(OperationKind.JMP_IF_FALSE, targetPc, null, null, null, span);
        }

        public static Operation rawSyscall(final int syscallId) {
            return new Operation(OperationKind.RAW_SYSCALL, syscallId, null, null, null, null);
        }
    }
}
