package spike

import (
	"fmt"

	ort "github.com/yalue/onnxruntime_go"
)

const (
	poseInputSize   = 640
	poseInputCount  = 1 * 3 * poseInputSize * poseInputSize
	poseOutputCount = 1 * 56 * 8400
)

// RunPose executes the locked YOLO pose ONNX graph once using the CPU runtime.
// The caller supplies absolute model and DLL paths so no camera credential or
// machine-specific path is retained in V2 configuration or source.
func RunPose(modelPath, runtimeDLLPath string, input []float32) ([]float32, error) {
	if len(input) != poseInputCount {
		return nil, fmt.Errorf("pose input length = %d, want %d", len(input), poseInputCount)
	}

	ort.SetSharedLibraryPath(runtimeDLLPath)
	if err := ort.InitializeEnvironment(); err != nil {
		return nil, fmt.Errorf("initialize ONNX Runtime: %w", err)
	}
	defer func() { _ = ort.DestroyEnvironment() }()

	inputTensor, err := ort.NewTensor(ort.NewShape(1, 3, poseInputSize, poseInputSize), input)
	if err != nil {
		return nil, fmt.Errorf("create pose input tensor: %w", err)
	}
	defer func() { _ = inputTensor.Destroy() }()

	outputTensor, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 56, 8400))
	if err != nil {
		return nil, fmt.Errorf("create pose output tensor: %w", err)
	}
	defer func() { _ = outputTensor.Destroy() }()

	session, err := ort.NewAdvancedSession(
		modelPath,
		[]string{"images"},
		[]string{"output0"},
		[]ort.Value{inputTensor},
		[]ort.Value{outputTensor},
		nil,
	)
	if err != nil {
		return nil, fmt.Errorf("open pose ONNX session: %w", err)
	}
	defer func() { _ = session.Destroy() }()

	if err := session.Run(); err != nil {
		return nil, fmt.Errorf("run pose ONNX session: %w", err)
	}

	output := outputTensor.GetData()
	if len(output) != poseOutputCount {
		return nil, fmt.Errorf("pose output length = %d, want %d", len(output), poseOutputCount)
	}
	return append([]float32(nil), output...), nil
}