package autogen

import (
	"strings"
	"testing"
)

// Issue #57: an offloaded extra model places its VAE through the one --backend
// spec, not the deprecated --vae-on-cpu beside it.
func TestExtraImageCmdLines_OffloadVaeInBackendSpec(t *testing.T) {
	s := Settings{SdServerExe: "sd-server"}
	m := ExtraImageModel{ModelPath: "C:/models/dit.safetensors", OffloadToCpu: "on"}
	cmd := strings.Join(extraImageCmdLines(s, m).Lines, " ")
	if !strings.Contains(cmd, "--backend te=cpu,vae=cpu") || !strings.Contains(cmd, "--offload-to-cpu") {
		t.Errorf("offloaded extra model should carry te=cpu,vae=cpu:\n%s", cmd)
	}
	if strings.Contains(cmd, "--vae-on-cpu") {
		t.Errorf("deprecated --vae-on-cpu still emitted:\n%s", cmd)
	}
	// Both knobs on must not duplicate the assignment.
	m.VaeOnCpu = "on"
	if cmd = strings.Join(extraImageCmdLines(s, m).Lines, " "); strings.Count(cmd, "vae=cpu") != 1 {
		t.Errorf("vae=cpu should appear once:\n%s", cmd)
	}
}

// A hand-declared extra image model carries its own backend pin (the model
// editor offers the same sd rows), so the class default must not move it.
func TestExtraImageCmdLines_BackendPin(t *testing.T) {
	rows := []BackendEntry{
		{ID: "managed-rocm", Kind: "sd", Path: "C:/rocm/sd-server.exe", Default: true},
		{ID: "build-vulkan", Kind: "sd", Path: "C:/vulkan/sd-server.exe"},
	}
	s := Settings{SdServerExe: "C:/rocm/sd-server.exe", Backends: rows}
	m := ExtraImageModel{ModelPath: "C:/models/dit.safetensors", Backend: "build-vulkan"}

	if got := extraImageCmdLines(s, m).Lines[0]; got != "C:/vulkan/sd-server.exe" {
		t.Errorf("pinned exe = %q, want the vulkan build", got)
	}
	m.Backend = ""
	if got := extraImageCmdLines(s, m).Lines[0]; got != "C:/rocm/sd-server.exe" {
		t.Errorf("auto exe = %q, want the class default", got)
	}

	// The editor reaches the same cmd by applying the override onto the model.
	m = ApplyOverrideToExtraImage(m, &Override{Backend: "build-vulkan"})
	if got := RenderExtraImageCmd(s, m); !strings.HasPrefix(got, "C:/vulkan/sd-server.exe ") {
		t.Errorf("preview = %q, want the pinned build first", got)
	}
}

func TestIsImageArch(t *testing.T) {
	for _, a := range []string{"flux", "FLUX", "flux.1", "sd3", "sd3.5", "qwen_image", "z_image", " stable-diffusion "} {
		if !isImageArch(a) {
			t.Errorf("isImageArch(%q) = false, want true", a)
		}
	}
	for _, a := range []string{"", "qwen3moe", "llama", "gemma4"} {
		if isImageArch(a) {
			t.Errorf("isImageArch(%q) = true, want false", a)
		}
	}
}

func TestEmitImageModel(t *testing.T) {
	var b strings.Builder
	var emitted []string
	onDisk(t, "ae.safetensors", "clip_l.safetensors", "t5xxl.gguf", "zae.safetensors", "qwen3.gguf")
	pool := EncoderSet{FluxVae: "ae.safetensors", ClipL: "clip_l.safetensors", T5: "t5xxl.gguf", ZimageVae: "zae.safetensors", QwenLlm: "qwen3.gguf"}
	s := Settings{SdServerExe: "sd-server", TtlSec: 600, TargetVramGB: 7, VramOverheadGB: 0.5, Threads: 7, Encoders: pool}

	// Big model (6.5GB + 1.5 compute > 6.5 budget) → offload path. No per-model
	// override: VAE + CLIP-L + T5 must be auto-attached from the arch (flux) pool.
	big := GgufRow{FullPath: `C:\models\flux.gguf`, SizeGB: 6.5}
	emitImageModel(&b, s, big, &Override{}, "flux-q4", "flux", false, 0, &emitted)
	out := b.String()
	for _, want := range []string{"sd-server", "--diffusion-model C:/models/flux.gguf", "--listen-port ${PORT}", "--max-vram 6.5", "--vae ae.safetensors", "--clip_l clip_l.safetensors", "--t5xxl t5xxl.gguf", "--diffusion-fa", "--vae-tiling", "--offload-to-cpu", "--backend te=cpu,vae=cpu", "offload=true", "checkEndpoint: /", "out: [image]", "in: [text]", "ttl: 600"} {
		if !strings.Contains(out, want) {
			t.Errorf("emit missing %q:\n%s", want, out)
		}
	}
	// One --backend spec carries the VAE placement; the deprecated flag is gone.
	if strings.Contains(out, "--vae-on-cpu") {
		t.Errorf("offload should place the VAE via --backend, not --vae-on-cpu:\n%s", out)
	}
	if strings.Contains(out, "WARNING") || strings.Contains(out, "missingComponents") {
		t.Errorf("complete pool should not warn:\n%s", out)
	}
	if len(emitted) != 1 || emitted[0] != "flux-q4" {
		t.Errorf("emitted = %v, want [flux-q4]", emitted)
	}

	// Override knobs: component paths + gen defaults emit, toggled-off savers omit,
	// offload pinned off wins over the auto "fits" decision.
	var b3 strings.Builder
	var em3 []string
	ov := &Override{
		VaePath: `C:\models\ae.safetensors`, ClipLPath: "clip_l.gguf", T5Path: "t5xxl.gguf",
		TextEncoderPath: "qwen3.gguf", VaeTiling: "off", DiffusionFa: "off", TeOnCpu: "off",
		OffloadToCpu: "off", DefaultSteps: 8, DefaultCfg: 1.0, DefaultSampler: "euler",
		DefaultWidth: 768, DefaultHeight: 512, ExtraArgs: "--clip-on-cpu",
	}
	emitImageModel(&b3, s, big, ov, "flux-tuned", "flux", false, 0, &em3)
	out3 := b3.String()
	for _, want := range []string{"--vae C:/models/ae.safetensors", "--clip_l clip_l.gguf", "--t5xxl t5xxl.gguf", "--llm qwen3.gguf", "--steps 8", "--cfg-scale 1", "--sampling-method euler", "--width 768", "--height 512", "--clip-on-cpu", "offload=false"} {
		if !strings.Contains(out3, want) {
			t.Errorf("tuned emit missing %q:\n%s", want, out3)
		}
	}
	for _, unwant := range []string{"--vae-tiling", "--diffusion-fa", "--backend te=cpu", "--offload-to-cpu"} {
		if strings.Contains(out3, unwant) {
			t.Errorf("tuned emit should omit %q:\n%s", unwant, out3)
		}
	}

	// Small model (2GB + 1.5 < 6.5 budget) → fits resident, no offload flags.
	var b2 strings.Builder
	var em2 []string
	small := GgufRow{FullPath: `C:\models\sd15.gguf`, SizeGB: 2.0}
	emitImageModel(&b2, s, small, &Override{}, "sd15-q4", "sd1", true, 0, &em2)
	out2 := b2.String()
	// SD1 is a full-checkpoint arch → -m, no external components.
	if !strings.Contains(out2, "-m C:/models/sd15.gguf") || strings.Contains(out2, "--diffusion-model") {
		t.Errorf("sd1 should load via -m, not --diffusion-model:\n%s", out2)
	}
	if strings.Contains(out2, "--offload-to-cpu") || strings.Contains(out2, "offload=true") {
		t.Errorf("small model should not offload:\n%s", out2)
	}
	// vae=cpu rides with offload only; a resident model keeps VAE on GPU.
	if strings.Contains(out2, "vae=cpu") || strings.Contains(out2, "--vae-on-cpu") {
		t.Errorf("resident model should not force vae-on-cpu:\n%s", out2)
	}
	// vae-tiling is always on (caps the VAE decode VRAM spike), even when resident.
	if !strings.Contains(out2, "--vae-tiling") {
		t.Errorf("small model should still vae-tile:\n%s", out2)
	}
	// te=cpu is unconditional (external encoder always parked on CPU), even when the
	// diffusion weights fit resident.
	if !strings.Contains(out2, "--backend te=cpu") {
		t.Errorf("small model should still set te=cpu:\n%s", out2)
	}
}

func TestFluxGuidance(t *testing.T) {
	cases := map[string]struct {
		val float64
		ok  bool
	}{
		"flux1-kontext-dev-Q4_K_M": {2.5, true},
		"flux1-fill-dev-Q4_K_S":    {30, true},
		"flux1-schnell-Q4_K_S":     {0, false},
		"flux1-dev":                {0, false},
	}
	for name, want := range cases {
		if v, ok := fluxGuidance(name); v != want.val || ok != want.ok {
			t.Errorf("fluxGuidance(%q) = (%g,%t), want (%g,%t)", name, v, ok, want.val, want.ok)
		}
	}
}

func TestResolveComponents(t *testing.T) {
	pool := EncoderSet{FluxVae: "ae", ClipL: "cl", ClipG: "cg", T5: "t5", ZimageVae: "zae", QwenLlm: "qwen"}

	// flux: vae + clip_l + t5, nothing missing.
	if c, m := resolveComponents(pool, nil, "flux", "flux1-schnell", nil, 0, videoInfo{}); c.vae != "ae" || c.clipL != "cl" || c.t5 != "t5" || c.clipG != "" || c.llm != "" || len(m) != 0 {
		t.Errorf("flux: got %+v missing=%v", c, m)
	}
	// chroma (arch flux, name-detected): vae + t5, NO clip_l.
	if c, m := resolveComponents(pool, nil, "flux", "Chroma1-HD-Q5_K_M", nil, 0, videoInfo{}); c.vae != "ae" || c.t5 != "t5" || c.clipL != "" || len(m) != 0 {
		t.Errorf("chroma should be vae+t5 only: got %+v missing=%v", c, m)
	}
	// z-image (arch lumina2): vae + llm.
	if c, m := resolveComponents(pool, nil, "lumina2", "Z-Image-Turbo", nil, 0, videoInfo{}); c.vae != "zae" || c.llm != "qwen" || c.clipL != "" || len(m) != 0 {
		t.Errorf("z-image should be vae+llm: got %+v missing=%v", c, m)
	}
	// sdxl: full-checkpoint arch → never resolves external components (loads via -m).
	// resolveComponents isn't called for it in emit, and falls through to attach nothing.
	if c, m := resolveComponents(pool, nil, "sdxl", "animagineXLV31", nil, 0, videoInfo{}); c.vae != "" || c.clipL != "" || c.clipG != "" || len(m) != 0 {
		t.Errorf("sdxl should attach nothing (served via -m): got %+v missing=%v", c, m)
	}
	// empty pool: flux reports every required role missing.
	if _, m := resolveComponents(EncoderSet{}, nil, "flux", "flux1-fill-dev", nil, 0, videoInfo{}); len(m) != 3 {
		t.Errorf("empty-pool flux should miss 3 roles, got %v", m)
	}
	// override supplies what the pool lacks → wins and clears the missing role.
	ov := &Override{VaePath: "ov-ae", ClipLPath: "ov-cl", T5Path: "ov-t5"}
	if c, m := resolveComponents(EncoderSet{}, ov, "flux", "flux1-fill-dev", nil, 0, videoInfo{}); c.vae != "ov-ae" || c.clipL != "ov-cl" || c.t5 != "ov-t5" || len(m) != 0 {
		t.Errorf("override should win and clear missing: got %+v missing=%v", c, m)
	}
}

// LongCat ships its DiT in BFL (flux) tensor layout, so it reports
// general.architecture "flux" and would otherwise be wired like flux.1:
// clip_l + t5xxl and no --llm. It actually wants flux.1's AE plus a Qwen2.5-VL
// text encoder, and the edit variant additionally wants that encoder's vision
// projector. Nothing here is declared: the components come from the scanned
// pool, matched on the width the DiT itself asks for.
func TestResolveComponents_LongCatFromPool(t *testing.T) {
	pool := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/ae.safetensors", Role: RoleVae, Family: VaeFamilyFlux},
		{Path: "/m/clip_l.safetensors", Role: RoleClip, Width: 768},
		{Path: "/m/t5xxl.gguf", Role: RoleT5, Width: 4096},
		{Path: "/m/qwen25vl/model-Q8_0.gguf", Role: RoleLlm, Width: 3584, SizeGB: 8,
			Mmproj: "/m/qwen25vl/mmproj-F16.gguf", Vision: true},
		{Path: "/m/qwen3-4b.gguf", Role: RoleLlm, Width: 2560, SizeGB: 3},
	}}

	// Edit variant: vae + llm + the paired projector, and NO clip_l/t5.
	c, m := resolveComponents(EncoderSet{}, nil, "flux", "LongCat-Image-Edit-Turbo-Q8_0", pool, 3584, videoInfo{})
	if len(m) != 0 {
		t.Fatalf("missing = %v, want none", m)
	}
	if c.vae != "/m/ae.safetensors" {
		t.Errorf("vae = %q", c.vae)
	}
	if c.llm != "/m/qwen25vl/model-Q8_0.gguf" {
		t.Errorf("llm = %q", c.llm)
	}
	if c.llmVision != "/m/qwen25vl/mmproj-F16.gguf" {
		t.Errorf("llm_vision = %q", c.llmVision)
	}
	if c.clipL != "" || c.t5 != "" || c.clipG != "" {
		t.Errorf("longcat must not take clip/t5: %+v", c)
	}

	// Base (non-edit) LongCat: same encoder, no projector.
	c2, _ := resolveComponents(EncoderSet{}, nil, "flux", "LongCat-Image-Q8_0", pool, 3584, videoInfo{})
	if c2.llm != "/m/qwen25vl/model-Q8_0.gguf" || c2.llmVision != "" {
		t.Errorf("base longcat: %+v", c2)
	}

	// A different caption width picks a different encoder with no name table.
	c3, _ := resolveComponents(EncoderSet{}, nil, "lumina2", "Z-Image-Turbo", pool, 2560, videoInfo{})
	if c3.llm != "/m/qwen3-4b.gguf" {
		t.Errorf("z-image llm = %q", c3.llm)
	}

	// A declared encoder still wins, and drags its own neighbour projector.
	ov := &Override{TextEncoderPath: "/m/qwen25vl/model-Q8_0.gguf"}
	c4, _ := resolveComponents(EncoderSet{}, ov, "flux", "LongCat-Image-Edit", pool, 3584, videoInfo{})
	if c4.llm != ov.TextEncoderPath || c4.llmVision != "/m/qwen25vl/mmproj-F16.gguf" {
		t.Errorf("declared encoder: %+v", c4)
	}

	// llmVision "off" drops the projector without touching the encoder pick.
	c5, _ := resolveComponents(EncoderSet{}, &Override{LlmVision: "off"}, "flux", "LongCat-Image-Edit", pool, 3584, videoInfo{})
	if c5.llm == "" || c5.llmVision != "" {
		t.Errorf("llmVision off: %+v", c5)
	}
}

// Width narrows the field, the declaration names the file. A pool holding both a
// Qwen3-4B and a Qwen3-VL-4B of the same caption width (both are 2560) used to
// resolve to the VL file on the size/path tiebreak no matter what
// settings.encoders.qwenLlm pinned, which wired Z-Image to the encoder Krea-2
// wants. Regression: 6aea9f5 (discover image encoders instead of declaring them).
func TestResolveComponents_DeclaredLlmWins(t *testing.T) {
	const pin = "/m/lmstudio/Qwen3-4B-Instruct-2507-Q8_0.gguf"
	pool := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/QwenVL/Qwen3VL-4B-Instruct-Q8_0.gguf", Role: RoleLlm, Width: 2560, SizeGB: 3.99, Vision: true,
			Mmproj: "/m/QwenVL/mmproj-Qwen3VL-4B-Instruct-F16.gguf"},
		{Path: pin, Role: RoleLlm, Width: 2560, SizeGB: 3.99},
	}}
	onDisk(t, "ae.safetensors")
	enc := EncoderSet{ZimageVae: "ae.safetensors", QwenLlm: pin}
	c, m := resolveComponents(enc, nil, "lumina2", "Z-Image-Turbo", pool, 2560, videoInfo{})
	if len(m) != 0 {
		t.Fatalf("missing = %v, want none", m)
	}
	if c.llm != pin {
		t.Errorf("llm = %q, want the declared %q", c.llm, pin)
	}
	if c.vae != enc.ZimageVae {
		t.Errorf("vae = %q", c.vae)
	}
	// Krea-2 is the model the VL file is right for, and a per-model override
	// still beats both the declaration and the scan.
	ov := &Override{TextEncoderPath: "/m/QwenVL/Qwen3VL-4B-Instruct-Q8_0.gguf"}
	c2, _ := resolveComponents(enc, ov, "qwen_image", "krea2_turbo-Q8_0", pool, 2560, videoInfo{})
	if c2.llm != ov.TextEncoderPath {
		t.Errorf("per-model override lost to the global pin: %q", c2.llm)
	}
}

func TestMergeImageVariant(t *testing.T) {
	base := Override{
		VaePath: "ae.safetensors", TextEncoderPath: "qwen3.gguf",
		DefaultSteps: 30, DefaultCfg: 7, VramTargetGB: 6,
	}
	// A "fast" preset overrides only steps/cfg.
	v := VariantSpec{Name: "fast", DefaultSteps: 8, DefaultCfg: 1}
	got := mergeImageVariant(base, v)
	if got.VaePath != "ae.safetensors" || got.TextEncoderPath != "qwen3.gguf" {
		t.Errorf("preset should inherit component paths, got %+v", got)
	}
	if got.DefaultSteps != 8 || got.DefaultCfg != 1 {
		t.Errorf("preset should override steps/cfg, got steps=%d cfg=%g", got.DefaultSteps, got.DefaultCfg)
	}
	if got.VramTargetGB != 6 {
		t.Errorf("preset should inherit vram budget, got %g", got.VramTargetGB)
	}
	if got.Variants != nil {
		t.Errorf("merged override must clear Variants to avoid re-emit")
	}
}

// Flux.2 is the only family whose caption width is not its encoder width: it
// concatenates 3 LLM hidden layers, so klein 9B states txt_in 12288 for a
// 4096-wide Qwen3-8B. The variants differ only in that arithmetic, which is why
// none of them is named here -- the width match resolves all three.
func TestResolveComponents_Flux2Klein(t *testing.T) {
	onDisk(t, "f2ae", "elsewhere-qwen3-8b.gguf")
	enc := EncoderSet{Flux2Vae: "f2ae", QwenLlm: "/m/qwen3-4b.gguf"}
	pool := &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},
	}}
	// klein 9B: the flux.2 VAE and the 8B encoder, no clip_l/t5.
	if c, m := resolveComponents(enc, nil, "flux", "Flux2-Klein-9B-True-v2-Q8_0", pool, 12288, videoInfo{}); c.vae != "f2ae" || c.llm != "/m/qwen3-8b.gguf" || c.clipL != "" || c.t5 != "" || len(m) != 0 {
		t.Errorf("klein 9B: got %+v missing=%v", c, m)
	}
	// klein 4B: same branch, 7680 = 3 x 2560, resolves to the smaller encoder
	// with no extra case in resolveComponents.
	if c, m := resolveComponents(enc, nil, "flux", "flux2-klein-4b-Q8_0", pool, 7680, videoInfo{}); c.llm != "/m/qwen3-4b.gguf" || len(m) != 0 {
		t.Errorf("klein 4B: got %+v missing=%v", c, m)
	}
	// The encoder is not on disk. The global qwenLlm pin is 2560 wide and the
	// scan classified it, so substituting it would wire a 12288-caption DiT to a
	// 2560-wide encoder and emit NOTHING -- req only warns on an empty path.
	// Report the role missing instead.
	only4b := &EncoderPool{Files: []ComponentFile{
		{Path: "/m/qwen3-4b.gguf", Role: RoleLlm, Family: "qwen3", Width: 2560, SizeGB: 3},
	}}
	if c, m := resolveComponents(enc, nil, "flux", "Flux2-Klein-9B-True-v2-Q8_0", only4b, 12288, videoInfo{}); c.llm != "" || len(m) != 1 || m[0] != "llm" {
		t.Errorf("klein 9B without its encoder should report llm missing: got %+v missing=%v", c, m)
	}
	// A pin the scan never saw (kept outside the models root) is still honoured:
	// that is the case the fallback exists for.
	unseen := EncoderSet{Flux2Vae: "f2ae", QwenLlm: "elsewhere-qwen3-8b.gguf"}
	if c, m := resolveComponents(unseen, nil, "flux", "Flux2-Klein-9B-True-v2-Q8_0", only4b, 12288, videoInfo{}); c.llm != "elsewhere-qwen3-8b.gguf" || len(m) != 0 {
		t.Errorf("an unscanned pin should still be honoured: got %+v missing=%v", c, m)
	}
}

func TestImageDefaultsFor(t *testing.T) {
	cases := []struct {
		name, arch, model string
		cond              int64
		steps             int
		cfg               float64
	}{
		{"2.1 turbo", "qwen_image", "qwen_image_2.1_turbo-q8_0", 4096, 8, 1.0},
		{"krea2 turbo", "qwen_image", "krea2_turbo-q8_0", 2560, 8, 1.0},
		{"2.1 base", "qwen_image", "qwen_image_2.1-q8_0", 4096, 40, 1.0},
		{"20B line untouched", "qwen_image", "qwen-image-edit-q4_k_m", 3584, 0, 0},
		{"other arch untouched", "flux", "flux1-schnell-turbo", 4096, 0, 0},
	}
	for _, c := range cases {
		d := imageDefaultsFor(c.arch, c.model, c.cond)
		if d.steps != c.steps || d.cfg != c.cfg {
			t.Errorf("%s: got steps %d cfg %g, want %d / %g", c.name, d.steps, d.cfg, c.steps, c.cfg)
		}
	}
	// The override still wins over the arch profile.
	d := resolveGenDefaults(videoInfo{}, "qwen_image", "qwen_image_2.1_turbo-q8_0", 4096, &Override{DefaultSteps: 12})
	if d.steps != 12 || d.cfg != 1.0 {
		t.Errorf("override: got steps %d cfg %g, want 12 / 1", d.steps, d.cfg)
	}
}