package eval

import (
	"fmt"
	"os"
	"sort"
	"strconv"
	"sync"
	"testing"

	"github.com/kami/maven/internal/router"
)

const requiredONNXGatesEnv = "MAVEN_ONNX_REQUIRED_GATES"

var executedONNXGates struct {
	sync.Mutex
	names []string
}

// recordONNXGateExecuted is called only after a gate has loaded every model it
// needs. TestMain uses the names as proof that an aggregate command did not
// turn a second initialization error into a green SKIP.
func recordONNXGateExecuted(t *testing.T) {
	t.Helper()
	executedONNXGates.Lock()
	executedONNXGates.names = append(executedONNXGates.names, t.Name())
	executedONNXGates.Unlock()
}

func TestMain(m *testing.M) {
	required, err := requiredONNXGateCount()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}

	lease, err := configuredONNXTestRuntime(required > 0)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}

	code := m.Run()
	if required > 0 {
		executedONNXGates.Lock()
		names := append([]string(nil), executedONNXGates.names...)
		executedONNXGates.Unlock()
		sort.Strings(names)
		fmt.Fprintf(os.Stderr, "ONNX aggregate proof: verified %d required model gates: %v\n", len(names), names)
		if len(names) < required {
			fmt.Fprintf(os.Stderr, "ONNX aggregate gate failed: executed %d, require at least %d\n", len(names), required)
			code = 1
		}
	}

	if lease != nil {
		if err := lease.Close(); err != nil {
			fmt.Fprintf(os.Stderr, "ONNX test runtime cleanup: %v\n", err)
			code = 1
		}
	}
	os.Exit(code)
}

func requiredONNXGateCount() (int, error) {
	raw := os.Getenv(requiredONNXGatesEnv)
	if raw == "" {
		return 0, nil
	}
	n, err := strconv.Atoi(raw)
	if err != nil || n < 0 {
		return 0, fmt.Errorf("%s must be a non-negative integer, got %q", requiredONNXGatesEnv, raw)
	}
	return n, nil
}

func configuredONNXTestRuntime(required bool) (*router.ONNXRuntimeLease, error) {
	lib := os.Getenv("MAVEN_ONNX_LIB")
	if lib == "" {
		if required {
			return nil, fmt.Errorf("ONNX aggregate gate requires MAVEN_ONNX_LIB")
		}
		return nil, nil
	}
	if _, err := os.Stat(lib); err != nil {
		if required {
			return nil, fmt.Errorf("ONNX aggregate gate runtime %s: %w", lib, err)
		}
		// Preserve the ordinary portable test path: individual model tests
		// report a skip when optional local model dependencies are absent.
		return nil, nil
	}
	lease, err := router.AcquireONNXRuntime(lib)
	if err != nil {
		return nil, fmt.Errorf("initialize shared ONNX test runtime: %w", err)
	}
	return lease, nil
}