package com.codeflash.profiler;

import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.AdviceAdapter;

/**
 * ASM MethodVisitor that injects line-level profiling probes.
 *
 * <p>At each {@code LineNumber} table entry within the target method:
 * <ol>
 *   <li>Registers the line with {@link ProfilerRegistry} (happens once at class-load time)</li>
 *   <li>Injects bytecode: {@code LDC globalId; INVOKESTATIC ProfilerData.hit(I)V}</li>
 * </ol>
 *
 * <p>At method entry: injects a warmup self-call loop (if warmup is configured) followed by
 * {@code ProfilerData.enterMethod(entryLineId)}.
 * <p>At method exit (every RETURN/ATHROW): injects {@code ProfilerData.exitMethod()}.
 */
public class LineProfilingMethodVisitor extends AdviceAdapter {

    private static final String PROFILER_DATA = "com/codeflash/profiler/ProfilerData";

    private final String internalClassName;
    private final String sourceFile;
    private final String methodName;
    private boolean firstLineVisited = false;

    protected LineProfilingMethodVisitor(
            MethodVisitor mv, int access, String name, String descriptor,
            String internalClassName, String sourceFile) {
        super(Opcodes.ASM9, mv, access, name, descriptor);
        this.internalClassName = internalClassName;
        this.sourceFile = sourceFile;
        this.methodName = name;
    }

    /**
     * Inject a warmup self-call loop at method entry.
     *
     * <p>Generated bytecode equivalent:
     * <pre>
     * if (ProfilerData.isWarmupNeeded()) {
     *     ProfilerData.startWarmup();
     *     for (int i = 0; i &lt; ProfilerData.getWarmupThreshold(); i++) {
     *         thisMethod(originalArgs);
     *     }
     *     ProfilerData.finishWarmup();
     * }
     * </pre>
     *
     * <p>Recursive warmup calls re-enter this method but {@code isWarmupNeeded()} returns
     * {@code false} (guard flag set by {@code startWarmup()}), so they execute the normal
     * instrumented body. After the loop, {@code finishWarmup()} zeros all counters so the
     * next real execution records clean data.
     */
    @Override
    protected void onMethodEnter() {
        Label skipWarmup = new Label();

        // if (!ProfilerData.isWarmupNeeded()) goto skipWarmup
        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "isWarmupNeeded", "()Z", false);
        mv.visitJumpInsn(IFEQ, skipWarmup);

        // ProfilerData.startWarmup()
        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "startWarmup", "()V", false);

        // int _warmupIdx = 0
        int counterLocal = newLocal(Type.INT_TYPE);
        mv.visitInsn(ICONST_0);
        mv.visitVarInsn(ISTORE, counterLocal);

        Label loopCheck = new Label();
        Label loopBody = new Label();

        mv.visitJumpInsn(GOTO, loopCheck);

        // loop body: call self with original arguments
        mv.visitLabel(loopBody);

        boolean isStatic = (methodAccess & Opcodes.ACC_STATIC) != 0;
        if (!isStatic) {
            loadThis();
        }
        loadArgs();

        int invokeOp;
        if (isStatic) {
            invokeOp = INVOKESTATIC;
        } else if ((methodAccess & Opcodes.ACC_PRIVATE) != 0) {
            invokeOp = INVOKESPECIAL;
        } else {
            invokeOp = INVOKEVIRTUAL;
        }
        mv.visitMethodInsn(invokeOp, internalClassName, methodName, methodDesc, false);

        // Discard return value
        Type returnType = Type.getReturnType(methodDesc);
        switch (returnType.getSort()) {
            case Type.VOID:
                break;
            case Type.LONG:
            case Type.DOUBLE:
                mv.visitInsn(POP2);
                break;
            default:
                mv.visitInsn(POP);
                break;
        }

        // _warmupIdx++
        mv.visitIincInsn(counterLocal, 1);

        // loop check: _warmupIdx < ProfilerData.getWarmupThreshold()
        mv.visitLabel(loopCheck);
        mv.visitVarInsn(ILOAD, counterLocal);
        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "getWarmupThreshold", "()I", false);
        mv.visitJumpInsn(IF_ICMPLT, loopBody);

        // ProfilerData.finishWarmup()
        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "finishWarmup", "()V", false);

        mv.visitLabel(skipWarmup);
    }

    @Override
    public void visitLineNumber(int line, Label start) {
        super.visitLineNumber(line, start);

        // Register this line and get its global ID (happens once at class-load time)
        String dotClassName = internalClassName.replace('/', '.');
        int globalId = ProfilerRegistry.register(sourceFile, dotClassName, methodName, line);

        if (!firstLineVisited) {
            firstLineVisited = true;
            // Inject enterMethod call at the first line of the method
            mv.visitLdcInsn(globalId);
            mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "enterMethod", "(I)V", false);
        }

        // Inject: ProfilerData.hit(globalId)
        mv.visitLdcInsn(globalId);
        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "hit", "(I)V", false);
    }

    @Override
    protected void onMethodExit(int opcode) {
        // Before every RETURN or ATHROW, flush timing for the last line
        // This fixes the "last line always shows 0ms" bug
        mv.visitMethodInsn(INVOKESTATIC, PROFILER_DATA, "exitMethod", "()V", false);
    }
}
