safesight-control/internal/service/template_composer.go

package service

import (
    "encoding/json"
    "fmt"
    "sort"
    "strings"
)

// templateNodeType classifies a node for composition purposes.
type templateNodeType int

const (
    nodeSource      templateNodeType = iota // input_* or role=source
    nodePreprocess                          // type=preprocess
    nodeOSD                                 // type=osd
    nodePublish                             // type=publish or role=sink
    nodeAlarm                               // type=alarm
    nodeDetection                           // everything else (AI, tracker, logic_gate, etc.)
)

// templateDAG holds the parsed structure of a single template.
type templateDAG struct {
    name     string
    executor map[string]any
    nodes    []map[string]any // original node objects
    edges    []any            // original edge tuples [from, to, meta?]
    // classification
    nodeByID    map[string]map[string]any
    nodeClasses map[string]templateNodeType
    // adjacency
    successors map[string][]string
    predecessors map[string][]string
}

// ComposeTemplates merges multiple template RAWs into a single combined template.
//
// The first template in the list is treated as the "primary" — its infrastructure
// nodes (input, preprocess, publish, osd, alarm) are kept as-is. Detection nodes
// from all templates are combined, with supplementary templates' nodes renamed to
// avoid ID collisions.
//
// Returns the merged template as a map suitable for saving.
func ComposeTemplates(templates []map[string]any) (map[string]any, error) {
    if len(templates) == 0 {
        return nil, fmt.Errorf("至少需要一个模板")
    }
    if len(templates) == 1 {
        // Nothing to compose — return a deep copy.
        return deepCopyMap(templates[0]), nil
    }

    dags := make([]*templateDAG, 0, len(templates))
    for i, raw := range templates {
        dag, err := parseTemplateDAG(raw)
        if err != nil {
            return nil, fmt.Errorf("parse template %d: %w", i, err)
        }
        dags = append(dags, dag)
    }

    return composeDAGs(dags)
}

func parseTemplateDAG(raw map[string]any) (*templateDAG, error) {
    name := stringValue(raw["name"])
    if name == "" {
        return nil, fmt.Errorf("template name is required")
    }

    templateMap, _ := raw["template"].(map[string]any)
    if templateMap == nil {
        return nil, fmt.Errorf("template %q: missing template body", name)
    }

    executor, _ := templateMap["executor"].(map[string]any)

    rawNodes, ok := templateMap["nodes"].([]any)
    if !ok || len(rawNodes) == 0 {
        return nil, fmt.Errorf("template %q: nodes are required", name)
    }

    rawEdges, ok := templateMap["edges"].([]any)
    if !ok || len(rawEdges) == 0 {
        return nil, fmt.Errorf("template %q: edges are required", name)
    }

    dag := &templateDAG{
        name:         name,
        executor:     executor,
        nodeByID:     make(map[string]map[string]any, len(rawNodes)),
        nodeClasses:  make(map[string]templateNodeType, len(rawNodes)),
        successors:   make(map[string][]string),
        predecessors: make(map[string][]string),
    }

    for _, item := range rawNodes {
        node, _ := item.(map[string]any)
        if node == nil {
            return nil, fmt.Errorf("template %q: node entry must be an object", name)
        }
        id := stringValue(node["id"])
        if id == "" {
            return nil, fmt.Errorf("template %q: node id is required", name)
        }
        dag.nodes = append(dag.nodes, node)
        dag.nodeByID[id] = node
        dag.nodeClasses[id] = classifyNode(node)
    }

    for _, item := range rawEdges {
        edge, _ := item.([]any)
        if len(edge) < 2 {
            return nil, fmt.Errorf("template %q: edge must have at least 2 elements [from, to]", name)
        }
        from := stringValue(edge[0])
        to := stringValue(edge[1])
        if from == "" || to == "" {
            return nil, fmt.Errorf("template %q: edge from/to cannot be empty", name)
        }
        if _, ok := dag.nodeByID[from]; !ok {
            return nil, fmt.Errorf("template %q: edge references unknown node %q", name, from)
        }
        if _, ok := dag.nodeByID[to]; !ok {
            return nil, fmt.Errorf("template %q: edge references unknown node %q", name, to)
        }
        dag.edges = append(dag.edges, edge)
        dag.successors[from] = append(dag.successors[from], to)
        dag.predecessors[to] = append(dag.predecessors[to], from)
    }

    return dag, nil
}

func classifyNode(node map[string]any) templateNodeType {
    nodeType := stringValue(node["type"])
    role := stringValue(node["role"])

    if role == "source" || strings.HasPrefix(strings.ToLower(nodeType), "input_") {
        return nodeSource
    }
    if nodeType == "preprocess" {
        return nodePreprocess
    }
    if nodeType == "osd" {
        return nodeOSD
    }
    if nodeType == "publish" || (role == "sink" && nodeType == "publish") {
        return nodePublish
    }
    if nodeType == "alarm" {
        return nodeAlarm
    }
    return nodeDetection
}

// composeDAGs merges multiple DAGs into one.
func composeDAGs(dags []*templateDAG) (map[string]any, error) {
    primary := dags[0]
    supplemental := dags[1:]

    // Re-classify: only the first preprocess after input is "infrastructure".
    // Later preprocess nodes (e.g. prepare_publish) are pipeline nodes.
    for _, dag := range dags {
        fixPreprocessClassification(dag)
    }

    mergedNodes := make([]any, 0)
    mergedEdges := make([]any, 0)

    // Track which node IDs map to which renamed versions.
    // For the primary, IDs stay as-is except detection nodes get a prefix.
    // For supplemental templates, all kept nodes get a prefix.
    renamed := make(map[string]string) // origID → mergedID
    mergedIDs := make(map[string]bool) // prevents collisions

    // Reserve primary node IDs
    for id := range primary.nodeByID {
        mergedIDs[id] = true
    }

    // --- 1. Add ALL nodes from primary template ---
    for _, node := range primary.nodes {
        id := stringValue(node["id"])
        mergedNodes = append(mergedNodes, deepCopyMap(node))
        renamed[id] = id
    }

    // --- 2. Add detection-only nodes from supplemental templates ---
    for _, dag := range supplemental {
        prefix := sanitizeIDPrefix(dag.name) + "_"
        for _, node := range dag.nodes {
            id := stringValue(node["id"])
            if isInfraNode(dag.nodeClasses[id]) {
                continue // infra nodes come from primary only
            }
            newID := uniqueID(prefix+id, mergedIDs)
            mergedIDs[newID] = true
            renamed[dag.name+"::"+id] = newID

            cloned := deepCopyMap(node)
            cloned["id"] = newID
            mergedNodes = append(mergedNodes, cloned)
        }
    }

    // --- 3. Rewire edges ---
    // Primary edges: keep, but remap if target went through detection nodes.
    // We need to track the boundary between detection and infra.
    //
    // Approach: For the primary, keep all edges as-is.
    // For supplemental templates, we need to:
    //   a. Find the "entry" detection node (first detection after preprocess)
    //   b. Find the "exit" detection node (last detection before OSD/publish)
    //   c. Add edges: primary_preprocess → entry, exit → primary_osd

    // --- 3a. Primary edges: keep as-is ---
    for _, edge := range primary.edges {
        mergedEdges = append(mergedEdges, deepCopyEdge(edge))
    }

    // --- 3b. Supplemental template detection chains ---
    // Find the preprocess and osd nodes in the primary.
    primaryPrepID := findNodeByClass(primary, nodePreprocess)
    primaryOSDID := findNodeByClass(primary, nodeOSD)

    for _, dag := range supplemental {
        // Find detection chain in this dag
        entryID, exitID := findDetectionChainEnds(dag)

        if entryID == "" {
            // No detection nodes — nothing to add.
            continue
        }

        // Find the preprocess's successor in this dag (the first detection)
        prepID := findNodeByClass(dag, nodePreprocess)
        if prepID != "" {
            for _, succ := range dag.successors[prepID] {
                if dag.nodeClasses[succ] == nodeDetection {
                    entryID = succ
                    break
                }
            }
        }

        // Find the node before OSD/publish (the last detection)
        osdID := findNodeByClass(dag, nodeOSD)
        if osdID == "" {
            osdID = findNodeByClass(dag, nodePublish)
        }
        if osdID != "" && len(dag.predecessors[osdID]) > 0 {
            exitID = dag.predecessors[osdID][0]
        }

        if entryID == "" && exitID == "" {
            continue
        }

        // Build the detection chain from entry to exit.
        chain := buildDetectionChain(dag, entryID, exitID)
        if len(chain) == 0 {
            continue
        }

        // Connect: primary_preprocess → first detection in chain
        if primaryPrepID != "" {
            newEntry := renamed[dag.name+"::"+chain[0]]
            if newEntry == "" {
                newEntry = chain[0] // primary detection node, keep as-is
            }
            // Only add if not already connected
            if !hasEdge(mergedEdges, primaryPrepID, newEntry) {
                mergedEdges = append(mergedEdges, newEdge(primaryPrepID, newEntry))
            }
        }

        // Wire internal chain edges
        for i := 0; i < len(chain)-1; i++ {
            a := nodeIDInMerged(chain[i], dag.name, renamed)
            b := nodeIDInMerged(chain[i+1], dag.name, renamed)
            if a != "" && b != "" && !hasEdge(mergedEdges, a, b) {
                mergedEdges = append(mergedEdges, newEdge(a, b))
            }
        }

        // Connect: last detection in chain → primary_osd
        if primaryOSDID != "" {
            lastID := nodeIDInMerged(chain[len(chain)-1], dag.name, renamed)
            if !hasEdge(mergedEdges, lastID, primaryOSDID) {
                mergedEdges = append(mergedEdges, newEdge(lastID, primaryOSDID))
            }
        }
    }

    // --- 4. Merge alarm rules and OSD labels ---
    for _, dag := range supplemental {
        mergeAlarmRules(mergedNodes, dag)
        mergeOSDLabels(mergedNodes, dag)
    }

    // --- 5. Build result ---
    executor := map[string]any{
        "batch_size": 2,
        "run_budget": estimateRunBudget(dags),
    }

    templateBody := map[string]any{
        "executor": executor,
        "nodes":    mergedNodes,
        "edges":    mergedEdges,
    }

    name := "auto_" + strings.Join(dagNames(dags), "_")
    result := map[string]any{
        "name":        name,
        "description": "自动组合模板: " + strings.Join(dagNames(dags), " + "),
        "source":      "composed",
        "template":    templateBody,
    }

    return result, nil
}

func isInfraNode(cls templateNodeType) bool {
    switch cls {
    case nodeSource, nodePreprocess, nodeOSD, nodePublish, nodeAlarm:
        return true
    }
    return false
}

func findNodeByClass(dag *templateDAG, cls templateNodeType) string {
    for id, c := range dag.nodeClasses {
        if c == cls {
            return id
        }
    }
    return ""
}

// findDetectionChainEnds finds the first and last detection nodes in the DAG
// by traversing from preprocess forward and from OSD/publish backward.
func findDetectionChainEnds(dag *templateDAG) (entryID, exitID string) {
    prepID := findNodeByClass(dag, nodePreprocess)
    if prepID != "" {
        for _, succ := range dag.successors[prepID] {
            if dag.nodeClasses[succ] == nodeDetection {
                entryID = succ
                break
            }
        }
    }

    osdID := findNodeByClass(dag, nodeOSD)
    if osdID == "" {
        osdID = findNodeByClass(dag, nodePublish)
    }
    if osdID != "" {
        for _, pred := range dag.predecessors[osdID] {
            if dag.nodeClasses[pred] == nodeDetection {
                exitID = pred
                break
            }
        }
    }

    return entryID, exitID
}

// buildDetectionChain walks from entry to exit through detection nodes.
func buildDetectionChain(dag *templateDAG, entryID, exitID string) []string {
    if entryID == "" || exitID == "" {
        return nil
    }
    chain := []string{entryID}
    if entryID == exitID {
        return chain
    }
    visited := map[string]bool{entryID: true}
    current := entryID
    for current != exitID {
        var next string
        for _, succ := range dag.successors[current] {
            cls := dag.nodeClasses[succ]
            if cls == nodeDetection && !visited[succ] {
                next = succ
                break
            }
        }
        if next == "" {
            // Try the first unvisited successor regardless of type
            for _, succ := range dag.successors[current] {
                if !visited[succ] {
                    next = succ
                    break
                }
            }
        }
        if next == "" {
            break
        }
        chain = append(chain, next)
        visited[next] = true
        if next == exitID {
            break
        }
        current = next
    }
    return chain
}

func nodeIDInMerged(origID, dagName string, renamed map[string]string) string {
    if v, ok := renamed[dagName+"::"+origID]; ok {
        return v
    }
    if v, ok := renamed[origID]; ok {
        return v
    }
    return origID
}

func hasEdge(edges []any, from, to string) bool {
    for _, e := range edges {
        edge, _ := e.([]any)
        if len(edge) >= 2 && stringValue(edge[0]) == from && stringValue(edge[1]) == to {
            return true
        }
    }
    return false
}

func newEdge(from, to string) []any {
    return []any{from, to}
}

func deepCopyEdge(edge any) any {
    // Deep copy via JSON round-trip for safety.
    body, err := json.Marshal(edge)
    if err != nil {
        // Fallback: shallow copy.
        if arr, ok := edge.([]any); ok {
            out := make([]any, len(arr))
            copy(out, arr)
            return out
        }
        return edge
    }
    var out any
    json.Unmarshal(body, &out)
    return out
}

func mergeAlarmRules(nodes []any, dag *templateDAG) {
    for _, dagNode := range dag.nodes {
        if dag.nodeClasses[stringValue(dagNode["id"])] != nodeAlarm {
            continue
        }
        // Find the alarm node in merged nodes.
        for _, n := range nodes {
            merged, _ := n.(map[string]any)
            if merged == nil {
                continue
            }
            if stringValue(merged["type"]) == "alarm" {
                dagRules, _ := dagNode["rules"].([]any)
                mergedRules, _ := merged["rules"].([]any)
                if len(dagRules) > 0 {
                    // Merge rules, avoiding duplicates by name.
                    mergedNames := make(map[string]bool)
                    for _, r := range mergedRules {
                        if rm, ok := r.(map[string]any); ok {
                            mergedNames[stringValue(rm["name"])] = true
                        }
                    }
                    for _, r := range dagRules {
                        if rm, ok := r.(map[string]any); ok {
                            if !mergedNames[stringValue(rm["name"])] {
                                mergedRules = append(mergedRules, deepCopyAny(r))
                            }
                        }
                    }
                    merged["rules"] = mergedRules
                }
                break
            }
        }
    }
}

func mergeOSDLabels(nodes []any, dag *templateDAG) {
    for _, dagNode := range dag.nodes {
        if dag.nodeClasses[stringValue(dagNode["id"])] != nodeOSD {
            continue
        }
        for _, n := range nodes {
            merged, _ := n.(map[string]any)
            if merged == nil {
                continue
            }
            if stringValue(merged["type"]) == "osd" {
                dagLabels, _ := dagNode["labels"].([]any)
                mergedLabels, _ := merged["labels"].([]any)
                if len(dagLabels) > 0 {
                    mergedSet := make(map[string]bool)
                    for _, l := range mergedLabels {
                        if s, ok := l.(string); ok {
                            mergedSet[s] = true
                        }
                    }
                    for _, l := range dagLabels {
                        if s, ok := l.(string); ok {
                            if !mergedSet[s] {
                                mergedLabels = append(mergedLabels, s)
                            }
                        }
                    }
                    merged["labels"] = mergedLabels
                }
                break
            }
        }
    }
}

func estimateRunBudget(dags []*templateDAG) int {
    total := 0
    for _, dag := range dags {
        if dag.executor != nil {
            if b, ok := dag.executor["run_budget"].(float64); ok {
                total += int(b)
            }
        }
    }
    if total < 4 {
        total = 4
    }
    if total > 32 {
        total = 32
    }
    return total
}

func sanitizeIDPrefix(name string) string {
    // Keep alphanumeric and underscore, replace others.
    var b strings.Builder
    for _, r := range name {
        if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '_' {
            b.WriteRune(r)
        } else {
            b.WriteByte('_')
        }
    }
    return strings.Trim(b.String(), "_")
}

func uniqueID(candidate string, used map[string]bool) string {
    if !used[candidate] {
        return candidate
    }
    for i := 2; ; i++ {
        alt := fmt.Sprintf("%s_%d", candidate, i)
        if !used[alt] {
            return alt
        }
    }
}

// fixPreprocessClassification marks only the entry preprocess (directly after input)
// as infrastructure. Later preprocess nodes (prepare_publish) are reclassified as
// detection so they stay in the pipeline chain.
func fixPreprocessClassification(dag *templateDAG) {
    // Find source nodes.
    for id, cls := range dag.nodeClasses {
        if cls != nodeSource {
            continue
        }
        // The first preprocess successor of the source is the entry.
        for _, succ := range dag.successors[id] {
            if dag.nodeClasses[succ] == nodePreprocess {
                // Keep as nodePreprocess (infra).
                // All OTHER preprocess nodes become detection.
                for otherID, otherCls := range dag.nodeClasses {
                    if otherCls == nodePreprocess && otherID != succ {
                        dag.nodeClasses[otherID] = nodeDetection
                    }
                }
                return
            }
        }
    }
}

func dagNames(dags []*templateDAG) []string {
    out := make([]string, len(dags))
    for i, d := range dags {
        out[i] = d.name
    }
    return out
}

// sortStringSet returns sorted, deduplicated strings.
func sortStringSet(in []string) []string {
    seen := map[string]struct{}{}
    out := make([]string, 0, len(in))
    for _, s := range in {
        s = strings.TrimSpace(s)
        if s == "" {
            continue
        }
        if _, ok := seen[s]; ok {
            continue
        }
        seen[s] = struct{}{}
        out = append(out, s)
    }
    sort.Strings(out)
    return out
}