xyj-figures/pkg/batch/batch.go

package batch

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"xyj-figures/pkg/entity"
	"xyj-figures/pkg/relationship"

	findfont "github.com/flopp/go-findfont"
	graphviz "github.com/goccy/go-graphviz"
	"github.com/goccy/go-graphviz/cgraph"
	"github.com/golang/freetype/truetype"
	log "github.com/sirupsen/logrus"
	"golang.org/x/image/font"
)

const (
	CLUSTER_PREFIX = "cluster_"
)

func LoadFolder(path string) ([]entity.Entity, []relationship.Relationship, error) {
	es := []entity.Entity{}
	rs := []relationship.Relationship{}
	errs := []string{}

	// Triverse all files in path and load all entities and relationships
	err := filepath.Walk(path, func(subpath string, info os.FileInfo, err error) error {
		if err != nil {
			errs = append(errs, err.Error())
			return nil
		}
		if !info.IsDir() {
			e, err := entity.Load(subpath)
			if err != nil {
				errs = append(errs, err.Error())
				return nil
			}
			es = append(es, *e)
			rs = append(rs, e.Details.(entity.EntityDetails).GetRelationships()...)
		}
		return nil
	})
	if len(errs) > 0 {
		return es, rs, fmt.Errorf("multiple errors: %v", strings.Join(errs, "; "))
	}
	return es, rs, err
}

// DistillData takes loaded data and do all needed optimizations and extract only needed info for maps
func DistillData(es []entity.Entity, rs []relationship.Relationship) ([]entity.Entity, []relationship.Relationship, error) {
	// Create a set of all entities, including characters and every members
	// TODO: Duplicated name
	nodes := map[string]entity.Entity{}
	edges := map[string]relationship.Relationship{}
	for _, e := range es {
		switch e.Type {
		case entity.EntityTypeCharacter:
			if val, ok := nodes[e.Details.(entity.Character).Name]; ok {
				if err := val.Merge(e); err != nil {
					log.Warnf("failed to merge: %v", err)
				}
				continue
			}
			nodes[e.Details.(entity.Character).Name] = e
		case entity.EntityTypeGroup:
			nodes[e.Details.(entity.Group).Name] = e
			for _, m := range e.Details.(entity.Group).Members {
				nodes[m] = entity.Entity{
					Schema: "v1",
					Type:   entity.EntityTypeCharacter,
					Details: entity.Character{
						Name: m,
					},
				}
			}
		}
	}

	for _, r := range rs {
		key1, key2 := r.GetKey()
		val1, ok1 := edges[key1]
		val2, ok2 := edges[key2]
		// Both keys don't exist
		if !ok1 && !ok2 {
			edges[key1] = r
			continue
		}
		// Logs all unexpected edges
		if ok1 && !r.IsSame(val1) {
			log.Warningf("Found ambigious edges, existing: %+v, new: %+v", val1, r)
		}
		if ok2 && !r.IsSame(val2) {
			log.Warningf("Found ambigious edges, existing: %+v, new: %+v", val2, r)
		}
	}

	retns := make([]entity.Entity, 0, len(nodes))
	for _, n := range nodes {
		retns = append(retns, n)
	}
	retes := make([]relationship.Relationship, 0, len(edges))
	for _, v := range edges {
		retes = append(retes, v)
	}

	return retns, retes, nil
}

// Save distilled info into DOT type files
func Save(path string, ents []entity.Entity, rels []relationship.Relationship) error {
	g := graphviz.New()
	g.SetLayout(graphviz.TWOPI)
	useSystemFont(g, "SimSum.ttc")
	rootgraph, err := g.Graph()
	rootgraph.SafeSet("fontname", "NSimSun", "Times-Roman")
	if err != nil {
		return fmt.Errorf("failed to create graph: %v", err)
	}
	defer func() {
		if err := rootgraph.Close(); err != nil {
			log.Errorf("failed to close graph: %v", err)
		}
		g.Close()
	}()

	nodesPtr := map[string]*cgraph.Node{}
	subGraphPtr := map[string]*cgraph.Graph{}
	nodeGraphMap := map[string]string{}
	for _, e := range ents {
		if e.Type == entity.EntityTypeGroup {
			subGraph := rootgraph.SubGraph(CLUSTER_PREFIX+e.Details.(entity.Group).Name, 1)
			subGraph.SetLabel(e.Details.(entity.Group).Name)
			// subGraph.SafeSet("fontname", "NSimSun", "Times-Roman")
			subGraphPtr[e.Details.(entity.Group).Name] = subGraph
			// Record all members
			for _, m := range e.Details.(entity.Group).Members {
				nodeGraphMap[m] = e.Details.(entity.Group).Name
			}
		}
	}
	fmt.Println(subGraphPtr)
	fmt.Println(nodeGraphMap)

	for _, e := range ents {
		if e.Type == entity.EntityTypeCharacter {
			log.Debugf("entity: %+v", e)
			if group, ok := nodeGraphMap[e.Details.(entity.Character).Name]; ok {
				ptr, ok := subGraphPtr[group]
				if !ok {
					log.Debugf("can't find sub graph %s", group)
					continue
				}
				if ptr == nil {
					log.Debugf("sub graph %s is nil", group)
					continue
				}
				n, err := ptr.CreateNode(e.Details.(entity.Character).Name)
				if err != nil {
					log.Errorf("failed to create node in sub graph %s: %v", group, err)
					continue
				}
				nodesPtr[e.Details.(entity.Character).Name] = n
				continue
			}
			node, err := rootgraph.CreateNode(e.Details.(entity.Character).Name)
			if err != nil {
				log.Errorf("failed to create node: %v", err)
				continue
			}
			nodesPtr[e.Details.(entity.Character).Name] = node
			continue
		}
	}
	fmt.Println(nodesPtr)

	for _, r := range rels {
		key, _ := r.GetKey()
		srcPtr := nodesPtr[r.Src]
		dstPtr := nodesPtr[r.Dst]
		fmt.Printf("src: %s, dst: %s\n", r.Src, r.Dst)
		e, err := rootgraph.CreateEdge(key, srcPtr, dstPtr)
		if err != nil {
			log.Errorf("failed to create edge: %v", err)
			continue
		}
		e.SetLabel(r.Description)
		e.SafeSet("fontname", "NSimSun", "Times-Roman")
		e.SafeSet("fontsize", "12", "10")
		if r.Bidirection {
			e.SetDir(cgraph.NoneDir)
		}
	}
	if err := g.RenderFilename(rootgraph, "dot", path); err != nil {
		log.Fatal(err)
	}

	if err := g.RenderFilename(rootgraph, graphviz.PNG, path+".png"); err != nil {
		return fmt.Errorf("failed to save graph to disk: %v", err)
	}
	return nil
}

func useSystemFont(g *graphviz.Graphviz, name string) {
	fontPath, err := findfont.Find(name)
	if err != nil {
		panic(err)
	}

	fontData, err := os.ReadFile(fontPath)
	if err != nil {
		panic(err)
	}
	ft, err := truetype.Parse(fontData)
	if err != nil {
		panic(err)
	}
	g.SetFontFace(func(size float64) (font.Face, error) {
		opt := &truetype.Options{
			Size: size,
		}
		return truetype.NewFace(ft, opt), nil
	})
}