package onnx

/*
#cgo CXXFLAGS: -std=c++17
#cgo linux LDFLAGS: -lonnxruntime -ldl
#cgo darwin LDFLAGS: -lonnxruntime -framework CoreFoundation

#include <stdlib.h>
#include <string.h>

// ONNX Runtime C API 声明
#ifdef __cplusplus
extern "C" {
#endif

typedef void* OrtSession;
typedef void* OrtMemoryInfo;
typedef void* OrtValue;

// 初始化 ONNX 会话
OrtSession onnx_create_session(const char* model_path);
void onnx_destroy_session(OrtSession session);

// 运行推理
int onnx_run(OrtSession session, const float* input_data, int input_size, float* output_data, int output_size);

// 错误信息
const char* onnx_get_last_error();

#ifdef __cplusplus
}
#endif
*/
import "C"
import (
	"errors"
	"fmt"
	"sync"
	"unsafe"
)

type Session struct {
	session C.OrtSession
	mu      sync.Mutex
}

var (
	sessions = make(map[string]*Session)
	mu       sync.RWMutex
)

// LoadModel 加载 ONNX 模型
func LoadModel(name, path string) error {
	cPath := C.CString(path)
	defer C.free(unsafe.Pointer(cPath))

	session := C.onnx_create_session(cPath)
	if session == nil {
		return fmt.Errorf("加载模型失败: %s", C.GoString(C.onnx_get_last_error()))
	}

	mu.Lock()
	sessions[name] = &Session{session: session}
	mu.Unlock()

	return nil
}

// Run 执行模型推理
func (s *Session) Run(input []float32, inputSize int) ([]float32, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	output := make([]float32, inputSize)
	
	result := C.onnx_run(
		s.session,
		(*C.float)(unsafe.Pointer(&input[0])),
		C.int(len(input)),
		(*C.float)(unsafe.Pointer(&output[0])),
		C.int(len(output)),
	)

	if result != 0 {
		return nil, errors.New(C.GoString(C.onnx_get_last_error()))
	}

	return output, nil
}

// Close 关闭会话
func (s *Session) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	
	if s.session != nil {
		C.onnx_destroy_session(s.session)
		s.session = nil
	}
}

// GetSession 获取已加载的会话
func GetSession(name string) (*Session, bool) {
	mu.RLock()
	s, ok := sessions[name]
	mu.RUnlock()
	return s, ok
}

// CloseAll 关闭所有会话
func CloseAll() {
	mu.Lock()
	for _, s := range sessions {
		s.Close()
	}
	sessions = make(map[string]*Session)
	mu.Unlock()
}