package autogen

import (
	"strings"
	"testing"

	"gopkg.in/yaml.v3"
)

func TestAutogen_promptEnhancerName(t *testing.T) {
	cases := []struct {
		name       string
		id         string
		wantOK     bool
		wantDir    string
		wantFamily string
	}{
		// The pair this feature was built for, as prithivMLmods ships them.
		{"qwen i2i", "qwen-image-2.1-pe-i2i-q5_k_m", true, PEDirEdit, "qwen-image-2.1"},
		{"qwen t2i", "Qwen-Image-2.1-PE-T2I-Q8_0", true, PEDirText, "qwen-image-2.1"},
		// A filename, which is what the encoder pool holds.
		{"full path", "/d/llm/models/Qwen-Image-2.1-PE-I2I-BF16.gguf", true, PEDirEdit, "qwen-image-2.1"},
		// Underscores must fold, or the pairing misses its own image model.
		{"underscores", "qwen_image_2.1_pe_t2i", true, PEDirText, "qwen-image-2.1"},
		// Spelled out, no direction: usable, just not direction-specific.
		{"spelled", "someones-prompt-enhancer-7b", true, "", "someones"},
		{"spelled joined", "flux-promptenhancer-v2", true, "", "flux"},

		// The false positives that would cost a user their text encoder.
		{"pegasus", "pegasus-7b-q4_k_m", false, "", ""},
		{"pe alone", "gpt-pe-7b", false, "", ""},
		{"direction alone", "flux-t2i-turbo", false, "", ""},
		{"ordinary chat model", "qwen3.6-27b-q4_k_m", false, "", ""},
		{"the image model itself", "qwen-image-2.1-q8_0", false, "", ""},
		{"empty", "", false, "", ""},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			dir, family, ok := promptEnhancerName(c.id)
			if ok != c.wantOK {
				t.Fatalf("ok = %v, want %v", ok, c.wantOK)
			}
			if !ok {
				return
			}
			if dir != c.wantDir {
				t.Errorf("dir = %q, want %q", dir, c.wantDir)
			}
			if family != c.wantFamily {
				t.Errorf("family = %q, want %q", family, c.wantFamily)
			}
		})
	}
}

// The pairing has to survive the thing that makes ids differ in practice: the
// same model at another quant, on either side.
func TestAutogen_detectEnhancers_PairsAcrossQuants(t *testing.T) {
	rows := []GgufRow{
		{ID: "qwen-image-2.1-pe-t2i-q5_k_m"},
		{ID: "qwen-image-2.1-pe-i2i-q5_k_m"},
		{ID: "qwen3.6-27b-q4_k_m"},
	}
	auto := detectEnhancers(rows)

	got := auto.For("qwen-image-2.1-q8_0")
	if got == nil {
		t.Fatal("no enhancers paired to the image model")
	}
	if got.Text != "qwen-image-2.1-pe-t2i-q5_k_m" {
		t.Errorf("text = %q", got.Text)
	}
	if got.Edit != "qwen-image-2.1-pe-i2i-q5_k_m" {
		t.Errorf("edit = %q", got.Edit)
	}
	if auto.For("flux-dev-q8_0") != nil {
		t.Error("paired an unrelated image model")
	}
}

// Enhancers used to get a "-vision" twin and an i2i one was paired to it. The
// twin is gone (the base profile holds the projector in VRAM), so a pin saved
// back then must land on the base id instead of naming a model that no longer
// exists. A non-enhancer "-vision" id is a real twin and is left alone.
func TestAutogen_resolveEnhancerIDs_LegacyVisionPin(t *testing.T) {
	text, edit := resolveEnhancerIDs(Settings{}, "qwen-image-2.1-q8_0", "qwen3-vl-8b-vision", "qwen-image-2.1-pe-i2i-q5_k_m-vision")
	if edit != "qwen-image-2.1-pe-i2i-q5_k_m" {
		t.Errorf("edit = %q, want the base id", edit)
	}
	if text != "qwen3-vl-8b-vision" {
		t.Errorf("text = %q, a non-enhancer twin must be kept", text)
	}
}

// A PE model must never be offered as a text encoder: it outranks the real
// encoder on file size and the failure is a confidently unrelated image.
func TestAutogen_IsPromptEnhancerFile(t *testing.T) {
	if !IsPromptEnhancerFile("/d/models/Qwen-Image-2.1-PE-I2I-BF16.gguf") {
		t.Error("PE gguf not recognised")
	}
	// The quant tag joined by "." to the direction token, as prithivMLmods
	// names them. Missed, the file is a qwen35 4096-wide hybrid, the same shape
	// MOSS-Image's real encoder has.
	if !IsPromptEnhancerFile("/d/models/Qwen-Image-2.1-PE-I2I.Q5_K_M.gguf") {
		t.Error("PE gguf with a dotted quant tag not recognised")
	}
	if IsPromptEnhancerFile("/d/models/Qwen3-VL-8B-Instruct-Q8_0.gguf") {
		t.Error("a real text encoder was withheld from the pool")
	}
}

// Tier 2: what a model ends up carrying. The three cases that matter are the
// empty field (fill it), a configured one (never touch it) and the explicit
// refusal (which only exists because auto-fill gave "" a second meaning).
func TestAutogen_resolveEnhancerIDs(t *testing.T) {
	s := Settings{autoEnhancers: detectEnhancers([]GgufRow{
		{ID: "qwen-image-2.1-pe-t2i-q5_k_m"},
		{ID: "qwen-image-2.1-pe-i2i-q5_k_m"},
	})}

	text, edit := resolveEnhancerIDs(s, "qwen-image-2.1-q8_0", "", "")
	if text != "qwen-image-2.1-pe-t2i-q5_k_m" || edit != "qwen-image-2.1-pe-i2i-q5_k_m" {
		t.Errorf("auto-fill: text=%q edit=%q", text, edit)
	}

	text, edit = resolveEnhancerIDs(s, "qwen-image-2.1-q8_0", "my-own-rewriter", "")
	if text != "my-own-rewriter" {
		t.Errorf("a configured id must win, got %q", text)
	}
	if edit != "qwen-image-2.1-pe-i2i-q5_k_m" {
		t.Errorf("the other direction still fills, got %q", edit)
	}

	// PEDisabled survives resolution (writeEnhancerKey drops it), which is what
	// keeps a regeneration from re-adding an enhancer the user removed.
	text, _ = resolveEnhancerIDs(s, "qwen-image-2.1-q8_0", PEDisabled, "")
	if text != PEDisabled {
		t.Errorf("an explicit refusal must not be overwritten, got %q", text)
	}

	// Nothing detected beside it => unchanged, not invented.
	if text, edit = resolveEnhancerIDs(Settings{}, "flux-dev-q8_0", "", ""); text != "" || edit != "" {
		t.Errorf("no detection table should mean no enhancer, got %q/%q", text, edit)
	}
}

// Tier 3: the per-model system prompt, which is the one field a user pastes a
// multi-KB document into.
func TestAutogen_writeEnhancerPrompt(t *testing.T) {
	var b strings.Builder
	writeEnhancerPrompt(&b, "promptEnhancerPrompt", "pe-t2i", trickyPrompt)

	var parsed map[string]string
	if err := yaml.Unmarshal([]byte(b.String()), &parsed); err != nil {
		t.Fatalf("emitted prompt does not parse: %v\n%s", err, b.String())
	}
	if parsed["promptEnhancerPrompt"] != trickyPrompt {
		t.Errorf("prompt mangled:\n got %q\nwant %q", parsed["promptEnhancerPrompt"], trickyPrompt)
	}

	// A prompt with no enhancer to run under is not config, it is a leftover.
	b.Reset()
	writeEnhancerPrompt(&b, "promptEnhancerPrompt", "", trickyPrompt)
	writeEnhancerPrompt(&b, "promptEnhancerPrompt", PEDisabled, trickyPrompt)
	writeEnhancerPrompt(&b, "promptEnhancerPrompt", "pe-t2i", "")
	if b.String() != "" {
		t.Errorf("expected nothing, got %q", b.String())
	}
}

// An enhancer is emitted as ONE profile shaped for the job: no fleet tiers, no
// vision twin, a capped window, and the projector placed by direction (i2i in
// VRAM because every call carries an image, t2i none at all). Gated on the real
// models tree holding a Qwen PE pair.
func TestAutogen_Generate_PromptEnhancerSingleProfile(t *testing.T) {
	if realModelsRoot == "" {
		t.Skip("no real models tree")
	}
	gf := GenerateFile{Settings: Settings{ModelsRoot: realModelsRoot}}
	gf.Settings.applyDefaults()
	gf.Settings.DefaultVariants = []VariantSpec{{Name: "32k", Ctx: 32768}, {Name: "game", Ctx: 16384}}
	out, err := Generate(gf, "T")
	if err != nil {
		t.Fatalf("generate: %v", err)
	}
	seen := map[string]bool{}
	for id, cmd := range modelBlocks(out) {
		dir, _, ok := promptEnhancerName(id)
		if !ok {
			continue
		}
		if !strings.Contains(cmd, "-c 32768") {
			t.Errorf("%s: want the enhancer window -c 32768", id)
		}
		if strings.HasSuffix(id, "-vision") || strings.HasSuffix(id, "-32k") || strings.HasSuffix(id, "-game") {
			t.Errorf("%s: enhancer must emit no extra profiles", id)
		}
		switch dir {
		case PEDirEdit:
			if strings.Contains(cmd, "--mmproj ") && strings.Contains(cmd, "--no-mmproj-offload") {
				t.Errorf("%s: i2i projector must be in VRAM", id)
			}
		case PEDirText:
			if strings.Contains(cmd, "--mmproj ") {
				t.Errorf("%s: t2i is never shown an image, want no projector", id)
			}
		}
		seen[dir] = true
	}
	if !seen[PEDirEdit] && !seen[PEDirText] {
		t.Skip("no prompt enhancer in the models tree")
	}
}