import java.lang.foreign.Arena;
import java.lang.foreign.FunctionDescriptor;
import java.lang.foreign.Linker;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.SymbolLookup;
import java.lang.invoke.MethodHandle;
import java.nio.charset.StandardCharsets;
import java.nio.file.Path;

import static java.lang.foreign.ValueLayout.ADDRESS;
import static java.lang.foreign.ValueLayout.JAVA_BYTE;
import static java.lang.foreign.ValueLayout.JAVA_INT;
import static java.lang.foreign.ValueLayout.JAVA_LONG;

/// Panama downcalls into `libgreeter_ffi.so`. No upcalls: Rust never calls Java.
///
/// Every function is a plain synchronous call that any thread may make. The
/// Rust side enforces one-at-a-time per handle and reports `BUSY` otherwise.
final class Native {
    static final int AWAITING = 0;
    static final int COMPLETE = 1;
    static final int STALLED = 2;
    static final int BUSY = -1;
    static final int FINISHED = -2;
    static final int REPLY_MISSING = -3;
    static final int PANICKED = -4;
    static final int BAD_HANDLE = -5;
    static final int BAD_INPUT = -6;

    private final MethodHandle bufFree;
    private final MethodHandle capabilitiesSpawn;
    private final MethodHandle effectRoutineSpawn;
    private final MethodHandle effectRoutineResume;
    private final MethodHandle effectRoutineFree;
    private final MethodHandle replyString;
    private final MethodHandle replyU32;
    private final MethodHandle replyUnit;
    private final MethodHandle sansioNew;
    private final MethodHandle sansioStep;
    private final MethodHandle sansioFree;

    Native(Path library) {
        var linker = Linker.nativeLinker();
        var lookup = SymbolLookup.libraryLookup(library, Arena.global());

        bufFree = bind(linker, lookup, "greeter_buf_free", FunctionDescriptor.ofVoid(ADDRESS, JAVA_LONG));
        capabilitiesSpawn = bind(linker, lookup, "capabilities_spawn", FunctionDescriptor.of(JAVA_LONG));
        effectRoutineSpawn = bind(linker, lookup, "effect_routine_spawn", FunctionDescriptor.of(JAVA_LONG));
        effectRoutineResume = bind(linker, lookup, "effect_routine_resume", FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, ADDRESS));
        effectRoutineFree = bind(linker, lookup, "effect_routine_free", FunctionDescriptor.of(JAVA_INT, JAVA_LONG));
        replyString = bind(linker, lookup, "effect_routine_reply_string", FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, JAVA_LONG));
        replyU32 = bind(linker, lookup, "effect_routine_reply_u32", FunctionDescriptor.of(JAVA_INT, JAVA_LONG, JAVA_INT));
        replyUnit = bind(linker, lookup, "effect_routine_reply_unit", FunctionDescriptor.of(JAVA_INT, JAVA_LONG));
        sansioNew = bind(linker, lookup, "sansio_new", FunctionDescriptor.of(JAVA_LONG));
        sansioStep = bind(linker, lookup, "sansio_step", FunctionDescriptor.of(JAVA_INT, JAVA_LONG, ADDRESS, JAVA_LONG, ADDRESS, ADDRESS));
        sansioFree = bind(linker, lookup, "sansio_free", FunctionDescriptor.of(JAVA_INT, JAVA_LONG));
    }

    private static MethodHandle bind(Linker linker, SymbolLookup lookup, String name, FunctionDescriptor descriptor) {
        var symbol = lookup.find(name).orElseThrow(() -> new IllegalStateException("missing symbol " + name));
        return linker.downcallHandle(symbol, descriptor);
    }

    /// A status code plus the bytes Rust handed back (already copied and freed).
    record Out(int status, byte[] bytes) {}

    /// Spawn the capabilities greeter. The handle is driven with the `effectRoutine*` calls.
    long capabilitiesSpawn() {
        return (long) invoke(capabilitiesSpawn);
    }

    long effectRoutineSpawn() {
        return (long) invoke(effectRoutineSpawn);
    }

    Out effectRoutineResume(long handle) {
        try (var arena = Arena.ofConfined()) {
            var outPtr = arena.allocate(ADDRESS);
            var outLen = arena.allocate(JAVA_LONG);
            int status = (int) invoke(effectRoutineResume, handle, outPtr, outLen);
            return new Out(status, status >= 0 ? takeBuffer(outPtr, outLen) : new byte[0]);
        }
    }

    int effectRoutineFree(long handle) {
        return (int) invoke(effectRoutineFree, handle);
    }

    int effectRoutineReplyString(long handle, String value) {
        try (var arena = Arena.ofConfined()) {
            var bytes = value.getBytes(StandardCharsets.UTF_8);
            var segment = arena.allocateArray(JAVA_BYTE, bytes);
            return (int) invoke(replyString, handle, segment, (long) bytes.length);
        }
    }

    int effectRoutineReplyU32(long handle, int value) {
        return (int) invoke(replyU32, handle, value);
    }

    int effectRoutineReplyUnit(long handle) {
        return (int) invoke(replyUnit, handle);
    }

    long sansioNew() {
        return (long) invoke(sansioNew);
    }

    Out sansioStep(long handle, byte[] input) {
        try (var arena = Arena.ofConfined()) {
            var in = arena.allocateArray(JAVA_BYTE, input);
            var outPtr = arena.allocate(ADDRESS);
            var outLen = arena.allocate(JAVA_LONG);
            int status = (int) invoke(sansioStep, handle, in, (long) input.length, outPtr, outLen);
            return new Out(status, status >= 0 ? takeBuffer(outPtr, outLen) : new byte[0]);
        }
    }

    int sansioFree(long handle) {
        return (int) invoke(sansioFree, handle);
    }

    /// Copy a Rust-owned buffer out of the out-pointers and free it.
    private byte[] takeBuffer(MemorySegment outPtr, MemorySegment outLen) {
        long len = outLen.get(JAVA_LONG, 0);
        var ptr = outPtr.get(ADDRESS, 0);
        var bytes = ptr.reinterpret(len).toArray(JAVA_BYTE);
        invoke(bufFree, ptr, len);
        return bytes;
    }

    private static Object invoke(MethodHandle handle, Object... args) {
        try {
            return handle.invokeWithArguments(args);
        } catch (Throwable t) {
            throw new IllegalStateException("native call failed", t);
        }
    }

    static String describe(int code) {
        return switch (code) {
            case AWAITING -> "AWAITING";
            case COMPLETE -> "COMPLETE";
            case STALLED -> "STALLED";
            case BUSY -> "BUSY";
            case FINISHED -> "FINISHED";
            case REPLY_MISSING -> "REPLY_MISSING";
            case PANICKED -> "PANICKED";
            case BAD_HANDLE -> "BAD_HANDLE";
            case BAD_INPUT -> "BAD_INPUT";
            default -> "code " + code;
        };
    }
}