package catalog

import (
	"strings"
	"testing"
)

func comfyModel(source string, c *ComfyUI, profiles ...Profile) Model {
	if len(profiles) == 0 {
		profiles = []Profile{{Name: "gpu", Engine: EngineComfyUI, GPUClass: "6000-24g", TP: 1}}
	}
	return Model{Name: "flux", Source: source, ComfyUI: c, Profiles: profiles}
}

func TestAComfyUIModelMakesPicturesFromItsSourcesFile(t *testing.T) {
	cat, err := New([]Model{comfyModel("hf:Comfy-Org/flux1-schnell/flux1-schnell-fp8.safetensors", nil)})
	if err != nil {
		t.Fatal(err)
	}
	m, _ := cat.Resolve("flux")
	if !m.MakesImages() || !m.ServesComfyUI() {
		t.Error("a ComfyUI model does not say it makes pictures")
	}
	if w, ckpt := m.ComfyWorkflow(); w != "" || ckpt != "flux1-schnell-fp8.safetensors" {
		t.Errorf("workflow %q, checkpoint %q", w, ckpt)
	}
}

func TestAComfyUIModelIsRefusedWhatItCannotDo(t *testing.T) {
	for name, c := range map[string]struct {
		m    Model
		want string
	}{
		"a set without a checkpoint": {comfyModel("https://mirror/flux-dev/v1/", nil), "comfyui.checkpoint"},
		"an unknown workflow":        {comfyModel("/w/f.safetensors", &ComfyUI{Workflow: "sdxl"}), "no built-in workflow"},
		"mixed engines": {comfyModel("/w/f.safetensors", nil,
			Profile{Name: "a", Engine: EngineComfyUI, GPUClass: "6000-24g", TP: 1},
			Profile{Name: "b", Engine: "vllm", GPUClass: "6000-24g", TP: 1}), "every profile"},
		"comfyui settings without the engine": {Model{Name: "x", Source: "/w/x", ComfyUI: &ComfyUI{},
			Profiles: []Profile{{Name: "a", Engine: "vllm", GPUClass: "6000-24g", TP: 1}}}, "engine comfyui"},
		"no card": {comfyModel("/w/f.safetensors", nil, Profile{Name: "a", Engine: EngineComfyUI, TP: 1}), "gpu_class"},
	} {
		if _, err := New([]Model{c.m}); err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: %v", name, err)
		}
	}
	ok := comfyModel("https://mirror/flux-dev/v1/", &ComfyUI{Workflow: "flux-dev", Checkpoint: "flux1-dev-fp8.safetensors"})
	if _, err := New([]Model{ok}); err != nil {
		t.Errorf("a set with its checkpoint named was refused: %v", err)
	}
}