package modules

import (
	"regexp"
	"strings"
	"testing"
)

func load(t *testing.T) *Catalog {
	t.Helper()
	c, err := Load()
	if err != nil {
		t.Fatal(err)
	}
	return c
}

func TestCatalogLoads(t *testing.T) {
	c := load(t)
	want := []string{"audio.transcribe", "image.edit", "image.generate", "video.edit", "video.generate"}
	got := c.IDs()
	if strings.Join(got, ",") != strings.Join(want, ",") {
		t.Fatalf("capabilities = %v, want %v", got, want)
	}
}

func TestSizesDedupe(t *testing.T) {
	c := load(t)
	all := c.Sizes(c.IDs())
	var sum int64
	for _, id := range c.IDs() {
		cp, _ := c.Capability(id)
		sum += cp.ModelBytes()
	}
	if all.Models != sum {
		t.Fatalf("models = %d, want %d (no file is shared today)", all.Models, sum)
	}
	// three ComfyUI capabilities share one engine; the image is also shared with whisper (same digest)
	e, _ := c.Engine("comfyui")
	w, _ := c.Engine("whisper")
	if want := e.ImageBytes + e.WorkBytes + w.WorkBytes; all.Engines != want {
		t.Fatalf("engines = %d, want %d", all.Engines, want)
	}
	one := c.Sizes([]string{"image.generate", "image.generate"})
	cp, _ := c.Capability("image.generate")
	if one.Models != cp.ModelBytes() {
		t.Fatalf("a repeated id counted twice: %d", one.Models)
	}
}

func TestValidate(t *testing.T) {
	c := load(t)
	ig, _ := c.Capability("image.generate")
	if _, err := ig.Validate(map[string]any{}); err == nil {
		t.Fatal("missing prompt accepted")
	}
	if _, err := ig.Validate(map[string]any{"prompt": "x", "bogus": 1}); err == nil {
		t.Fatal("unknown input accepted")
	}
	if _, err := ig.Validate(map[string]any{"prompt": "x", "count": 9}); err == nil {
		t.Fatal("count above max accepted")
	}
	if _, err := ig.Validate(map[string]any{"prompt": "x", "size": "9999x9999"}); err == nil {
		t.Fatal("size outside the choices accepted")
	}
	out, err := ig.Validate(map[string]any{"prompt": "  a cat ", "count": float64(2)})
	if err != nil {
		t.Fatal(err)
	}
	if out["prompt"] != "a cat" || out["count"] != int64(2) || out["size"] != "1024x1024" {
		t.Fatalf("got %v", out)
	}

	ve, _ := c.Capability("video.edit")
	if _, err := ve.Validate(map[string]any{"operation": "extend", "video": "a.mp4"}); err == nil {
		t.Fatal("unknown operation accepted")
	}
	if _, err := ve.Validate(map[string]any{"video": "../a.mp4"}); err == nil {
		t.Fatal("path in a file name accepted")
	}
	// multiplier only matters for interpolate: it is dropped for upscale
	out, err = ve.Validate(map[string]any{"operation": "upscale", "video": "a.mp4"})
	if err != nil {
		t.Fatal(err)
	}
	if _, has := out["multiplier"]; has {
		t.Fatal("multiplier kept for upscale")
	}
}

func TestVideoGenerateImageNeedsFrame(t *testing.T) {
	c := load(t)
	vg, _ := c.Capability("video.generate")
	if _, err := vg.Validate(map[string]any{"operation": "image", "prompt": "x"}); err == nil {
		t.Fatal("image mode without a picture accepted")
	}
	if _, err := vg.Validate(map[string]any{"prompt": "x"}); err != nil {
		t.Fatalf("text mode needs no picture: %v", err)
	}
}

func TestValuesDerived(t *testing.T) {
	c := load(t)
	vg, _ := c.Capability("video.generate")
	in, err := vg.Validate(map[string]any{"prompt": "x", "seconds": 3})
	if err != nil {
		t.Fatal(err)
	}
	v, err := vg.Values(in)
	if err != nil {
		t.Fatal(err)
	}
	if v["length"] != int64(73) || v["fps"] != int64(24) || v["width"] != int64(832) || v["height"] != int64(480) {
		t.Fatalf("got %v", v)
	}
	if s, ok := v["seed"].(int64); !ok || s < 0 {
		t.Fatalf("seed = %v", v["seed"])
	}
	ve, _ := c.Capability("video.edit")
	in, _ = ve.Validate(map[string]any{"video": "a.mp4", "scale": "2"})
	v, _ = ve.Values(in)
	if v["down"] != 0.5 {
		t.Fatalf("x2 should halve the 4x result, got %v", v["down"])
	}
}

func TestInstantiate(t *testing.T) {
	out, err := Instantiate([]byte(`{"1":{"inputs":{"text":"@prompt","n":"@n","keep":"plain"}}}`), map[string]any{"prompt": "@n", "n": int64(3)})
	if err != nil {
		t.Fatal(err)
	}
	in := out["1"].(map[string]any)["inputs"].(map[string]any)
	if in["text"] != "@n" || in["n"] != int64(3) || in["keep"] != "plain" {
		t.Fatalf("got %v", in) // user text is inserted as is, never expanded again
	}
	if _, err := Instantiate([]byte(`{"1":"@missing"}`), nil); err == nil {
		t.Fatal("a missing value was not reported")
	}
}

// Every @name in every workflow must be something the capability can supply for that operation.
func TestWorkflowPlaceholdersSatisfied(t *testing.T) {
	c := load(t)
	ref := regexp.MustCompile(`"@([a-z_]+)"`)
	for _, cp := range c.Capabilities {
		for op, file := range cp.Workflows {
			if file == BuiltinStory {
				continue // built in code: story_test.go checks it
			}
			raw, err := Workflow(file)
			if err != nil {
				t.Fatal(err)
			}
			in := map[string]any{"operation": op}
			for _, f := range cp.Inputs {
				switch f.Type {
				case "text":
					in[f.Name] = "x"
				case "image", "video", "audio":
					if f.shownFor(op) {
						in[f.Name] = "in.bin"
					}
				}
			}
			if op == "default" {
				delete(in, "operation")
			}
			ok, err := cp.Validate(in)
			if err != nil {
				t.Fatalf("%s/%s: %v", cp.ID, op, err)
			}
			vals, err := cp.Values(ok)
			if err != nil {
				t.Fatal(err)
			}
			vals["prefix"] = "wxch/job/out" // supplied by the agent
			for _, m := range ref.FindAllSubmatch(raw, -1) {
				if _, has := vals[string(m[1])]; !has {
					t.Errorf("%s/%s (%s): @%s cannot be filled", cp.ID, op, file, m[1])
				}
			}
			if _, err := Instantiate(raw, vals); err != nil {
				t.Errorf("%s/%s: %v", cp.ID, op, err)
			}
		}
	}
}

func TestVideoLengthUpToTwentySeconds(t *testing.T) {
	c := load(t)
	vg, _ := c.Capability("video.generate")
	for _, ok := range []float64{1, 3, 10, 12.5, 20} {
		if _, err := vg.Validate(map[string]any{"prompt": "x", "seconds": ok}); err != nil {
			t.Errorf("%v seconds must be accepted: %v", ok, err)
		}
	}
	for _, bad := range []float64{0, 0.5, 20.5, 30, 60} {
		if _, err := vg.Validate(map[string]any{"prompt": "x", "seconds": bad}); err == nil {
			t.Errorf("%v seconds must be refused", bad)
		}
	}
	in, _ := vg.Validate(map[string]any{"prompt": "x", "seconds": float64(20)})
	v, err := vg.Values(in)
	if err != nil {
		t.Fatal(err)
	}
	if v["length"] != int64(481) { // 4n+1 frames at 24 per second
		t.Errorf("20 seconds = %v frames", v["length"])
	}
}

// A video is decoded in tiles: a plain VAEDecode of 700 frames at once is what filled the memory of the machine at the end of a job.
func TestVideoWorkflowsDecodeInTiles(t *testing.T) {
	for _, f := range []string{"video.generate.text.json", "video.generate.image.json"} {
		raw, err := Workflow(f)
		if err != nil {
			t.Fatal(err)
		}
		s := string(raw)
		if strings.Contains(s, `"VAEDecode"`) || !strings.Contains(s, `"VAEDecodeTiled"`) || !strings.Contains(s, `"temporal_size"`) {
			t.Errorf("%s must decode in tiles in space and time", f)
		}
	}
}

func TestEveryEngineSaysHowMuchRoomItsOwnPythonEnvironmentNeeds(t *testing.T) {
	c, err := Load()
	if err != nil {
		t.Fatal(err)
	}
	for _, e := range c.Engines {
		// engines run directly on the machine by default: the page shows this size instead of the image's
		if e.HostBytes < 1e9 || e.HostBytes >= e.ImageBytes {
			t.Errorf("%s: host_bytes %d should be a few GB, far below the image (%d)", e.ID, e.HostBytes, e.ImageBytes)
		}
	}
}

func TestModelHashesLookRight(t *testing.T) {
	cat, err := Load()
	if err != nil {
		t.Fatal(err)
	}
	big := 0
	for _, c := range cat.Capabilities {
		for _, m := range c.Models {
			if m.SHA256 != "" && !regexp.MustCompile(`^[0-9a-f]{64}$`).MatchString(m.SHA256) {
				t.Errorf("%s: the sha256 is not 64 hex digits", m.Dest)
			}
			if m.Bytes > 100<<20 && m.SHA256 == "" {
				t.Errorf("%s is %d bytes and has no sha256: a damaged copy of it would pass the size check", m.Dest, m.Bytes)
			}
			if m.SHA256 != "" {
				big++
			}
		}
	}
	if big == 0 {
		t.Error("no model has a sha256")
	}
}