Luminary/internal/prediction/workflow_analyze.go

package prediction

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"strings"

	"github.com/rose_cat707/luminary/internal/model"
)

type WorkflowNode struct {
	ID               string   `json:"id"`
	ClassType        string   `json:"class_type"`
	Title            string   `json:"title"`
	InjectableFields []string `json:"injectable_fields"`
}

type BindingCandidate struct {
	Node  string `json:"node"`
	Field string `json:"field"`
}

type BindingSuggestion struct {
	Param               string             `json:"param"`
	Node                string             `json:"node,omitempty"`
	Field               string             `json:"field,omitempty"`
	Type                ParamType          `json:"type"`
	DefaultFromWorkflow any                `json:"default_from_workflow,omitempty"`
	Confidence          string             `json:"confidence"`
	Candidates          []BindingCandidate `json:"candidates"`
}

type AnalyzeResult struct {
	VersionHashPreview string              `json:"version_hash_preview"`
	Nodes              []WorkflowNode      `json:"nodes"`
	Suggestions        []BindingSuggestion `json:"suggestions"`
}

func HashWorkflowJSON(raw string) string {
	sum := sha256.Sum256([]byte(raw))
	return hex.EncodeToString(sum[:])
}

func AnalyzeWorkflow(workflowJSON string, wt model.WorkflowType) (*AnalyzeResult, error) {
	var workflow map[string]any
	if err := json.Unmarshal([]byte(workflowJSON), &workflow); err != nil {
		return nil, fmt.Errorf("invalid workflow json: %w", err)
	}
	nodes := parseNodes(workflow)
	suggestions := suggestBindings(workflow, nodes, wt)
	return &AnalyzeResult{
		VersionHashPreview: HashWorkflowJSON(workflowJSON),
		Nodes:              nodes,
		Suggestions:        suggestions,
	}, nil
}

func parseNodes(workflow map[string]any) []WorkflowNode {
	var out []WorkflowNode
	for id, raw := range workflow {
		node, ok := raw.(map[string]any)
		if !ok {
			continue
		}
		classType, _ := node["class_type"].(string)
		title := ""
		if meta, ok := node["_meta"].(map[string]any); ok {
			title, _ = meta["title"].(string)
		}
		var fields []string
		if inputs, ok := node["inputs"].(map[string]any); ok {
			for k, v := range inputs {
				if isInjectableValue(v) {
					fields = append(fields, k)
				}
			}
		}
		out = append(out, WorkflowNode{
			ID: id, ClassType: classType, Title: title, InjectableFields: fields,
		})
	}
	return out
}

func isInjectableValue(v any) bool {
	switch v.(type) {
	case string, float64, bool, int, int64:
		return true
	default:
		return false
	}
}

func suggestBindings(workflow map[string]any, nodes []WorkflowNode, wt model.WorkflowType) []BindingSuggestion {
	defs := ParamsForWorkflowType(wt)
	var clipNodes []WorkflowNode
	var samplerNodes []WorkflowNode
	var latentNodes []WorkflowNode
	var loadImageNodes []WorkflowNode
	for _, n := range nodes {
		switch n.ClassType {
		case "CLIPTextEncode":
			clipNodes = append(clipNodes, n)
		case "KSampler", "KSamplerAdvanced", "RandomNoise":
			samplerNodes = append(samplerNodes, n)
		case "EmptyLatentImage":
			latentNodes = append(latentNodes, n)
		case "LoadImage":
			loadImageNodes = append(loadImageNodes, n)
		}
	}
	var out []BindingSuggestion
	for _, d := range defs {
		s := BindingSuggestion{Param: d.Name, Type: d.Type, Confidence: "low"}
		switch d.Name {
		case "prompt":
			s = suggestCLIP(clipNodes, workflow, false)
		case "negative_prompt":
			s = suggestCLIP(clipNodes, workflow, true)
		case "seed":
			s = suggestSamplerField(samplerNodes, workflow, "seed", d.Type)
			if s.Node == "" {
				s = suggestSamplerField(samplerNodes, workflow, "noise_seed", d.Type)
			}
		case "steps":
			s = suggestSamplerField(samplerNodes, workflow, "steps", d.Type)
		case "cfg_scale":
			s = suggestSamplerField(samplerNodes, workflow, "cfg", d.Type)
		case "denoise":
			s = suggestSamplerField(samplerNodes, workflow, "denoise", d.Type)
		case "width":
			s = suggestLatentField(latentNodes, workflow, "width", d.Type)
		case "height":
			s = suggestLatentField(latentNodes, workflow, "height", d.Type)
		case "image":
			s = suggestLoadImage(loadImageNodes, workflow, d.Type)
		default:
			s.Param = d.Name
			s.Type = d.Type
		}
		if s.Param == "" {
			s.Param = d.Name
			s.Type = d.Type
		}
		if s.Node != "" {
			s.Confidence = "high"
		}
		out = append(out, s)
	}
	return out
}

func suggestCLIP(nodes []WorkflowNode, workflow map[string]any, negative bool) BindingSuggestion {
	s := BindingSuggestion{Param: "prompt", Type: ParamString, Confidence: "low"}
	if negative {
		s.Param = "negative_prompt"
	}
	var candidates []BindingCandidate
	var picked *WorkflowNode
	for i := range nodes {
		n := nodes[i]
		titleLower := strings.ToLower(n.Title)
		isNeg := strings.Contains(titleLower, "negative") || strings.Contains(titleLower, "neg")
		if !hasField(n, "text") {
			continue
		}
		candidates = append(candidates, BindingCandidate{Node: n.ID, Field: "text"})
		if negative && isNeg && picked == nil {
			picked = &nodes[i]
		}
		if !negative && !isNeg && picked == nil {
			picked = &nodes[i]
		}
	}
	if picked == nil && len(candidates) > 0 {
		idx := 0
		if negative && len(candidates) > 1 {
			idx = 1
		}
		for i := range nodes {
			if nodes[i].ID == candidates[idx].Node {
				picked = &nodes[i]
				break
			}
		}
	}
	s.Candidates = candidates
	if picked != nil {
		s.Node = picked.ID
		s.Field = "text"
		s.DefaultFromWorkflow = nodeFieldValue(workflow, picked.ID, "text")
	}
	return s
}

func suggestSamplerField(nodes []WorkflowNode, workflow map[string]any, field string, pt ParamType) BindingSuggestion {
	paramName := field
	if field == "cfg" {
		paramName = "cfg_scale"
	}
	s := BindingSuggestion{Param: paramName, Type: pt, Confidence: "low"}
	var candidates []BindingCandidate
	for _, n := range nodes {
		if hasField(n, field) {
			candidates = append(candidates, BindingCandidate{Node: n.ID, Field: field})
			if s.Node == "" {
				s.Node = n.ID
				s.Field = field
				s.DefaultFromWorkflow = nodeFieldValue(workflow, n.ID, field)
			}
		}
	}
	s.Candidates = candidates
	return s
}

func suggestLatentField(nodes []WorkflowNode, workflow map[string]any, field string, pt ParamType) BindingSuggestion {
	s := BindingSuggestion{Param: field, Type: pt, Confidence: "low"}
	for _, n := range nodes {
		if hasField(n, field) {
			s.Candidates = append(s.Candidates, BindingCandidate{Node: n.ID, Field: field})
			if s.Node == "" {
				s.Node = n.ID
				s.Field = field
				s.DefaultFromWorkflow = nodeFieldValue(workflow, n.ID, field)
			}
		}
	}
	return s
}

func suggestLoadImage(nodes []WorkflowNode, workflow map[string]any, pt ParamType) BindingSuggestion {
	s := BindingSuggestion{Param: "image", Type: pt, Confidence: "low"}
	for _, n := range nodes {
		if hasField(n, "image") {
			s.Candidates = append(s.Candidates, BindingCandidate{Node: n.ID, Field: "image"})
			if s.Node == "" {
				s.Node = n.ID
				s.Field = "image"
				s.DefaultFromWorkflow = nodeFieldValue(workflow, n.ID, "image")
			}
		}
	}
	return s
}

func hasField(n WorkflowNode, field string) bool {
	for _, f := range n.InjectableFields {
		if f == field {
			return true
		}
	}
	return false
}

func nodeFieldValue(workflow map[string]any, nodeID, field string) any {
	node, ok := workflow[nodeID].(map[string]any)
	if !ok {
		return nil
	}
	inputs, ok := node["inputs"].(map[string]any)
	if !ok {
		return nil
	}
	return inputs[field]
}