package autogen

import (
	"encoding/binary"
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

// writeSafetensors lays down a real .safetensors container carrying only the
// header (8-byte LE length + JSON table). Classification never reads past it, so
// a synthetic file with no payload exercises the same path as a 10 GB encoder.
func writeSafetensors(t *testing.T, path string, shapes map[string][]int64) {
	t.Helper()
	tbl := map[string]stTensor{}
	for k, v := range shapes {
		tbl[k] = stTensor{DType: "F16", Shape: v}
	}
	j, err := json.Marshal(tbl)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		t.Fatal(err)
	}
	f, err := os.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	if err := binary.Write(f, binary.LittleEndian, uint64(len(j))); err != nil {
		t.Fatal(err)
	}
	if _, err := f.Write(j); err != nil {
		t.Fatal(err)
	}
}

func TestAutogen_classifySafetensors(t *testing.T) {
	cases := []struct {
		name       string
		shapes     map[string][]int64
		wantRole   ComponentRole
		wantFamily string
		wantWidth  int64
		wantVision bool
	}{
		{
			name:       "sd vae",
			shapes:     map[string][]int64{"decoder.conv_in.weight": {512, 4, 3, 3}},
			wantRole:   RoleVae,
			wantFamily: VaeFamilySD,
		},
		{
			name:       "flux ae",
			shapes:     map[string][]int64{"decoder.conv_in.weight": {512, 16, 3, 3}},
			wantRole:   RoleVae,
			wantFamily: VaeFamilyFlux,
		},
		{
			name:       "flux2 ae",
			shapes:     map[string][]int64{"decoder.conv_in.weight": {512, 32, 3, 3}},
			wantRole:   RoleVae,
			wantFamily: VaeFamilyFlux2,
		},
		{
			name:       "wan 3d causal vae",
			shapes:     map[string][]int64{"conv1.weight": {96, 3, 3, 3, 3}},
			wantRole:   RoleVae,
			wantFamily: VaeFamilyWan3D,
		},
		{
			// Wan 2.1 and the Qwen-Image 20B VAE: 16-channel latent, and the
			// decoder projection is what states it.
			name: "wan 3d causal vae, 16-channel latent",
			shapes: map[string][]int64{
				"conv1.weight":         {32, 32, 1, 1, 1},
				"decoder.conv1.weight": {384, 16, 3, 3, 3},
			},
			wantRole:   RoleVae,
			wantFamily: VaeFamilyWan3D,
			wantWidth:  16,
		},
		{
			// Qwen-Image 2.1's RGBA autoencoder: same family, 64-channel latent.
			name: "qwen-image 2.1 rgba vae",
			shapes: map[string][]int64{
				"conv1.weight":         {128, 128, 1, 1, 1},
				"decoder.conv1.weight": {1152, 64, 1, 3, 3},
				"encoder.conv1.weight": {96, 4, 1, 3, 3},
			},
			wantRole:   RoleVae,
			wantFamily: VaeFamilyWan3D,
			wantWidth:  64,
		},
		{
			name:      "clip-l",
			shapes:    map[string][]int64{"text_model.embeddings.token_embedding.weight": {49408, 768}},
			wantRole:  RoleClip,
			wantWidth: 768,
		},
		{
			name:      "clip-g",
			shapes:    map[string][]int64{"text_model.embeddings.token_embedding.weight": {49408, 1280}},
			wantRole:  RoleClip,
			wantWidth: 1280,
		},
		{
			name:      "t5xxl",
			shapes:    map[string][]int64{"encoder.block.0.layer.0.SelfAttention.q.weight": {4096, 4096}},
			wantRole:  RoleT5,
			wantWidth: 4096,
		},
		{
			name:      "plain decoder llm",
			shapes:    map[string][]int64{"model.embed_tokens.weight": {131072, 3072}},
			wantRole:  RoleLlm,
			wantWidth: 3072,
		},
		{
			name: "vision llm",
			shapes: map[string][]int64{
				"model.embed_tokens.weight":      {152064, 3584},
				"visual.patch_embed.proj.weight": {1280, 3, 2, 14, 14},
			},
			wantRole:   RoleLlm,
			wantWidth:  3584,
			wantVision: true,
		},
		{
			// A DiT is not a component: it must not land in the pool, or the
			// diffusion model would be offered as its own text encoder.
			name:     "dit is not a component",
			shapes:   map[string][]int64{"double_blocks.0.img_attn.qkv.weight": {9216, 3072}},
			wantRole: RoleNone,
		},
		{
			// An unknown latent width is a VAE we cannot vouch for; wiring it
			// would produce a broken decode, so it is dropped entirely.
			name:     "unknown latent width dropped",
			shapes:   map[string][]int64{"decoder.conv_in.weight": {512, 7, 3, 3}},
			wantRole: RoleNone,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			tbl := map[string]stTensor{}
			for k, v := range tc.shapes {
				tbl[k] = stTensor{Shape: v}
			}
			got := classifySafetensors(tbl, 1.0, "x.safetensors")
			if got.Role != tc.wantRole {
				t.Fatalf("role = %q, want %q", got.Role, tc.wantRole)
			}
			if got.Family != tc.wantFamily {
				t.Errorf("family = %q, want %q", got.Family, tc.wantFamily)
			}
			if got.Width != tc.wantWidth {
				t.Errorf("width = %d, want %d", got.Width, tc.wantWidth)
			}
			if got.Vision != tc.wantVision {
				t.Errorf("vision = %v, want %v", got.Vision, tc.wantVision)
			}
		})
	}
}

func TestAutogen_ScanEncoderPool(t *testing.T) {
	root := t.TempDir()
	writeSafetensors(t, filepath.Join(root, "vae", "ae.safetensors"),
		map[string][]int64{"decoder.conv_in.weight": {512, 16, 3, 3}})
	writeSafetensors(t, filepath.Join(root, "vae", "sdxl_vae.safetensors"),
		map[string][]int64{"decoder.conv_in.weight": {512, 4, 3, 3}})
	writeSafetensors(t, filepath.Join(root, "clip", "clip_l.safetensors"),
		map[string][]int64{"text_model.embeddings.token_embedding.weight": {49408, 768}})
	writeSafetensors(t, filepath.Join(root, "clip", "clip_g.safetensors"),
		map[string][]int64{"text_model.embeddings.token_embedding.weight": {49408, 1280}})
	writeSafetensors(t, filepath.Join(root, "junk", "readme.safetensors"),
		map[string][]int64{"nothing.weight": {1, 1}})
	// Not a safetensors container at all: must be skipped, not fatal.
	if err := os.WriteFile(filepath.Join(root, "junk", "half.safetensors"), []byte("no"), 0o644); err != nil {
		t.Fatal(err)
	}

	p := ScanEncoderPool([]string{root})
	if len(p.Files) != 4 {
		t.Fatalf("files = %d, want 4: %+v", len(p.Files), p.Files)
	}
	if got := p.Vae(VaeFamilyFlux); !strings.HasSuffix(got, "ae.safetensors") {
		t.Errorf("flux vae = %q", got)
	}
	if got := p.Vae(VaeFamilySD, "sdxl"); !strings.Contains(got, "sdxl") {
		t.Errorf("sdxl vae = %q", got)
	}
	if got := p.Clip(768); !strings.Contains(got, "clip_l") {
		t.Errorf("clip-l = %q", got)
	}
	if got := p.Clip(1280); !strings.Contains(got, "clip_g") {
		t.Errorf("clip-g = %q", got)
	}
	if got := p.T5(); got != "" {
		t.Errorf("t5 = %q, want none", got)
	}
	// A hint that matches nothing still yields a stable pick, never a blank.
	if got := p.Vae(VaeFamilyFlux, "no-such-thing"); got == "" {
		t.Error("hint miss should fall back to the sorted-first candidate")
	}
}

func TestAutogen_EncoderPoolLlm(t *testing.T) {
	p := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/qwen25vl-q4.gguf", Role: RoleLlm, Width: 3584, SizeGB: 4, Mmproj: "/m/mmproj.gguf", Vision: true},
		{Path: "/m/qwen25vl-q8.gguf", Role: RoleLlm, Width: 3584, SizeGB: 8, Mmproj: "/m/mmproj.gguf", Vision: true},
		{Path: "/m/qwen25-7b.gguf", Role: RoleLlm, Width: 3584, SizeGB: 9},
		{Path: "/m/qwen3-4b.gguf", Role: RoleLlm, Width: 2560, SizeGB: 3},
		{Path: "/m/t5.gguf", Role: RoleT5, Width: 4096, SizeGB: 5},
	}}
	// Widest file wins at equal width...
	if got, _ := p.Llm(3584, false, ""); got != "/m/qwen25-7b.gguf" {
		t.Errorf("llm(3584) = %q, want the largest", got)
	}
	// ...unless a vision tower is required, which excludes the bigger text-only
	// file and drags the paired projector along.
	got, proj := p.Llm(3584, true, "")
	if got != "/m/qwen25vl-q8.gguf" || proj != "/m/mmproj.gguf" {
		t.Errorf("llm(3584, vision) = %q/%q", got, proj)
	}
	if got, _ := p.Llm(2560, false, ""); got != "/m/qwen3-4b.gguf" {
		t.Errorf("llm(2560) = %q", got)
	}
	// An unmatched width picks nothing rather than the closest: a mismatched
	// encoder does not degrade, it fails to load.
	if got, _ := p.Llm(4096, false, ""); got != "" {
		t.Errorf("llm(4096) = %q, want none (t5 is not an llm)", got)
	}
	if got, _ := p.Llm(0, false, ""); got != "" {
		t.Errorf("llm(0) = %q, want none", got)
	}
	var nilPool *EncoderPool
	if got, _ := nilPool.Llm(3584, true, ""); got != "" {
		t.Error("nil pool must be inert")
	}
	if nilPool.Vae(VaeFamilyFlux) != "" || nilPool.Clip(768) != "" || nilPool.T5() != "" {
		t.Error("nil pool must be inert for every getter")
	}
}

// The real-world tie this pool has to break: Qwen3-4B-Instruct-2507 and
// Qwen3-VL-4B-Instruct are both 2560 wide, and the Q8 quants of each round to
// the same size, so the scan's own tiebreak (path) lands on the VL file. Only
// the declared pin says which of them a given DiT was trained against.
func TestAutogen_EncoderPoolLlmPrefer(t *testing.T) {
	const (
		vl   = "/m/QwenVL/Qwen3VL-4B-Instruct-Q8_0.gguf"
		q3q8 = "/m/lmstudio-community/Qwen3-4B-Instruct-2507-GGUF/Qwen3-4B-Instruct-2507-Q8_0.gguf"
		q3q4 = "/m/lmstudio-community/Qwen3-4B-Instruct-2507-GGUF/Qwen3-4B-Instruct-2507-Q4_K_M.gguf"
	)
	p := &EncoderPool{Files: []ComponentFile{
		{Path: vl, Role: RoleLlm, Width: 2560, SizeGB: 3.99, Vision: true,
			Mmproj: "/m/QwenVL/mmproj-Qwen3VL-4B-Instruct-F16.gguf"},
		{Path: q3q8, Role: RoleLlm, Width: 2560, SizeGB: 3.99},
		{Path: q3q4, Role: RoleLlm, Width: 2560, SizeGB: 2.32},
	}}
	// No declaration: the path tiebreak takes the VL file, which is how
	// Z-Image ended up conditioned on Qwen3-VL.
	if got, _ := p.Llm(2560, false, ""); got != vl {
		t.Errorf("unpinned llm = %q, want %q", got, vl)
	}
	// A declared file of the matching width wins outright, even when the scan
	// finds a larger candidate.
	if got, _ := p.Llm(2560, false, q3q4); got != q3q4 {
		t.Errorf("pinned llm = %q, want %q", got, q3q4)
	}
	// A pin without a vision tower cannot serve a vision model (the caller then
	// retries text-only and gets it), and an unseen pin changes nothing.
	if got, _ := p.Llm(2560, true, q3q4); got != vl {
		t.Errorf("vision call took a text-only pin: %q", got)
	}
	if got, proj := p.Llm(2560, true, vl); got != vl || proj == "" {
		t.Errorf("vision pin = %q/%q", got, proj)
	}
	if got, _ := p.Llm(2560, false, "/m/elsewhere/qwen3-4b.gguf"); got != vl {
		t.Errorf("unseen pin = %q, want %q", got, vl)
	}
}

// onDisk moves the test into a fresh directory holding an empty file for each
// relative name, so a declared component path like "ae.safetensors" exists:
// fillEncoderSet clears declared files that are gone, and a fixture naming a
// file that is not there would read as a model missing that part.
func onDisk(t *testing.T, names ...string) {
	t.Helper()
	t.Chdir(t.TempDir())
	for _, n := range names {
		if err := os.WriteFile(n, nil, 0o644); err != nil {
			t.Fatal(err)
		}
	}
}

func TestAutogen_fillEncoderSet(t *testing.T) {
	p := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/ae.safetensors", Role: RoleVae, Family: VaeFamilyFlux},
		{Path: "/m/flux2-ae.safetensors", Role: RoleVae, Family: VaeFamilyFlux2},
		{Path: "/m/clip_l.safetensors", Role: RoleClip, Width: 768},
		{Path: "/m/t5xxl.gguf", Role: RoleT5, Width: 4096},
	}}
	// A declared file outside the scan still wins, as long as it exists.
	hand := filepath.Join(t.TempDir(), "clip_l.safetensors")
	if err := os.WriteFile(hand, []byte("x"), 0o644); err != nil {
		t.Fatal(err)
	}
	got := fillEncoderSet(EncoderSet{ClipL: hand}, p)
	if got.ClipL != hand {
		t.Errorf("declared ClipL was overwritten: %q", got.ClipL)
	}
	// A declared file that is gone is cleared, so the pool fills the slot
	// (the flux2 VAE moved out of the models folder must not read as present).
	gone := fillEncoderSet(EncoderSet{Flux2Vae: filepath.Join(t.TempDir(), "moved", "ae.safetensors")}, p)
	if gone.Flux2Vae != "/m/flux2-ae.safetensors" {
		t.Errorf("missing declared flux2Vae = %q, want the pool's", gone.Flux2Vae)
	}
	// ...and a declared path the pool lists is kept without touching the disk.
	if kept := fillEncoderSet(EncoderSet{T5: "/m/t5xxl.gguf"}, p); kept.T5 != "/m/t5xxl.gguf" {
		t.Errorf("pool-listed declared t5 = %q", kept.T5)
	}
	if got.FluxVae != "/m/ae.safetensors" || got.Flux2Vae != "/m/flux2-ae.safetensors" {
		t.Errorf("vae fill = %q / %q", got.FluxVae, got.Flux2Vae)
	}
	// Z-Image ships an AE byte-identical to flux.1's, so it falls back to it.
	if got.ZimageVae != "/m/ae.safetensors" {
		t.Errorf("zimage vae = %q", got.ZimageVae)
	}
	if got.T5 != "/m/t5xxl.gguf" {
		t.Errorf("t5 = %q", got.T5)
	}
	if got.SdxlVae != "" {
		t.Errorf("sdxl vae = %q, want blank (none on disk)", got.SdxlVae)
	}
	// A nil pool must leave a hand-written set exactly as it was.
	in := EncoderSet{FluxVae: "a", ClipL: "b"}
	if fillEncoderSet(in, nil) != in {
		t.Error("nil pool must not alter the declared set")
	}
}

func TestAutogen_wantsVisionEncoder(t *testing.T) {
	cases := []struct {
		name string
		ov   *Override
		want bool
	}{
		{name: "LongCat-Image-Edit-Turbo-Q8_0", want: true},
		{name: "Qwen-Rapid-NSFW", want: true},
		{name: "flux1-kontext-dev", want: true},
		{name: "sd15-inpaint", want: true},
		{name: "LongCat-Image-Q8_0", want: false},
		{name: "Z-Image-Turbo", want: false},
		// "edit" only counts as its own token: "credit" is not an edit model.
		{name: "credit-model", want: false},
		{name: "LongCat-Image-Edit", ov: &Override{LlmVision: "off"}, want: false},
		{name: "LongCat-Image", ov: &Override{LlmVision: "on"}, want: true},
		// Turning on reference editing asks for the projector too: without it
		// sd.cpp refuses the prompt on 2.1 rather than degrading.
		{name: "some-unified-editor", ov: &Override{RefEdit: "on"}, want: true},
		// ...but refEdit off does not mean no vision: an inpaint model wants the
		// masked img2img route AND the projector.
		{name: "sd15-inpaint", ov: &Override{RefEdit: "off"}, want: true},
		// An explicit vision off still wins over the refEdit implication.
		{name: "some-unified-editor", ov: &Override{RefEdit: "on", LlmVision: "off"}, want: false},
		// Qwen-Image 2.1: one checkpoint does generation AND reference editing,
		// so nothing in the name says "edit" and the version is the only tell.
		{name: "qwen_image_2.1-q8_0", want: true},
		{name: "Qwen-Image-2.1-Q4_K", want: true},
		// The 20B line is a separate text-only checkpoint; its edit sibling is
		// named for it, so 2.1 must not widen the match to every qwen image.
		{name: "qwen_image-q8_0", want: false},
		{name: "LongCat-Image", ov: &Override{LlmVision: "on"}, want: true},
	}
	for _, tc := range cases {
		if got := wantsVisionEncoder("flux", tc.name, tc.ov); got != tc.want {
			t.Errorf("wantsVisionEncoder(%q) = %v, want %v", tc.name, got, tc.want)
		}
	}
}

func TestAutogen_condHiddenFrom(t *testing.T) {
	if got := condHiddenFrom(nil); got != 0 {
		t.Errorf("empty = %d, want 0", got)
	}
	if got := condHiddenFrom(map[string]int64{"txt_in.weight": 3584}); got != 3584 {
		t.Errorf("txt_in = %d", got)
	}
	// Highest-ranked present tensor wins when a model carries several.
	dims := map[string]int64{
		"context_embedder.weight": 4096,
		"txt_in.weight":           3584,
	}
	if got := condHiddenFrom(dims); got != 3584 {
		t.Errorf("ranked pick = %d, want 3584", got)
	}
	if got := condHiddenFrom(map[string]int64{"unrelated.weight": 99}); got != 0 {
		t.Errorf("unrelated = %d, want 0", got)
	}
	// Qwen-Image 2.1 projects the caption through a two-layer MLP, so the width
	// is on txt_in.in_layer and the bare txt_in the 20B line uses is absent.
	// Missing this entry leaves the model with no stated encoder width at all.
	if got := condHiddenFrom(map[string]int64{"txt_in.in_layer.weight": 4096}); got != 4096 {
		t.Errorf("txt_in.in_layer = %d, want 4096", got)
	}
}

func TestAutogen_wan3dLatent(t *testing.T) {
	cases := []struct {
		arch       string
		condHidden int64
		want       int64
	}{
		{"qwen_image", 4096, 64}, // Qwen-Image 2.1 + RGBA VAE
		{"qwen_image", 3584, 16}, // Qwen-Image 20B / Qwen-Image-Edit
		{"qwen_image", 2560, 16}, // Krea2 reports arch qwen_image too
		{"wan", 0, 0},            // Wan video: no opinion, 2.2 TI2V is 48-channel
		{"", 4096, 0},
	}
	for _, c := range cases {
		if got := wan3dLatent(c.arch, c.condHidden); got != c.want {
			t.Errorf("wan3dLatent(%q, %d) = %d, want %d", c.arch, c.condHidden, got, c.want)
		}
	}
}

func TestAutogen_EncoderPoolVaeOfWidth(t *testing.T) {
	// Both files are family wan3d and both match the "qwen_image" hint, so
	// before latent width was recorded the pick was decided by sorted path,
	// which puts the 2.1 VAE ahead of the one the 20B line needs.
	pool := &EncoderPool{Files: []ComponentFile{
		{Role: RoleVae, Family: VaeFamilyWan3D, Width: 64, Path: "m/qwen_image_2.1_vae_bf16.safetensors"},
		{Role: RoleVae, Family: VaeFamilyWan3D, Width: 16, Path: "m/qwen_image_vae.safetensors"},
	}}
	if got := pool.VaeOfWidth(VaeFamilyWan3D, 16, "qwen_image"); !strings.HasSuffix(got, "qwen_image_vae.safetensors") {
		t.Errorf("want 16 = %q", got)
	}
	if got := pool.VaeOfWidth(VaeFamilyWan3D, 64, "qwen_image"); !strings.HasSuffix(got, "qwen_image_2.1_vae_bf16.safetensors") {
		t.Errorf("want 64 = %q", got)
	}
	// want 0 keeps the old hint-then-sort behaviour.
	if got := pool.Vae(VaeFamilyWan3D, "qwen_image"); !strings.HasSuffix(got, "qwen_image_2.1_vae_bf16.safetensors") {
		t.Errorf("no opinion = %q", got)
	}
	// A width nothing carries falls back to the full list rather than to "",
	// so a file that predates Width being recorded never drops the --vae flag.
	if got := pool.VaeOfWidth(VaeFamilyWan3D, 48, "qwen_image"); got == "" {
		t.Error("unmatched width returned empty, want fallback")
	}
}

func TestAutogen_IsReferenceEditModel(t *testing.T) {
	cases := []struct {
		name string
		ov   *Override
		want bool
	}{
		{name: "LongCat-Image-Edit-Turbo-Q8_0", want: true},
		{name: "Qwen-Rapid-NSFW", want: true},
		{name: "flux1-kontext-dev", want: true},
		{name: "LongCat-Image-Q8_0", want: false},
		{name: "Z-Image-Turbo", want: false},
		{name: "credit-model", want: false},
		// Narrower than wantsVisionEncoder on purpose: an inpaint model is an
		// edit, but it wants the MASKED img2img route. Sending its source as a
		// reference would redraw the whole frame and ignore the mask.
		{name: "sd15-inpaint", want: false},
		{name: "flux1-fill-dev", want: false},
		{name: "LongCat-Image-Edit", ov: &Override{LlmVision: "off"}, want: false},
		{name: "LongCat-Image", ov: &Override{LlmVision: "on"}, want: true},
		// Unified checkpoint: edit-capable with no edit token in the name.
		{name: "qwen_image_2.1-q8_0", want: true},
		{name: "Qwen-Image-2.1-Q4_K", want: true},
		{name: "qwen_image-q8_0", want: false},
		// An explicit off still wins over the version match.
		{name: "qwen_image_2.1-q8_0", ov: &Override{RefEdit: "off"}, want: false},
		// FLUX.2 / FLUX.3: one checkpoint generates and edits, so no edit token.
		{name: "flux-3-image-klein-UD-Q8_K_XL", want: true},
		{name: "flux-2-klein-4b-Q8_0", want: true},
		{name: "flux2-dev-Q4_K_M", want: true},
		{name: "FLUX.2-klein-9B", want: true},
		{name: "flux-3-image-klein-Q8_0", ov: &Override{RefEdit: "off"}, want: false},
		// flux.1 base is text-only, and "flux-1" must not read as "flux".
		{name: "flux1-dev-Q8_0", want: false},
		{name: "flux-1-schnell", want: false},
		{name: "fluxed-23-merge", want: false},
	}
	for _, tc := range cases {
		if got := IsReferenceEditModel(tc.name, tc.ov); got != tc.want {
			t.Errorf("IsReferenceEditModel(%q) = %v, want %v", tc.name, got, tc.want)
		}
	}
}

// Flux.2 states a caption width that is a MULTIPLE of its encoder's hidden size
// (klein 9B: txt_in 12288 against a 4096-wide Qwen3-8B), so the pick has to
// divide rather than compare. Exactness still outranks concatenation, and the
// factor cap keeps the rule from matching arbitrarily narrow encoders.
func TestAutogen_EncoderPoolLlmConcat(t *testing.T) {
	p := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/qwen3-8b.gguf", Role: RoleLlm, Family: "qwen3", Width: 4096, SizeGB: 8},
		{Path: "/m/qwen3-4b.gguf", Role: RoleLlm, Family: "qwen3", Width: 2560, SizeGB: 3},
		{Path: "/m/tiny-1024.gguf", Role: RoleLlm, Family: "qwen3", Width: 1024, SizeGB: 1},
	}}
	// klein 9B: 12288 = 3 x 4096.
	if got, _ := p.Llm(12288, false, ""); got != "/m/qwen3-8b.gguf" {
		t.Errorf("llm(12288) = %q, want the 4096-wide encoder", got)
	}
	// The global pin is 2560 wide: 12288 is not a multiple of it, so it must not
	// be dragged in. This is the exact shape of the klein-9B mis-wire.
	if got, _ := p.Llm(12288, false, "/m/qwen3-4b.gguf"); got != "/m/qwen3-8b.gguf" {
		t.Errorf("llm(12288, pin=4b) = %q, want the pin ignored", got)
	}
	// 12288/1024 = 12, not a known caption layout: a narrow encoder is not a candidate
	// just because the arithmetic works.
	narrow := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/tiny-1024.gguf", Role: RoleLlm, Family: "qwen3", Width: 1024, SizeGB: 1},
	}}
	if got, _ := narrow.Llm(12288, false, ""); got != "" {
		t.Errorf("llm(12288) over a 1024-wide pool = %q, want none", got)
	}
	// An exact match beats a concatenated one even when the concat candidate is
	// bigger and would win the size tiebreak within a tier.
	exact := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/wide-12288.gguf", Role: RoleLlm, Family: "qwen3", Width: 12288, SizeGB: 2},
		{Path: "/m/qwen3-8b.gguf", Role: RoleLlm, Family: "qwen3", Width: 4096, SizeGB: 8},
	}}
	if got, _ := exact.Llm(12288, false, ""); got != "/m/wide-12288.gguf" {
		t.Errorf("llm(12288) = %q, want the exact-width file", got)
	}
	// A pin only reachable by concat does not outrank an exact-width scan hit.
	if got, _ := exact.Llm(12288, false, "/m/qwen3-8b.gguf"); got != "/m/wide-12288.gguf" {
		t.Errorf("llm(12288, pin=concat) = %q, want the exact-width file", got)
	}
	// ...but a pin at the same tier still wins outright.
	if got, _ := p.Llm(12288, false, "/m/qwen3-8b.gguf"); got != "/m/qwen3-8b.gguf" {
		t.Errorf("llm(12288, pin=8b) = %q, want the pin", got)
	}
}

func TestAutogen_captionFactor(t *testing.T) {
	cases := []struct {
		hidden, width, want int64
	}{
		{2560, 2560, 1},  // the ordinary case: one encoder hidden state
		{12288, 4096, 3}, // flux.2 klein 9B
		{7680, 2560, 3},  // flux.2 klein 4B
		{15360, 5120, 3}, // flux.2 dev (Mistral-Small-3)
		{12288, 2560, 0}, // not divisible
		{12288, 1024, 0}, // divisible but past the concat cap
		{4096, 0, 0},     // unmeasured encoder
		{0, 4096, 0},     // unmeasured DiT
		{-1, 4096, 0},    // defensive
	}
	for _, c := range cases {
		if got := captionFactor(c.hidden, c.width); got != c.want {
			t.Errorf("captionFactor(%d, %d) = %d, want %d", c.hidden, c.width, got, c.want)
		}
	}
}

func TestAutogen_knowsLlm(t *testing.T) {
	p := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/qwen3-4b.gguf", Role: RoleLlm, Width: 2560},
		{Path: "/m/t5.gguf", Role: RoleT5, Width: 4096},
	}}
	// Classified as an LLM: known at ANY width, which is what separates "the scan
	// rejected this pin" from "the scan never saw it".
	if !p.knowsLlm("/m/qwen3-4b.gguf") {
		t.Error("a scanned llm should be known")
	}
	if p.knowsLlm("/m/t5.gguf") {
		t.Error("a t5 is not an llm")
	}
	if p.knowsLlm("/m/never-scanned.gguf") || p.knowsLlm("") {
		t.Error("an unscanned or empty path is not known")
	}
	var nilPool *EncoderPool
	if nilPool.knowsLlm("/m/qwen3-4b.gguf") {
		t.Error("nil pool knows nothing")
	}
}

// writeGgufKVs lays down a gguf carrying only its header and the given KVs (no
// tensors). String values are written as gguf strings, int64 as uint32, which is
// all the pool's classification ladder reads. Keys are written in order, so
// general.architecture must come first for the arch-prefixed keys to match.
func writeGgufKVs(t *testing.T, path string, kvs [][2]any) {
	t.Helper()
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		t.Fatal(err)
	}
	f, err := os.Create(path)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	le := binary.LittleEndian
	str := func(s string) {
		binary.Write(f, le, uint64(len(s)))
		f.WriteString(s)
	}
	f.WriteString("GGUF")
	binary.Write(f, le, uint32(3))
	binary.Write(f, le, uint64(0))
	binary.Write(f, le, uint64(len(kvs)))
	for _, kv := range kvs {
		str(kv[0].(string))
		switch v := kv[1].(type) {
		case string:
			binary.Write(f, le, uint32(8))
			str(v)
		case int:
			binary.Write(f, le, uint32(4))
			binary.Write(f, le, uint32(v))
		case bool:
			binary.Write(f, le, uint32(7))
			var b uint8
			if v {
				b = 1
			}
			binary.Write(f, le, b)
		case []string:
			binary.Write(f, le, uint32(9))
			binary.Write(f, le, uint32(8))
			binary.Write(f, le, uint64(len(v)))
			for _, e := range v {
				str(e)
			}
		default:
			t.Fatalf("unsupported kv value %T", v)
		}
	}
}

// A recurrent LLM can never be a diffusion text encoder, so it must not enter
// the pool even when it would win its width class. The hybrid paths sort first,
// so with equal sizes and arch rank the old tiebreak handed Qwen-Image 2.1 the
// hybrid (the MiMo-V2.6-Distill-Qwen-9B incident).
func TestAutogen_EncoderPoolSkipsRecurrentLlm(t *testing.T) {
	root := t.TempDir()
	writeGgufKVs(t, filepath.Join(root, "a-hybrid", "MiMo-Distill-Qwen-9B-Q8_0.gguf"), [][2]any{
		{"general.architecture", "qwen35"},
		{"qwen35.embedding_length", 4096},
		{"qwen35.full_attention_interval", 4},
	})
	writeGgufKVs(t, filepath.Join(root, "b-ssm", "Granite-Hybrid-Q8_0.gguf"), [][2]any{
		{"general.architecture", "granitehybrid"},
		{"granitehybrid.embedding_length", 4096},
		{"granitehybrid.ssm.state_size", 128},
	})
	writeGgufKVs(t, filepath.Join(root, "c-vl", "Qwen3VL-8B-Instruct-Q8_0.gguf"), [][2]any{
		{"general.architecture", "qwen3vl"},
		{"qwen3vl.embedding_length", 4096},
	})

	p := ScanEncoderPool([]string{root})
	for _, f := range p.Files {
		if f.Role == RoleLlm && !strings.Contains(f.Path, "Qwen3VL") {
			t.Errorf("recurrent llm in pool: %s", f.Path)
		}
	}
	if got, _ := p.Llm(4096, false, ""); !strings.Contains(got, "Qwen3VL") {
		t.Errorf("llm(4096) = %q, want the Qwen3-VL encoder", got)
	}
}