package modules

import (
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"strings"
)

// A story is several scenes in one video. The model makes about five seconds of one scene well and loses the picture when
// it is asked for much more, so each scene is a clip of its own: scene 1 starts from nothing, every later scene starts from
// the last picture of the one before it, and the clips are joined into one video. All of it is one ComfyUI graph built
// here, so the engine does it in one job and the models stay loaded between the scenes.

const (
	SceneMinSeconds   = 1.0
	SceneMaxSeconds   = 8.0
	MaxSceneText      = 1500
	StoryMaxScenes    = 12
	storyFPS          = 24
	storyBudgetPixels = 60 * 832 * 480 // seconds x width x height a story may hold (the joined frames are kept in memory)
)

// FramesFor is the number of frames the model wants for this many seconds: 4n+1 at 24 per second.
func FramesFor(seconds float64) int64 {
	n := int64(math.Round(seconds * storyFPS / 4))
	if n < 1 {
		n = 1
	}
	return 4*n + 1
}

// Scene is one validated scene.
type Scene struct {
	Prompt  string
	Seconds float64
}

// validateScenes checks the "scenes" input: a list of {prompt, seconds}.
func validateScenes(f Field, v any) ([]any, error) {
	list, ok := v.([]any)
	if !ok {
		return nil, fmt.Errorf("%q must be a list of scenes", f.Name)
	}
	min, max := 1, StoryMaxScenes
	if f.Min != nil {
		min = int(*f.Min)
	}
	if f.Max != nil {
		max = int(*f.Max)
	}
	if len(list) < min || len(list) > max {
		return nil, fmt.Errorf("a story needs between %d and %d scenes, not %d", min, max, len(list))
	}
	out := make([]any, 0, len(list))
	for i, e := range list {
		m, ok := e.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("scene %d must have a prompt and a length", i+1)
		}
		for k := range m {
			if k != "prompt" && k != "seconds" {
				return nil, fmt.Errorf("scene %d: unknown field %q", i+1, k)
			}
		}
		p, _ := m["prompt"].(string)
		p = strings.TrimSpace(p)
		if p == "" {
			return nil, fmt.Errorf("scene %d has no prompt", i+1)
		}
		if len([]rune(p)) > MaxSceneText {
			return nil, fmt.Errorf("scene %d: the prompt is longer than %d characters", i+1, MaxSceneText)
		}
		sec := 5.0
		if raw, has := m["seconds"]; has && raw != nil {
			n, ok := num(raw)
			if !ok || math.IsNaN(n) || math.IsInf(n, 0) {
				return nil, fmt.Errorf("scene %d: the length must be a number of seconds", i+1)
			}
			sec = n
		}
		if sec < SceneMinSeconds || sec > SceneMaxSeconds {
			return nil, fmt.Errorf("scene %d: the length must be between %v and %v seconds", i+1, SceneMinSeconds, SceneMaxSeconds)
		}
		out = append(out, map[string]any{"prompt": p, "seconds": sec})
	}
	return out, nil
}

// storyScenes reads the validated scenes back.
func storyScenes(v any) []Scene {
	var out []Scene
	list, _ := v.([]any)
	for _, e := range list {
		m, _ := e.(map[string]any)
		p, _ := m["prompt"].(string)
		s, _ := num(m["seconds"])
		out = append(out, Scene{p, s})
	}
	return out
}

// StorySeconds is the length of the whole story.
func StorySeconds(scenes any) float64 {
	var t float64
	for _, s := range storyScenes(scenes) {
		t += s.Seconds
	}
	return t
}

// checkStoryBudget refuses a story whose joined frames would not fit in memory: the bigger the picture, the shorter.
func checkStoryBudget(scenes any, size string) error {
	w, h, err := splitSize(size)
	if err != nil {
		return err
	}
	if sec := StorySeconds(scenes); sec*float64(w*h) > storyBudgetPixels {
		max := float64(storyBudgetPixels) / float64(w*h)
		return fmt.Errorf("a story of %.0f seconds at %dx%d is too big for the memory of the machine: at this size the scenes together may be at most %.0f seconds", sec, w, h, math.Floor(max))
	}
	return nil
}

// StoryMemBytes estimates the memory the joined frames need on top of the models: every frame as floats, kept a few times
// over while the clips are joined.
func StoryMemBytes(scenes any, size string) int64 {
	w, h, err := splitSize(size)
	if err != nil {
		return 0
	}
	return int64(StorySeconds(scenes) * storyFPS * float64(w*h) * 3 * 4 * 5)
}

// BuildStory builds the ComfyUI graph of a story. vals are the values of Capability.Values (scenes, style, width, height,
// seed, prefix). The models and the negative prompt come from the text-to-video workflow, so the two never drift apart.
func BuildStory(vals map[string]any) (map[string]any, error) {
	scenes := storyScenes(vals["scenes"])
	if len(scenes) == 0 {
		return nil, fmt.Errorf("a story needs at least one scene")
	}
	width, _ := num(vals["width"])
	height, _ := num(vals["height"])
	seed, _ := num(vals["seed"])
	style, _ := vals["style"].(string)
	prefix, _ := vals["prefix"].(string)
	if width < 64 || height < 64 || prefix == "" {
		return nil, fmt.Errorf("the story needs a size and an output name")
	}

	raw, err := Workflow("video.generate.text.json")
	if err != nil {
		return nil, err
	}
	var base map[string]map[string]any
	if err := json.Unmarshal(raw, &base); err != nil {
		return nil, err
	}
	wf := map[string]any{}
	for _, id := range []string{"1", "2", "3", "4", "6"} { // the loaders, the sampling shift and the negative prompt
		n, ok := base[id]
		if !ok {
			return nil, fmt.Errorf("the text-to-video workflow changed: node %s is missing", id)
		}
		wf[id] = n
	}
	node := func(id, class string, in map[string]any) { wf[id] = map[string]any{"class_type": class, "inputs": in} }
	ref := func(id string) []any { return []any{id, 0} }

	var joined, last []any
	for i, sc := range scenes {
		b := 10 + 10*i
		id := func(k int) string { return strconv.Itoa(b + k) }
		length := FramesFor(sc.Seconds)
		text := sc.Prompt
		if style = strings.TrimSpace(style); style != "" {
			text += "\n" + style
		}
		node(id(0), "CLIPTextEncode", map[string]any{"text": text, "clip": []any{"2", 0}})
		lat := map[string]any{"vae": []any{"3", 0}, "width": int64(width), "height": int64(height), "length": length, "batch_size": int64(1)}
		if i > 0 {
			lat["start_image"] = last // the last picture of the scene before: the story goes on from it
		}
		node(id(1), "Wan22ImageToVideoLatent", lat)
		node(id(2), "KSampler", map[string]any{"model": []any{"4", 0}, "positive": ref(id(0)), "negative": []any{"6", 0}, "latent_image": ref(id(1)),
			"seed": int64(seed) + int64(i), "steps": int64(20), "cfg": int64(5), "sampler_name": "uni_pc", "scheduler": "simple", "denoise": int64(1)})
		node(id(3), "VAEDecodeTiled", map[string]any{"samples": ref(id(2)), "vae": []any{"3", 0}, "tile_size": int64(512), "overlap": int64(64), "temporal_size": int64(64), "temporal_overlap": int64(8)})
		if i < len(scenes)-1 {
			node(id(4), "ImageFromBatch", map[string]any{"image": ref(id(3)), "batch_index": int64(-1), "length": int64(1)})
			last = ref(id(4))
		}
		clip := ref(id(3))
		if i > 0 { // the first picture of a later scene is the last one of the scene before: it is not shown twice
			node(id(5), "ImageFromBatch", map[string]any{"image": ref(id(3)), "batch_index": int64(1), "length": length - 1})
			clip = ref(id(5))
		}
		if joined == nil {
			joined = clip
		} else {
			node(id(6), "ImageBatch", map[string]any{"image1": joined, "image2": clip})
			joined = ref(id(6))
		}
	}
	node("900", "CreateVideo", map[string]any{"images": joined, "fps": int64(storyFPS)})
	node("901", "SaveVideo", map[string]any{"video": ref("900"), "filename_prefix": prefix, "format": "auto", "codec": "auto"})
	return wf, nil
}