// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

//go:build amd64

package verify

import (
	"encoding/binary"
	"fmt"
	"reflect"
	"syscall"
	"unsafe"
)

// savedSP holds the Go stack pointer while a JIT call is in flight.
// Written and read by the assembly trampoline (trampoline_amd64.s); no Go
// code references it, which staticcheck and GoLand cannot see.
//
// noinspection GoUnusedGlobalVariable
//
//lint:ignore U1000 written and read by the assembly
var savedSP uintptr

// enterJIT switches to the prepared stack and jumps to fn.
// It does not return normally; the JIT function's RET transfers control
// to leaveJIT, which restores the Go stack.  The body lives in
// trampoline_amd64.s and reads the parameters from the frame by name.
//
// noinspection GoUnusedParameter
//
//go:nosplit
func enterJIT(fn uintptr, stack uintptr)

// leaveJIT restores the Go stack after a JIT function returns.
// Its address is placed as the return address on the prepared stack.
//
//go:nosplit
func leaveJIT()

// leaveJITAddr is the machine address of leaveJIT, resolved once at init.
var leaveJITAddr uintptr

func init() {
	leaveJITAddr = reflect.ValueOf(leaveJIT).Pointer()
}

// stackPad is padding below the return address on the prepared stack.
// The ABIInternal wrapper that leaveJIT's address resolves to executes
// PUSHQ BP and CALL before reaching the raw assembly, writing up to 16
// bytes below the return-address slot.  64 bytes of headroom is ample.
const stackPad = 64

// Call invokes the assembled function at fnAddr with the given ABI0 argument
// block (the raw bytes that would appear at FP+0).  It returns the argument
// block after the call, which contains any results the function wrote back
// (the ABI0 convention shares the argument area for inputs and outputs).
//
// The function must be NOSPLIT (no stack growth) and must not reference
// external symbols, the image is self-contained.
//
// Not safe for concurrent use: only one JIT call may be in flight at a
// time, the trampolines keep the saved registers in package globals
// (savedSP and friends).
func Call(fnAddr uintptr, args []byte) ([]byte, error) {
	// Prepare the stack: [padding][leaveJIT addr][args...]
	stackSize := stackPad + 8 + len(args) + 64 // padding + ret + args + safety
	stackMem, err := syscall.Mmap(-1, 0, stackSize,
		syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON)
	if err != nil {
		return nil, fmt.Errorf("verify: stack mmap: %w", err)
	}
	defer syscall.Munmap(stackMem)

	// The return address sits after the padding; the function's SP will
	// point here, leaving stackPad bytes below for the wrapper's pushes.
	retOff := stackPad
	binary.LittleEndian.PutUint64(stackMem[retOff:retOff+8], uint64(leaveJITAddr))
	// The ABI0 argument area follows the return address.
	copy(stackMem[retOff+8:], args)

	stackBase := uintptr(unsafe.Pointer(&stackMem[retOff]))
	enterJIT(fnAddr, stackBase)

	// Copy out the (possibly modified) argument area.
	out := make([]byte, len(args))
	copy(out, stackMem[retOff+8:retOff+8+len(args)])
	return out, nil
}