modelRT/database/query_topologic.go

// Package database define database operation functions
package database

import (
	"context"
	"fmt"
	"time"

	"modelRT/constants"
	"modelRT/diagram"
	"modelRT/logger"
	"modelRT/orm"
	"modelRT/sql"

	"github.com/gofrs/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// QueryTopologic return the topologic info of the circuit diagram
func QueryTopologic(ctx context.Context, tx *gorm.DB) ([]orm.Topologic, error) {
	var topologics []orm.Topologic
	// ctx超时判断
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).Clauses(clause.Locking{Strength: "UPDATE"}).Raw(sql.RecursiveSQL, constants.UUIDNilStr).Scan(&topologics)
	if result.Error != nil {
		logger.Error(ctx, "query circuit diagram topologic info by start node uuid failed", "start_node_uuid", constants.UUIDNilStr, "error", result.Error)
		return nil, result.Error
	}
	return topologics, nil
}

// QueryTopologicByStartUUID returns all edges reachable from startUUID following
// directed uuid_from → uuid_to edges in the topologic table.
func QueryTopologicByStartUUID(ctx context.Context, tx *gorm.DB, startUUID uuid.UUID) ([]orm.Topologic, error) {
	var topologics []orm.Topologic

	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Clauses(clause.Locking{Strength: "UPDATE"}).
		Raw(sql.RecursiveSQL, startUUID).
		Scan(&topologics)
	if result.Error != nil {
		logger.Error(ctx, "query topologic by start uuid failed", "start_uuid", startUUID, "error", result.Error)
		return nil, result.Error
	}
	return topologics, nil
}

// QueryTopologicFromDB return the result of query topologic info from DB.
// Returns the root node and a flat nodeMap for O(1) lookup by UUID.
func QueryTopologicFromDB(ctx context.Context, tx *gorm.DB) (*diagram.MultiBranchTreeNode, map[uuid.UUID]*diagram.MultiBranchTreeNode, error) {
	topologicInfos, err := QueryTopologic(ctx, tx)
	if err != nil {
		logger.Error(ctx, "query topologic info failed", "error", err)
		return nil, nil, err
	}

	tree, nodeMap, err := BuildMultiBranchTree(topologicInfos)
	if err != nil {
		logger.Error(ctx, "init topologic failed", "error", err)
		return nil, nil, err
	}
	return tree, nodeMap, nil
}

// BuildMultiBranchTree return the multi branch tree by topologic info.
// Returns the root node and a flat nodeMap for O(1) lookup by UUID.
func BuildMultiBranchTree(topologics []orm.Topologic) (*diagram.MultiBranchTreeNode, map[uuid.UUID]*diagram.MultiBranchTreeNode, error) {
	nodeMap := make(map[uuid.UUID]*diagram.MultiBranchTreeNode, len(topologics)*2)

	for _, topo := range topologics {
		if _, exists := nodeMap[topo.UUIDFrom]; !exists {
			// UUIDNil is the virtual root sentinel — skip creating a regular node for it
			if topo.UUIDFrom != constants.UUIDNil {
				nodeMap[topo.UUIDFrom] = &diagram.MultiBranchTreeNode{
					ID:       topo.UUIDFrom,
					Children: make([]*diagram.MultiBranchTreeNode, 0),
				}
			}
		}

		if _, exists := nodeMap[topo.UUIDTo]; !exists {
			if topo.UUIDTo != constants.UUIDNil {
				nodeMap[topo.UUIDTo] = &diagram.MultiBranchTreeNode{
					ID:       topo.UUIDTo,
					Children: make([]*diagram.MultiBranchTreeNode, 0),
				}
			}
		}
	}

	for _, topo := range topologics {
		var parent *diagram.MultiBranchTreeNode
		if topo.UUIDFrom == constants.UUIDNil {
			if _, exists := nodeMap[constants.UUIDNil]; !exists {
				nodeMap[constants.UUIDNil] = &diagram.MultiBranchTreeNode{
					ID:       constants.UUIDNil,
					Children: make([]*diagram.MultiBranchTreeNode, 0),
				}
			}
			parent = nodeMap[constants.UUIDNil]
		} else {
			parent = nodeMap[topo.UUIDFrom]
		}

		var child *diagram.MultiBranchTreeNode
		if topo.UUIDTo == constants.UUIDNil {
			child = &diagram.MultiBranchTreeNode{
				ID: topo.UUIDTo,
			}
		} else {
			child = nodeMap[topo.UUIDTo]
		}
		child.Parent = parent
		parent.Children = append(parent.Children, child)
	}

	// return root vertex
	root, exists := nodeMap[constants.UUIDNil]
	if !exists {
		return nil, nil, fmt.Errorf("root node not found")
	}
	return root, nodeMap, nil
}