package onnx

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"

	ort "github.com/yalue/onnxruntime_go"
	"golang.org/x/sys/unix"
)

var (
	runtimeOnce sync.Once
	runtimeErr  error
)

// FindORTLibrary finds the libonnxruntime shared library: the explicit path,
// then $CHIRAMI_ORT_LIB (which the nix package sets), next to the executable,
// then the usual system locations.
func FindORTLibrary(explicitPath string) (string, error) {
	if explicitPath != "" {
		if _, err := os.Stat(explicitPath); err == nil {
			return explicitPath, nil
		}
		return "", fmt.Errorf("specified onnxruntime library not found: %s", explicitPath)
	}

	if envPath := os.Getenv("CHIRAMI_ORT_LIB"); envPath != "" {
		if _, err := os.Stat(envPath); err == nil {
			return envPath, nil
		}
		return "", fmt.Errorf("CHIRAMI_ORT_LIB points to non-existent path: %s", envPath)
	}

	libName := "libonnxruntime.so"
	if runtime.GOOS == "darwin" {
		libName = "libonnxruntime.dylib"
	}

	// Check next to current executable
	if exe, err := os.Executable(); err == nil {
		p := filepath.Join(filepath.Dir(exe), libName)
		if _, err := os.Stat(p); err == nil {
			return p, nil
		}
	}

	// Standard search locations
	candidates := []string{
		"/opt/homebrew/lib/" + libName,
		"/usr/local/lib/" + libName,
		"/usr/lib/" + libName,
		"/usr/lib64/" + libName,
		libName,
	}
	for _, c := range candidates {
		if _, err := os.Stat(c); err == nil {
			return c, nil
		}
	}

	return "", fmt.Errorf("could not find %s: set CHIRAMI_ORT_LIB", libName)
}

// InitRuntime initializes the ONNX Runtime environment once.
func InitRuntime(explicitLibPath string) error {
	runtimeOnce.Do(func() {
		if ort.IsInitialized() {
			return
		}
		libPath, err := FindORTLibrary(explicitLibPath)
		if err != nil {
			runtimeErr = err
			return
		}
		ort.SetSharedLibraryPath(libPath)
		err = ort.InitializeEnvironment()
		if err != nil {
			runtimeErr = fmt.Errorf("initialize onnxruntime (%s): %w", libPath, err)
			return
		}
	})
	return runtimeErr
}

// newSession opens one model so that its weights stay on disk: the model is
// kept in ORT format, memory-mapped, and ORT is told to use those bytes in
// place instead of copying them (and not to repack weights or keep arenas).
// The weights then count as page cache the OS can drop, not as chirami-ocr's
// own memory, and a model the image never needs is never read in: about
// 120 MB of real memory instead of 700 MB. The ORT-format file is made from
// the .onnx on first use, and again if this onnxruntime cannot load it.
func newSession(path string, inputs, outputs []string) (*ort.DynamicAdvancedSession, error) {
	ortPath := strings.TrimSuffix(path, ".onnx") + ".ort"
	s, err := mmapSession(ortPath, inputs, outputs)
	if err == nil {
		return s, nil
	}
	if err := convert(path, ortPath); err != nil {
		return nil, err
	}
	return mmapSession(ortPath, inputs, outputs)
}

func sessionOptions() (*ort.SessionOptions, error) {
	opts, err := ort.NewSessionOptions()
	if err != nil {
		return nil, err
	}
	for _, err := range []error{
		opts.SetCpuMemArena(false),
		opts.SetMemPattern(false),
		opts.AddSessionConfigEntry("session.disable_prepacking", "1"),
	} {
		if err != nil {
			opts.Destroy()
			return nil, err
		}
	}
	return opts, nil
}

// mmapSession opens an ORT-format model in place. The mapping lives as long
// as the helper process does.
func mmapSession(path string, inputs, outputs []string) (*ort.DynamicAdvancedSession, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	st, err := f.Stat()
	if err != nil {
		return nil, err
	}
	b, err := unix.Mmap(int(f.Fd()), 0, int(st.Size()), unix.PROT_READ, unix.MAP_SHARED)
	if err != nil {
		return nil, err
	}
	opts, err := sessionOptions()
	if err != nil {
		unix.Munmap(b)
		return nil, err
	}
	defer opts.Destroy()
	opts.AddSessionConfigEntry("session.use_ort_model_bytes_directly", "1")
	opts.AddSessionConfigEntry("session.use_ort_model_bytes_for_initializers", "1")
	s, err := ort.NewDynamicAdvancedSessionWithONNXData(b, inputs, outputs, opts)
	if err != nil {
		unix.Munmap(b)
		return nil, err
	}
	return s, nil
}

// convert saves the ONNX model at src as an optimised ORT-format model at
// dst, atomically.
func convert(src, dst string) error {
	opts, err := sessionOptions()
	if err != nil {
		return err
	}
	defer opts.Destroy()
	tmp := dst + ".tmp"
	opts.AddSessionConfigEntry("session.save_model_format", "ORT")
	if err := opts.SetOptimizedModelFilePath(tmp); err != nil {
		return err
	}
	s, err := ort.NewDynamicAdvancedSession(src, nil, nil, opts)
	if err != nil {
		return fmt.Errorf("convert %s: %w", filepath.Base(src), err)
	}
	s.Destroy()
	return os.Rename(tmp, dst)
}