package modules

import (
	"encoding/json"
	"strconv"
	"strings"
	"testing"
)

func storyInput(scenes ...map[string]any) map[string]any {
	var list []any
	for _, s := range scenes {
		list = append(list, s)
	}
	return map[string]any{"operation": "story", "scenes": list}
}

func scene(p string, sec float64) map[string]any { return map[string]any{"prompt": p, "seconds": sec} }

func TestStoryIsAnOperationOfVideoGenerate(t *testing.T) {
	c := load(t)
	vg, _ := c.Capability("video.generate")
	if vg.Workflows["story"] != BuiltinStory {
		t.Fatalf("workflows: %v", vg.Workflows)
	}
	shown := map[string][]string{}
	for _, f := range vg.Inputs {
		shown[f.Name] = f.ShowWhen
	}
	if strings.Join(shown["scenes"], ",") != "story" || strings.Join(shown["style"], ",") != "story" {
		t.Errorf("scenes and style belong to the story only: %v", shown)
	}
	if strings.Join(shown["prompt"], ",") != "text,image" || strings.Join(shown["seconds"], ",") != "text,image" {
		t.Errorf("a story has no single prompt or length: %v", shown)
	}
	if shown["size"] != nil || shown["seed"] != nil {
		t.Errorf("size and seed apply to every mode: %v", shown)
	}
}

func TestFramesForIsFourNPlusOne(t *testing.T) {
	for sec, want := range map[float64]int64{1: 25, 3: 73, 5: 121, 8: 193, 0.1: 5} {
		if got := FramesFor(sec); got != want {
			t.Errorf("%v s = %d frames, want %d", sec, got, want)
		}
	}
}

func TestStoryValidation(t *testing.T) {
	c := load(t)
	vg, _ := c.Capability("video.generate")
	out, err := vg.Validate(storyInput(scene("  a hallway  ", 4), map[string]any{"prompt": "a door"}))
	if err != nil {
		t.Fatal(err)
	}
	sc := out["scenes"].([]any)
	if len(sc) != 2 || sc[0].(map[string]any)["prompt"] != "a hallway" || sc[1].(map[string]any)["seconds"] != 5.0 {
		t.Errorf("a prompt is trimmed and a missing length is five seconds: %v", sc)
	}
	if _, ok := out["prompt"]; ok {
		t.Error("a story has no single prompt")
	}
	if out["size"] != "832x480" {
		t.Errorf("size default: %v", out["size"])
	}
	var many []map[string]any
	for i := 0; i < 13; i++ {
		many = append(many, scene("x", 1))
	}
	long := strings.Repeat("a", MaxSceneText+1)
	for name, in := range map[string]map[string]any{
		"no scenes":            {"operation": "story"},
		"an empty list":        {"operation": "story", "scenes": []any{}},
		"too many scenes":      storyInput(many...),
		"an empty prompt":      storyInput(scene("  ", 3)),
		"a scene too short":    storyInput(scene("x", 0.5)),
		"a scene too long":     storyInput(scene("x", 9)),
		"a text that is long":  storyInput(scene(long, 3)),
		"an unknown key":       storyInput(map[string]any{"prompt": "x", "seconds": 3, "image": "a.png"}),
		"a scene that is text": {"operation": "story", "scenes": []any{"just text"}},
		"scenes that are text": {"operation": "story", "scenes": "a; b"},
		"a length that is not": storyInput(map[string]any{"prompt": "x", "seconds": "abc"}),
	} {
		if _, err := vg.Validate(in); err == nil {
			t.Errorf("%s was accepted", name)
		}
	}
}

func TestAStoryMustFitInMemoryAtItsSize(t *testing.T) {
	c := load(t)
	vg, _ := c.Capability("video.generate")
	var scenes []map[string]any
	for i := 0; i < 8; i++ { // 8 x 7.5 = 60 seconds
		scenes = append(scenes, scene("x", 7.5))
	}
	in := storyInput(scenes...)
	if _, err := vg.Validate(in); err != nil {
		t.Fatalf("60 seconds at 832x480 must be allowed: %v", err)
	}
	in["size"] = "1280x704"
	if _, err := vg.Validate(in); err == nil || !strings.Contains(err.Error(), "at most") {
		t.Errorf("60 seconds at 1280x704 must be refused with the limit: %v", err)
	}
	in["scenes"] = storyInput(scene("x", 8), scene("x", 8), scene("x", 8))["scenes"]
	if _, err := vg.Validate(in); err != nil {
		t.Errorf("24 seconds at 1280x704 fits: %v", err)
	}
	if StoryMemBytes(in["scenes"], "1280x704") <= StoryMemBytes(in["scenes"], "832x480") {
		t.Error("a bigger picture needs more memory")
	}
}

// ---- the graph ----

func buildStory(t *testing.T, style string, scenes ...map[string]any) map[string]any {
	t.Helper()
	c := load(t)
	vg, _ := c.Capability("video.generate")
	in := storyInput(scenes...)
	if style != "" {
		in["style"] = style
	}
	ok, err := vg.Validate(in)
	if err != nil {
		t.Fatal(err)
	}
	vals, err := vg.Values(ok)
	if err != nil {
		t.Fatal(err)
	}
	vals["prefix"] = "wxch/j/out"
	vals["seed"] = int64(1000)
	wf, err := BuildStory(vals)
	if err != nil {
		t.Fatal(err)
	}
	return wf
}

func inputsOf(wf map[string]any, id string) map[string]any {
	return wf[id].(map[string]any)["inputs"].(map[string]any)
}

// every reference points at a node that exists, nothing is left as a placeholder, and the graph has no cycle
func checkGraph(t *testing.T, wf map[string]any) {
	t.Helper()
	deps := map[string][]string{}
	for id, n := range wf {
		for k, v := range n.(map[string]any)["inputs"].(map[string]any) {
			if r, ok := v.([]any); ok && len(r) == 2 {
				to := r[0].(string)
				if _, exists := wf[to]; !exists {
					t.Errorf("node %s input %s refers to the missing node %s", id, k, to)
				}
				deps[id] = append(deps[id], to)
			}
			if s, ok := v.(string); ok && strings.HasPrefix(s, "@") {
				t.Errorf("node %s input %s still holds a placeholder %s", id, k, s)
			}
		}
	}
	state := map[string]int{}
	var visit func(string) bool
	visit = func(id string) bool {
		if state[id] == 1 {
			return false
		}
		if state[id] == 2 {
			return true
		}
		state[id] = 1
		for _, d := range deps[id] {
			if !visit(d) {
				return false
			}
		}
		state[id] = 2
		return true
	}
	for id := range wf {
		if !visit(id) {
			t.Fatal("the graph has a cycle")
		}
	}
}

func TestStoryGraphChainsTheScenes(t *testing.T) {
	wf := buildStory(t, "pastel colours", scene("a hallway", 4), scene("a bedroom", 5), scene("a garden", 3))
	checkGraph(t, wf)
	if _, err := json.Marshal(wf); err != nil {
		t.Fatal(err)
	}
	count := map[string]int{}
	for _, n := range wf {
		count[n.(map[string]any)["class_type"].(string)]++
	}
	want := map[string]int{"UNETLoader": 1, "CLIPLoader": 1, "VAELoader": 1, "ModelSamplingSD3": 1, "KSampler": 3, "Wan22ImageToVideoLatent": 3,
		"VAEDecodeTiled": 3, "CLIPTextEncode": 4, "ImageFromBatch": 4, "ImageBatch": 2, "CreateVideo": 1, "SaveVideo": 1}
	for k, v := range want {
		if count[k] != v {
			t.Errorf("%s: %d nodes, want %d (%v)", k, count[k], v, count)
		}
	}
	prompts := []string{"a hallway", "a bedroom", "a garden"}
	for i, sec := range []float64{4, 5, 3} {
		b := 10 + 10*i
		id := func(k int) string { return strconv.Itoa(b + k) }
		lat := inputsOf(wf, id(1))
		if lat["length"] != FramesFor(sec) || lat["width"] != int64(832) || lat["height"] != int64(480) {
			t.Errorf("scene %d latent: %v", i+1, lat)
		}
		if i == 0 {
			if _, has := lat["start_image"]; has {
				t.Error("the first scene starts from nothing")
			}
		} else {
			pb := 10 + 10*(i-1)
			start := lat["start_image"].([]any)
			lastFrame := inputsOf(wf, strconv.Itoa(pb+4))
			if start[0] != strconv.Itoa(pb+4) || lastFrame["batch_index"] != int64(-1) || lastFrame["length"] != int64(1) || lastFrame["image"].([]any)[0] != strconv.Itoa(pb+3) {
				t.Errorf("scene %d must start from the last picture of scene %d: %v %v", i+1, i, start, lastFrame)
			}
			trim := inputsOf(wf, id(5))
			if trim["batch_index"] != int64(1) || trim["length"] != FramesFor(sec)-1 {
				t.Errorf("scene %d must drop its first picture (the one scene %d ended on): %v", i+1, i, trim)
			}
		}
		ks := inputsOf(wf, id(2))
		if ks["seed"] != int64(1000+i) || ks["steps"] != int64(20) || ks["sampler_name"] != "uni_pc" || ks["cfg"] != int64(5) {
			t.Errorf("scene %d sampler: %v", i+1, ks)
		}
		txt := inputsOf(wf, id(0))["text"].(string)
		if !strings.Contains(txt, "pastel colours") || !strings.HasPrefix(txt, prompts[i]) {
			t.Errorf("scene %d text: %q", i+1, txt)
		}
		if ks["negative"].([]any)[0] != "6" {
			t.Error("every scene uses the same negative prompt")
		}
	}
	if _, has := wf["34"]; has {
		t.Error("the last scene made a picture nobody uses")
	}
	// the clips are joined in order and saved
	j1, j2 := inputsOf(wf, "26"), inputsOf(wf, "36")
	if j1["image1"].([]any)[0] != "13" || j1["image2"].([]any)[0] != "25" || j2["image1"].([]any)[0] != "26" || j2["image2"].([]any)[0] != "35" {
		t.Errorf("join: %v %v", j1, j2)
	}
	if inputsOf(wf, "900")["images"].([]any)[0] != "36" || inputsOf(wf, "900")["fps"] != int64(24) {
		t.Errorf("video: %v", inputsOf(wf, "900"))
	}
	if sv := inputsOf(wf, "901"); sv["filename_prefix"] != "wxch/j/out" || sv["video"].([]any)[0] != "900" {
		t.Errorf("save: %v", sv)
	}
	// the models and the negative prompt are the ones of the text-to-video workflow
	raw, _ := Workflow("video.generate.text.json")
	var base map[string]map[string]any
	json.Unmarshal(raw, &base)
	for _, id := range []string{"1", "2", "3", "4", "6"} {
		a, _ := json.Marshal(base[id])
		b, _ := json.Marshal(wf[id])
		if string(a) != string(b) {
			t.Errorf("node %s differs from the text-to-video workflow", id)
		}
	}
}

func TestOneSceneStoryNeedsNoJoining(t *testing.T) {
	wf := buildStory(t, "", scene("a lone scene", 5))
	checkGraph(t, wf)
	for id, n := range wf {
		if c := n.(map[string]any)["class_type"]; c == "ImageFromBatch" || c == "ImageBatch" {
			t.Errorf("node %s (%v) is not needed for one scene", id, c)
		}
	}
	if inputsOf(wf, "900")["images"].([]any)[0] != "13" {
		t.Errorf("the video is the decoded scene: %v", inputsOf(wf, "900"))
	}
	if strings.Contains(inputsOf(wf, "10")["text"].(string), "\n") {
		t.Error("no style, no extra line")
	}
}

func TestTotalFramesOfAStory(t *testing.T) {
	wf := buildStory(t, "", scene("a", 4), scene("b", 5), scene("c", 3))
	// every scene but the first gives up one picture
	got := FramesFor(4)
	for _, id := range []string{"25", "35"} {
		got += inputsOf(wf, id)["length"].(int64)
	}
	if want := FramesFor(4) + FramesFor(5) - 1 + FramesFor(3) - 1; got != want || want != 289 {
		t.Errorf("frames %d, want %d", got, want)
	}
}

func TestStoryBuildRefusesWhatItCannotBuild(t *testing.T) {
	one := []any{map[string]any{"prompt": "x", "seconds": 3.0}}
	for name, vals := range map[string]map[string]any{
		"no scenes": {"width": int64(832), "height": int64(480), "prefix": "p"},
		"no size":   {"scenes": one, "prefix": "p"},
		"no prefix": {"scenes": one, "width": int64(832), "height": int64(480)},
	} {
		if _, err := BuildStory(vals); err == nil {
			t.Errorf("%s was built", name)
		}
	}
}