package modules

import (
	"crypto/rand"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"math"
	"strconv"
	"strings"

	"wxch.org/wxch/catalog"
)

const maxText = 4000

// Operation is the chosen "operation" (video.edit has several); "default" when the capability has none.
func (c Capability) Operation(in map[string]any) string {
	for _, f := range c.Inputs {
		if f.Name == "operation" {
			if s, ok := in["operation"].(string); ok && s != "" {
				return s
			}
			if s, ok := f.Default.(string); ok {
				return s
			}
		}
	}
	return "default"
}

func (f Field) shownFor(op string) bool {
	if len(f.ShowWhen) == 0 {
		return true
	}
	for _, o := range f.ShowWhen {
		if o == op {
			return true
		}
	}
	return false
}

func num(v any) (float64, bool) {
	switch x := v.(type) {
	case float64:
		return x, true
	case int:
		return float64(x), true
	case int64:
		return float64(x), true
	case json.Number:
		f, err := x.Float64()
		return f, err == nil
	case string:
		f, err := strconv.ParseFloat(strings.TrimSpace(x), 64)
		return f, err == nil
	}
	return 0, false
}

// Validate checks what a person asked for against the capability's fields. Unknown names are refused, a missing value
// gets its default, numbers must lie within their limits. A file field must carry the stored file's name (the caller
// stores the files); that is the only thing checked about it here.
func (c Capability) Validate(in map[string]any) (map[string]any, error) {
	known := map[string]bool{}
	for _, f := range c.Inputs {
		known[f.Name] = true
	}
	for k := range in {
		if !known[k] {
			return nil, fmt.Errorf("unknown input %q", k)
		}
	}
	op := c.Operation(in)
	if len(c.Workflows) > 0 {
		if _, ok := c.Workflows[op]; !ok {
			if _, ok := c.Workflows["default"]; !ok {
				return nil, fmt.Errorf("operation %q is not available", op)
			}
		}
	}
	out := map[string]any{}
	for _, f := range c.Inputs {
		if !f.shownFor(op) {
			continue
		}
		v, has := in[f.Name]
		if s, ok := v.(string); ok && strings.TrimSpace(s) == "" {
			has = false
		}
		if !has || v == nil {
			if f.Required {
				return nil, fmt.Errorf("%q is required", f.Name)
			}
			if f.Default != nil {
				v = f.Default
			} else {
				continue
			}
		}
		switch f.Type {
		case "text":
			s, ok := v.(string)
			if !ok {
				return nil, fmt.Errorf("%q must be text", f.Name)
			}
			s = strings.TrimSpace(s)
			if len([]rune(s)) > maxText {
				return nil, fmt.Errorf("%q is longer than %d characters", f.Name, maxText)
			}
			out[f.Name] = s
		case "int", "float":
			n, ok := num(v)
			if !ok || math.IsNaN(n) || math.IsInf(n, 0) {
				return nil, fmt.Errorf("%q must be a number", f.Name)
			}
			if f.Type == "int" {
				if n != math.Trunc(n) {
					return nil, fmt.Errorf("%q must be a whole number", f.Name)
				}
			}
			if (f.Min != nil && n < *f.Min) || (f.Max != nil && n > *f.Max) {
				return nil, fmt.Errorf("%q must be between %v and %v", f.Name, deref(f.Min), deref(f.Max))
			}
			if f.Type == "int" {
				out[f.Name] = int64(n)
			} else {
				out[f.Name] = n
			}
		case "bool":
			b, ok := v.(bool)
			if !ok {
				if s, isStr := v.(string); isStr && (s == "true" || s == "false") {
					b, ok = s == "true", true
				}
			}
			if !ok {
				return nil, fmt.Errorf("%q must be true or false", f.Name)
			}
			out[f.Name] = b
		case "scenes":
			sc, err := validateScenes(f, v)
			if err != nil {
				return nil, err
			}
			out[f.Name] = sc
		case "choice":
			s, ok := v.(string)
			if !ok {
				s = fmt.Sprint(v)
			}
			found := false
			for _, o := range f.Options {
				if o.Value == s {
					found = true
				}
			}
			if !found {
				return nil, fmt.Errorf("%q: %q is not one of the choices", f.Name, s)
			}
			out[f.Name] = s
		default: // a file
			s, ok := v.(string)
			if !ok || s == "" || strings.ContainsAny(s, "/\\") || strings.Contains(s, "..") || len(s) > 200 {
				return nil, fmt.Errorf("%q must be a file name", f.Name)
			}
			out[f.Name] = s
		}
	}
	if op == "story" {
		if err := checkStoryBudget(out["scenes"], fmt.Sprint(out["size"])); err != nil {
			return nil, err
		}
	}
	return out, nil
}

func deref(p *float64) any {
	if p == nil {
		return "any"
	}
	return *p
}

// FileFields names the file inputs this operation needs (the ones present in the validated input).
func (c Capability) FileFields(op string) []Field {
	var out []Field
	for _, f := range c.Inputs {
		if f.isFile() && f.shownFor(op) {
			out = append(out, f)
		}
	}
	return out
}

// BuiltinStory marks the "workflow" of the story operation: the graph is built in code (BuildStory), there is no file.
const BuiltinStory = "builtin:story"

// WorkflowFor returns the workflow file for the operation.
func (c Capability) WorkflowFor(op string) (string, error) {
	if wf, ok := c.Workflows[op]; ok {
		return wf, nil
	}
	if wf, ok := c.Workflows["default"]; ok {
		return wf, nil
	}
	return "", fmt.Errorf("no workflow for operation %q", op)
}

func randomSeed() int64 {
	var b [8]byte
	rand.Read(b[:])
	return int64(binary.BigEndian.Uint64(b[:]) & (1<<50 - 1)) // stays an exact number even where JSON uses floats
}

func splitSize(s string) (w, h int64, err error) {
	a, b, ok := strings.Cut(s, "x")
	if !ok {
		return 0, 0, fmt.Errorf("bad size %q", s)
	}
	w, e1 := strconv.ParseInt(a, 10, 64)
	h, e2 := strconv.ParseInt(b, 10, 64)
	if e1 != nil || e2 != nil || w < 64 || h < 64 || w > 4096 || h > 4096 {
		return 0, 0, fmt.Errorf("bad size %q", s)
	}
	return w, h, nil
}

// Values turns validated input into the values the workflow template asks for ("@name" entries): it keeps the input,
// adds a random seed when none was given, and derives what the workflow needs from friendlier fields (a size like
// 832x480 becomes width and height, seconds become a frame count).
func (c Capability) Values(in map[string]any) (map[string]any, error) {
	v := map[string]any{}
	for k, x := range in {
		v[k] = x
	}
	if s, ok := v["seed"].(int64); !ok || s < 0 {
		v["seed"] = randomSeed()
	}
	op := c.Operation(in)
	switch c.ID {
	case "image.generate":
		w, h, err := splitSize(fmt.Sprint(v["size"]))
		if err != nil {
			return nil, err
		}
		v["width"], v["height"] = w, h
		if _, ok := v["count"]; !ok {
			v["count"] = int64(1)
		}
	case "image.edit":
		if v["quality"] == "best" {
			v["steps"], v["cfg"], v["lora_strength"] = int64(40), 3.0, 0.0
		} else {
			v["steps"], v["cfg"], v["lora_strength"] = int64(4), 1.0, 1.0
		}
	case "video.generate":
		w, h, err := splitSize(fmt.Sprint(v["size"]))
		if err != nil {
			return nil, err
		}
		v["width"], v["height"] = w, h
		sec, _ := num(v["seconds"])
		v["length"], v["fps"] = FramesFor(sec), int64(storyFPS) // the model wants 4n+1 frames
	case "video.edit":
		switch op {
		case "upscale":
			// the model scales by 4; for x2 the result is halved afterwards
			if fmt.Sprint(v["scale"]) == "2" {
				v["down"] = 0.5
			} else {
				v["down"] = 1.0
			}
		case "interpolate":
			if _, ok := v["multiplier"]; !ok {
				v["multiplier"] = int64(2)
			}
		}
	}
	return v, nil
}

// Instantiate fills a workflow template: every string that is exactly "@name" is replaced by that value (keeping its
// type). Text from the user is only ever inserted as a value, never scanned for "@name" itself.
func Instantiate(wf []byte, vals map[string]any) (map[string]any, error) {
	var root map[string]any
	if err := json.Unmarshal(wf, &root); err != nil {
		return nil, err
	}
	var walk func(any) (any, error)
	walk = func(x any) (any, error) {
		switch t := x.(type) {
		case map[string]any:
			for k, e := range t {
				r, err := walk(e)
				if err != nil {
					return nil, err
				}
				t[k] = r
			}
			return t, nil
		case []any:
			for i, e := range t {
				r, err := walk(e)
				if err != nil {
					return nil, err
				}
				t[i] = r
			}
			return t, nil
		case string:
			if strings.HasPrefix(t, "@") && len(t) > 1 {
				v, ok := vals[t[1:]]
				if !ok {
					return nil, fmt.Errorf("the workflow needs %q, which was not given", t[1:])
				}
				return v, nil
			}
		}
		return x, nil
	}
	out, err := walk(root)
	if err != nil {
		return nil, err
	}
	return out.(map[string]any), nil
}

// Workflow reads an embedded workflow template.
func Workflow(file string) ([]byte, error) { return catalog.Workflows.ReadFile("workflows/" + file) }