package eval

import (
	"context"
	"os"
	"path/filepath"
	"testing"

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

// TestONNXRoutingHeads — the cascade with the routing heads wired, which is
// what V-664 deploys. Opt-in via MAVEN_ONNX_LIB, same as TestONNXBaseline, and
// one TestONNX* per process.
//
// The comparison worth reading is against TestONNXBaseline, which is the same
// cascade with the same grammars and the same classifier floor and no heads.
// Only the middle arm varies.
//
// It also checks the Go unigram tokenizer against the Python one, because the
// heads were trained through transformers and are read through a hand-written
// tokenizer. A mismatch shows up here as a score below what Python measured on
// the same weights, and nowhere else.
func TestONNXRoutingHeads(t *testing.T) {
	lib := os.Getenv("MAVEN_ONNX_LIB")
	if lib == "" {
		t.Skip("MAVEN_ONNX_LIB unset — see AGENTS.md § Embedder model for intent routing")
	}
	// Absolute, because onnxruntime resolves a graph's external weights file
	// against the model path it was given, and a relative one lands in the
	// test's working directory.
	root, err := filepath.Abs("../../..")
	if err != nil {
		t.Fatal(err)
	}
	model := filepath.Join(root, "models/embedder/multilingual-e5-small/model_quantized.onnx")
	tok := filepath.Join(root, "models/embedder/multilingual-e5-small/tokenizer.json")
	heads := filepath.Join(root, "models/embedder/router-heads/router_heads.onnx")
	for _, p := range []string{lib, model, tok, heads} {
		if _, err := os.Stat(p); err != nil {
			t.Skipf("missing %s: %v", p, err)
		}
	}
	emb, err2 := router.NewONNXEmbedder(model, tok, lib)
	if err2 != nil {
		t.Skipf("onnx embedder unavailable: %v", err2)
	}
	err = nil
	defer emb.Close()

	h, err := router.NewRouterHeads(heads, tok)
	if err != nil {
		t.Skipf("routing heads unavailable: %v", err)
	}
	defer h.Close()
	recordONNXGateExecuted(t)

	f, err := Load()
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	rep, err := Score(context.Background(), "heads+classifier", withHeads(t, emb, h), f)
	if err != nil {
		t.Fatalf("Score: %v", err)
	}
	t.Log("\n" + rep.String() + rep.Failures())
}

// withHeads mirrors newBaselineRouter and adds the one arm under test. It is a
// separate function rather than a parameter so the baseline's signature stays
// the shape every other test calls it with.
func withHeads(t *testing.T, emb router.Embedder, h *router.RouterHeads) *router.Router {
	t.Helper()
	acts := router.DefaultActMatcher{Fns: actFns}
	return router.New(router.Config{
		Grammars:   baselineGrammars(acts),
		Classifier: newBaselineClassifier(t, emb),
		Extractor: router.Extractor{
			Time:  router.StubDateTimeParser{},
			Acts:  acts,
			Facts: router.DefaultFactParser{},
		},
		Threshold: config.DefaultRouterThreshold,
		Heads:     h,
	})
}