{{ header }}package {{ package }};

import java.lang.foreign.Arena;
import java.lang.foreign.FunctionDescriptor;
import java.lang.foreign.Linker;
import java.lang.foreign.MemoryLayout;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.util.ArrayList;
import java.util.List;

/**
 * Allocates Panama FFM upcall stubs for a Visitor and assembles
 * the C HTMHtmVisitorCallbacks struct in native memory.
 */
final class VisitorBridge implements AutoCloseable {
    private static final Linker LINKER = Linker.nativeLinker();
    private static final MethodHandles.Lookup LOOKUP =
        MethodHandles.lookup();

    // VisitResult discriminant codes returned to C
    private static final int VISIT_RESULT_CONTINUE = 0;
    private static final int VISIT_RESULT_SKIP = 1;
    private static final int VISIT_RESULT_PRESERVE_HTML = 2;
    private static final int VISIT_RESULT_CUSTOM = 3;
    private static final int VISIT_RESULT_ERROR = 4;

    // HTMHtmVisitorCallbacks: user_data + {{ num_callbacks }} callbacks
    // = {{ num_fields }} pointer-sized slots
    private static final long CALLBACKS_STRUCT_SIZE =
        (long) ValueLayout.ADDRESS.byteSize() * {{ num_fields }}L;

    // HTMHtmNodeContext field offsets
    private static final long CTX_OFFSET_TAG_NAME = 8L;
    private static final long CTX_OFFSET_DEPTH = 16L;
    private static final long CTX_OFFSET_INDEX_IN_PARENT = 24L;
    private static final long CTX_OFFSET_PARENT_TAG = 32L;
    private static final long CTX_OFFSET_IS_INLINE = 40L;

    private final Arena arena;
    private final MemorySegment struct;
    private final Visitor visitor;
    /** Sticky exception captured during a visitor upcall; surfaced after conversion. */
    private volatile Throwable visitorError;

    VisitorBridge(Visitor visitor) {
        this.visitor = visitor;
        this.arena = Arena.ofConfined();
        this.struct = arena.allocate(CALLBACKS_STRUCT_SIZE);
        // Slot 0: user_data = NULL
        // (visitor captured via lambda closure)
        struct.set(ValueLayout.ADDRESS, 0L, MemorySegment.NULL);
        try {
            long offset = ValueLayout.ADDRESS.byteSize();
{% for stub_call in stub_calls %}
            offset = {{ stub_call }};
{% endfor %}
        } catch (ReflectiveOperationException e) {
            arena.close();
            throw new RuntimeException(
                "Failed to create visitor upcall stubs", e);
        }
    }

{% for stub_method in stub_methods %}
{{ stub_method }}
{% endfor %}

    /** Returns the native HTMHtmVisitorCallbacks struct pointer. */
    MemorySegment callbacksStruct() {
        return struct;
    }

    void rethrowVisitorError() throws Throwable {
        Throwable err = visitorError;
        if (err != null) {
            throw err;
        }
    }

{% for handle_method in handle_methods %}
{{ handle_method }}
{% endfor %}

    // HTMHtmNodeContext: int32 node_type, char* tag_name, uintptr depth,
    //   uintptr index_in_parent, char* parent_tag,
    //   int32 is_inline
    private static final MemoryLayout CTX_LAYOUT =
        MemoryLayout.structLayout(
            ValueLayout.JAVA_INT.withName("node_type"),
            MemoryLayout.paddingLayout(4),
            ValueLayout.ADDRESS.withName("tag_name"),
            ValueLayout.JAVA_LONG.withName("depth"),
            ValueLayout.JAVA_LONG.withName("index_in_parent"),
            ValueLayout.ADDRESS.withName("parent_tag"),
            ValueLayout.JAVA_INT.withName("is_inline"),
            MemoryLayout.paddingLayout(4)
    );

    private static NodeContext decodeNodeContext(
            final MemorySegment ctxPtr) {
        var ctx = ctxPtr.reinterpret(
            CTX_LAYOUT.byteSize());
        int nodeType = ctx.get(
            ValueLayout.JAVA_INT, 0L);
        var tagNamePtr = ctx.get(
            ValueLayout.ADDRESS, CTX_OFFSET_TAG_NAME);
        String tagName = tagNamePtr
            .reinterpret(Long.MAX_VALUE).getString(0);
        long depth = ctx.get(ValueLayout.JAVA_LONG, CTX_OFFSET_DEPTH);
        long indexInParent = ctx.get(ValueLayout.JAVA_LONG, CTX_OFFSET_INDEX_IN_PARENT);
        var parentTagPtr = ctx.get(ValueLayout.ADDRESS, CTX_OFFSET_PARENT_TAG);
        String parentTag = parentTagPtr.equals(MemorySegment.NULL) ? null
                : parentTagPtr.reinterpret(Long.MAX_VALUE).getString(0);
        int isInlineRaw = ctx.get(ValueLayout.JAVA_INT, CTX_OFFSET_IS_INLINE);
        return new NodeContext(nodeType, tagName, depth, indexInParent, parentTag, isInlineRaw != 0);
    }

    private static List<String> decodeCells(MemorySegment cellsPtr, long count) {
        var result = new ArrayList<String>((int) count);
        for (long i = 0; i < count; i++) {
            var ptr = cellsPtr.get(ValueLayout.ADDRESS, i * ValueLayout.ADDRESS.byteSize());
            result.add(ptr.reinterpret(Long.MAX_VALUE).getString(0));
        }
        return result;
    }

    private int encodeVisitResult(VisitResult result, MemorySegment outCustom, MemorySegment outLen) {
        return switch (result) {
            case VisitResult.Continue ignored -> VISIT_RESULT_CONTINUE;
            case VisitResult.Skip ignored -> VISIT_RESULT_SKIP;
            case VisitResult.PreserveHtml ignored -> VISIT_RESULT_PRESERVE_HTML;
            case VisitResult.Custom c -> {
                var buf = Arena.global().allocateFrom(c.markdown());
                outCustom.reinterpret(ValueLayout.ADDRESS.byteSize()).set(ValueLayout.ADDRESS, 0L, buf);
                long customLen = (long) c.markdown().getBytes(java.nio.charset.StandardCharsets.UTF_8).length;
                outLen.reinterpret(ValueLayout.JAVA_LONG.byteSize()).set(ValueLayout.JAVA_LONG, 0L, customLen);
                yield VISIT_RESULT_CUSTOM;
            }
            case VisitResult.Error e -> {
                var buf = Arena.global().allocateFrom(e.message());
                outCustom.reinterpret(ValueLayout.ADDRESS.byteSize()).set(ValueLayout.ADDRESS, 0L, buf);
                long errorLen = (long) e.message().getBytes(java.nio.charset.StandardCharsets.UTF_8).length;
                outLen.reinterpret(ValueLayout.JAVA_LONG.byteSize()).set(ValueLayout.JAVA_LONG, 0L, errorLen);
                yield VISIT_RESULT_ERROR;
            }
        };
    }

    @Override
    public void close() {
        arena.close();
    }
}
