package prediction

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"math/big"
)

func BuildPrompt(workflowJSON string, binding InputBinding, input map[string]any) (map[string]any, error) {
	var workflow map[string]any
	if err := json.Unmarshal([]byte(workflowJSON), &workflow); err != nil {
		return nil, fmt.Errorf("invalid workflow: %w", err)
	}
	prompt := deepCopyMap(workflow)
	for name, target := range binding {
		if target == nil || target.Node == "" || target.Field == "" {
			continue
		}
		val, ok := input[name]
		if !ok {
			continue
		}
		if name == "seed" {
			val = resolveSeed(val)
		}
		node, ok := prompt[target.Node].(map[string]any)
		if !ok {
			return nil, fmt.Errorf("node %s not found", target.Node)
		}
		inputs, ok := node["inputs"].(map[string]any)
		if !ok {
			return nil, fmt.Errorf("node %s inputs missing", target.Node)
		}
		inputs[target.Field] = val
	}
	sanitizeSamplerSeeds(prompt)
	return prompt, nil
}

// resolveSeed maps seed < 0 (random) to a non-negative value ComfyUI accepts.
func resolveSeed(val any) any {
	n, ok := seedAsInt64(val)
	if !ok || n >= 0 {
		return val
	}
	r, err := rand.Int(rand.Reader, big.NewInt(1<<31-1))
	if err != nil {
		return int(0)
	}
	return int(r.Int64())
}

func seedAsInt64(v any) (int64, bool) {
	switch n := v.(type) {
	case int:
		return int64(n), true
	case int64:
		return n, true
	case float64:
		return int64(n), true
	case json.Number:
		i, err := n.Int64()
		return i, err == nil
	default:
		return 0, false
	}
}

func sanitizeSamplerSeeds(prompt map[string]any) {
	seedFields := []string{"seed", "noise_seed"}
	for _, raw := range prompt {
		node, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		if class, _ := node["class_type"].(string); class != "KSampler" {
			continue
		}
		inputs, ok := node["inputs"].(map[string]any)
		if !ok {
			continue
		}
		for _, field := range seedFields {
			if v, ok := inputs[field]; ok {
				inputs[field] = resolveSeed(v)
			}
		}
	}
}

func deepCopyMap(src map[string]any) map[string]any {
	b, _ := json.Marshal(src)
	var dst map[string]any
	_ = json.Unmarshal(b, &dst)
	return dst
}