package p.studio.compiler.backend.irvm;

import org.junit.jupiter.api.Test;
import p.studio.compiler.backend.bytecode.BytecodeEmitter;
import p.studio.compiler.backend.bytecode.BytecodeModule;
import p.studio.compiler.backend.bytecode.BytecodePreloadVerifierService;
import p.studio.compiler.models.IRBackend;
import p.studio.compiler.models.IRBackendExecutableFunction;
import p.studio.compiler.source.Span;
import p.studio.compiler.source.identifiers.CallableId;
import p.studio.compiler.source.identifiers.FileId;
import p.studio.compiler.source.identifiers.IntrinsicId;
import p.studio.compiler.source.tables.IntrinsicReference;
import p.studio.utilities.structures.ReadOnlyList;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;

class OptimizeIRVMEquivalenceHarnessTest {

    @Test
    void optimizeOnOffMustPreserveObservableTraceForLoweredHostIntrinsicFixture() {
        final var backend = IRBackend.builder()
                .entryPointModuleId(new p.studio.compiler.source.identifiers.ModuleId(0))
                .executableFunctions(ReadOnlyList.from(
                        fn("main", 1, ReadOnlyList.from(
                                callHost("gfx", "draw_pixel", 1, 0, 0),
                                callIntrinsic("input.pad", 1, 0, 0, 0),
                                ret()))))
                .intrinsicPool(ReadOnlyList.from(new IntrinsicReference("input.pad", 1)))
                .build();

        final var lowered = new LowerToIRVMService().lower(backend);
        assertEquivalentObservableBehavior(lowered);
    }

    @Test
    void optimizeOnOffMustPreserveObservableTraceForConditionalJoinFixture() {
        final var program = singleFunctionProgram(
                "main",
                1,
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.INTRINSIC, 901),
                        new IRVMInstruction(IRVMOp.JMP_IF_FALSE, 24),
                        new IRVMInstruction(IRVMOp.HOSTCALL, 0),
                        new IRVMInstruction(IRVMOp.JMP, 30),
                        new IRVMInstruction(IRVMOp.HOSTCALL, 0),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.intrinsic(901, 0, 1, null),
                        BytecodeEmitter.Operation.jmpIfFalse(24, null),
                        hostcallOp("gfx", "draw_true_branch", 1, 0, 0),
                        BytecodeEmitter.Operation.jmp(30, null),
                        hostcallOp("gfx", "draw_false_branch", 1, 0, 0),
                        BytecodeEmitter.Operation.ret()));
        assertEquivalentObservableBehavior(program);
    }

    @Test
    void optimizeOnOffMustPreserveObservableTraceForSimpleLoopFixture() {
        final var program = singleFunctionProgram(
                "main",
                1,
                ReadOnlyList.from(
                        new IRVMInstruction(IRVMOp.INTRINSIC, 900),
                        new IRVMInstruction(IRVMOp.JMP_IF_FALSE, 24),
                        new IRVMInstruction(IRVMOp.HOSTCALL, 0),
                        new IRVMInstruction(IRVMOp.JMP, 0),
                        new IRVMInstruction(IRVMOp.RET, null)),
                ReadOnlyList.from(
                        BytecodeEmitter.Operation.intrinsic(900, 0, 1, null),
                        BytecodeEmitter.Operation.jmpIfFalse(24, null),
                        hostcallOp("gfx", "loop_tick", 1, 0, 0),
                        BytecodeEmitter.Operation.jmp(0, null),
                        BytecodeEmitter.Operation.ret()));
        assertEquivalentObservableBehavior(program);
    }

    @Test
    void optimizeOnOffMustPreserveObservableTraceForLinearCallFixture() {
        final var mainInstructions = ReadOnlyList.from(
                new IRVMInstruction(IRVMOp.CALL, 1),
                new IRVMInstruction(IRVMOp.HOSTCALL, 0),
                new IRVMInstruction(IRVMOp.RET, null));
        final var workerInstructions = ReadOnlyList.from(
                new IRVMInstruction(IRVMOp.HOSTCALL, 0),
                new IRVMInstruction(IRVMOp.RET, null));
        final var mainOps = ReadOnlyList.from(
                BytecodeEmitter.Operation.callFunc(1),
                hostcallOp("gfx", "draw_main", 1, 0, 0),
                BytecodeEmitter.Operation.ret());
        final var workerOps = ReadOnlyList.from(
                hostcallOp("gfx", "draw_worker", 1, 0, 0),
                BytecodeEmitter.Operation.ret());
        final var program = multiFunctionProgram(
                ReadOnlyList.from(
                        irvmFunction("main", 0, mainInstructions),
                        irvmFunction("worker", 0, workerInstructions)),
                ReadOnlyList.from(
                        functionPlan("main", 0, mainOps),
                        functionPlan("worker", 0, workerOps)));

        assertEquivalentObservableBehavior(program);
    }

    private void assertEquivalentObservableBehavior(final IRVMProgram input) {
        final var optimizedOff = new OptimizeIRVMService(new IRVMValidator(), List.of()).optimize(input);
        final var optimizedOn = new OptimizeIRVMService().optimize(input);
        final var optimizedOnAgain = new OptimizeIRVMService().optimize(input);

        final var offResult = execute(optimizedOff);
        final var onResult = execute(optimizedOn);
        assertEquals(offResult, onResult);
        assertEquals(optimizedOn, optimizedOnAgain);
        assertFalse(optimizedOn.hasInternalOpcodes());
        assertFalse(optimizedOff.hasInternalOpcodes());

        final var emitter = new BytecodeEmitter();
        final var verifier = new BytecodePreloadVerifierService();
        verifier.verify(emitter.emit(optimizedOff.coherentEmissionPlan()));
        verifier.verify(emitter.emit(optimizedOn.coherentEmissionPlan()));
    }

    private ExecutionResult execute(final IRVMProgram program) {
        final var module = program.module();
        if (module.functions().isEmpty()) {
            return new ExecutionResult(ReadOnlyList.empty(), true, ReadOnlyList.empty());
        }

        final var plan = program.coherentEmissionPlan();
        final var trace = new ArrayList<String>();
        final var stack = new ArrayDeque<Integer>();
        final var frames = new ArrayDeque<Frame>();
        final var indexByPcByFunction = new ArrayList<HashMap<Integer, Integer>>(module.functions().size());
        final var functionPlans = new ArrayList<BytecodeEmitter.FunctionPlan>(plan.functions().size());
        final var intrinsicInvocationCount = new HashMap<Integer, Integer>();
        for (final var function : module.functions()) {
            indexByPcByFunction.add(indexByPc(function.instructions()));
        }
        for (final var functionPlan : plan.functions()) {
            functionPlans.add(functionPlan);
        }
        frames.push(new Frame(0, 0));

        var steps = 0;
        while (!frames.isEmpty() && steps < 8_000) {
            steps++;
            final var frame = frames.peek();
            final var function = module.functions().get(frame.functionIndex());
            final var functionPlan = functionPlans.get(frame.functionIndex());
            final var instructions = function.instructions();
            if (frame.instructionIndex() < 0 || frame.instructionIndex() >= instructions.size()) {
                throw new IllegalStateException("invalid pc index during interpretation");
            }
            final var instruction = instructions.get(frame.instructionIndex());
            final var operation = functionPlan.operations().get(frame.instructionIndex());
            final var op = instruction.op();
            if (op == IRVMOp.HALT) {
                trace.add("HALT");
                return new ExecutionResult(ReadOnlyList.wrap(trace), true, stackSnapshot(stack));
            } else if (op == IRVMOp.RET) {
                trace.add("RET:" + function.name());
                frames.pop();
                if (frames.isEmpty()) {
                    trace.add("HALT");
                    return new ExecutionResult(ReadOnlyList.wrap(trace), true, stackSnapshot(stack));
                }
            } else if (op == IRVMOp.CALL) {
                trace.add("CALL#" + instruction.immediate());
                frame.advance();
                frames.push(new Frame(instruction.immediate(), 0));
            } else if (op == IRVMOp.JMP) {
                frame.jump(indexByPcByFunction.get(frame.functionIndex()), instruction.immediate());
            } else if (op == IRVMOp.JMP_IF_TRUE) {
                final var value = stack.isEmpty() ? 0 : stack.pop();
                if (value != 0) {
                    frame.jump(indexByPcByFunction.get(frame.functionIndex()), instruction.immediate());
                } else {
                    frame.advance();
                }
            } else if (op == IRVMOp.JMP_IF_FALSE) {
                final var value = stack.isEmpty() ? 0 : stack.pop();
                if (value == 0) {
                    frame.jump(indexByPcByFunction.get(frame.functionIndex()), instruction.immediate());
                } else {
                    frame.advance();
                }
            } else if (op == IRVMOp.HOSTCALL) {
                final var hostDecl = operation.syscallDecl();
                if (hostDecl != null) {
                    trace.add("HOSTCALL:" + hostDecl.module() + "." + hostDecl.name() + "@" + hostDecl.version());
                } else {
                    trace.add("HOSTCALL");
                }
                applyOperationStackEffect(stack, operation);
                frame.advance();
            } else if (op == IRVMOp.INTRINSIC) {
                trace.add("INTRINSIC#" + instruction.immediate());
                applyOperationStackEffect(stack, operation);
                final var pushes = operation.expectedRetSlots() == null ? 0 : operation.expectedRetSlots();
                for (var i = 0; i < pushes; i++) {
                    final var count = intrinsicInvocationCount.getOrDefault(instruction.immediate(), 0);
                    intrinsicInvocationCount.put(instruction.immediate(), count + 1);
                    stack.push(simulatedIntrinsicValue(instruction.immediate(), count));
                }
                frame.advance();
            } else {
                throw new IllegalStateException("unsupported op in equivalence harness: " + instruction.op().name());
            }
        }
        return new ExecutionResult(ReadOnlyList.wrap(trace), false, stackSnapshot(stack));
    }

    private HashMap<Integer, Integer> indexByPc(final ReadOnlyList<IRVMInstruction> instructions) {
        final var map = new HashMap<Integer, Integer>();
        var pc = 0;
        for (var i = 0; i < instructions.size(); i++) {
            map.put(pc, i);
            pc += instructions.get(i).encodedSize();
        }
        return map;
    }

    private void applyOperationStackEffect(
            final ArrayDeque<Integer> stack,
            final BytecodeEmitter.Operation operation) {
        final var pops = operation.expectedArgSlots() == null ? 0 : operation.expectedArgSlots();
        for (var i = 0; i < pops; i++) {
            if (stack.isEmpty()) {
                stack.push(0);
            }
            stack.pop();
        }
    }

    private int simulatedIntrinsicValue(
            final int intrinsicId,
            final int invocationIndex) {
        if (intrinsicId == 900) {
            return invocationIndex == 0 ? 1 : 0;
        }
        if (intrinsicId == 901) {
            return 0;
        }
        return 1;
    }

    private ReadOnlyList<Integer> stackSnapshot(final ArrayDeque<Integer> stack) {
        return ReadOnlyList.wrap(new ArrayList<>(stack));
    }

    private IRVMProgram singleFunctionProgram(
            final String name,
            final int maxStackSlots,
            final ReadOnlyList<IRVMInstruction> instructions,
            final ReadOnlyList<BytecodeEmitter.Operation> operations) {
        return multiFunctionProgram(
                ReadOnlyList.from(irvmFunction(name, maxStackSlots, instructions)),
                ReadOnlyList.from(functionPlan(name, maxStackSlots, operations)));
    }

    private IRVMProgram multiFunctionProgram(
            final ReadOnlyList<IRVMFunction> functions,
            final ReadOnlyList<BytecodeEmitter.FunctionPlan> functionPlans) {
        return new IRVMProgram(
                new IRVMModule("core-v1", functions),
                new BytecodeEmitter.EmissionPlan(
                        0,
                        ReadOnlyList.empty(),
                        ReadOnlyList.empty(),
                        functionPlans));
    }

    private IRVMFunction irvmFunction(
            final String name,
            final int maxStackSlots,
            final ReadOnlyList<IRVMInstruction> instructions) {
        return new IRVMFunction(
                name,
                0,
                0,
                0,
                maxStackSlots,
                instructions);
    }

    private BytecodeEmitter.FunctionPlan functionPlan(
            final String name,
            final int maxStackSlots,
            final ReadOnlyList<BytecodeEmitter.Operation> operations) {
        return new BytecodeEmitter.FunctionPlan(
                name,
                0,
                0,
                0,
                maxStackSlots,
                operations);
    }

    private BytecodeEmitter.Operation hostcallOp(
            final String module,
            final String name,
            final int version,
            final int argSlots,
            final int retSlots) {
        return BytecodeEmitter.Operation.hostcall(
                new BytecodeModule.SyscallDecl(module, name, version, argSlots, retSlots),
                argSlots,
                retSlots,
                null);
    }

    private IRBackendExecutableFunction fn(
            final String name,
            final int callableId,
            final ReadOnlyList<IRBackendExecutableFunction.Instruction> instructions) {
        return new IRBackendExecutableFunction(
                new FileId(1),
                new p.studio.compiler.source.identifiers.ModuleId(0),
                name,
                new CallableId(callableId),
                0,
                10,
                0,
                0,
                0,
                2,
                instructions,
                Span.none());
    }

    private IRBackendExecutableFunction.Instruction ret() {
        return new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.RET,
                "",
                null,
                null,
                Span.none());
    }

    private IRBackendExecutableFunction.Instruction callHost(
            final String module,
            final String name,
            final long version,
            final int argSlots,
            final int retSlots) {
        return new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.CALL_HOST,
                "",
                new IRBackendExecutableFunction.HostCallMetadata(module, name, version, argSlots, retSlots),
                null,
                Span.none());
    }

    private IRBackendExecutableFunction.Instruction callIntrinsic(
            final String canonicalName,
            final long canonicalVersion,
            final int intrinsicId,
            final int argSlots,
            final int retSlots) {
        return new IRBackendExecutableFunction.Instruction(
                IRBackendExecutableFunction.InstructionKind.CALL_INTRINSIC,
                "",
                null,
                null,
                new IRBackendExecutableFunction.IntrinsicCallMetadata(canonicalName, canonicalVersion, new IntrinsicId(intrinsicId)),
                argSlots,
                retSlots,
                Span.none());
    }

    private record ExecutionResult(
            ReadOnlyList<String> trace,
            boolean halted,
            ReadOnlyList<Integer> stack) {
    }

    private static final class Frame {
        private final int functionIndex;
        private int instructionIndex;

        private Frame(
                final int functionIndex,
                final int instructionIndex) {
            this.functionIndex = functionIndex;
            this.instructionIndex = instructionIndex;
        }

        int functionIndex() {
            return functionIndex;
        }

        int instructionIndex() {
            return instructionIndex;
        }

        void advance() {
            instructionIndex++;
        }

        void jump(
                final HashMap<Integer, Integer> indexByPc,
                final int targetPc) {
            final var targetIndex = indexByPc.get(targetPc);
            if (targetIndex == null) {
                throw new IllegalStateException("invalid jump target for harness");
            }
            instructionIndex = targetIndex;
        }
    }
}
