knowledge-graph-agent/backend/services/data_service.go

package services

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"sync"

	"knowledge-graph-backend/models"
)

type DataService struct {
	data     models.GraphData
	mutex    sync.RWMutex
	filePath string
}

func NewDataService(filePath string) (*DataService, error) {
	ds := &DataService{
		filePath: filePath,
		mutex:    sync.RWMutex{},
	}

	if err := ds.loadData(); err != nil {
		return nil, fmt.Errorf("failed to load data: %w", err)
	}

	return ds, nil
}

func (ds *DataService) loadData() error {
	ds.mutex.Lock()
	defer ds.mutex.Unlock()

	file, err := os.Open(ds.filePath)
	if err != nil {
		return fmt.Errorf("failed to open data file: %w", err)
	}
	defer file.Close()

	byteValue, err := ioutil.ReadAll(file)
	if err != nil {
		return fmt.Errorf("failed to read data file: %w", err)
	}

	if err := json.Unmarshal(byteValue, &ds.data); err != nil {
		return fmt.Errorf("failed to parse data file: %w", err)
	}

	log.Printf("Loaded %d nodes and %d edges from %s",
		len(ds.data.Nodes), len(ds.data.Edges), ds.filePath)

	return nil
}

func (ds *DataService) GetGraphData() models.GraphData {
	ds.mutex.RLock()
	defer ds.mutex.RUnlock()
	return ds.data
}

func (ds *DataService) GetNodeByID(id string) (models.Node, bool) {
	ds.mutex.RLock()
	defer ds.mutex.RUnlock()

	for _, node := range ds.data.Nodes {
		if node.ID == id {
			return node, true
		}
	}
	return models.Node{}, false
}

func (ds *DataService) SearchNodes(query string) []models.Node {
	ds.mutex.RLock()
	defer ds.mutex.RUnlock()

	if query == "" {
		return []models.Node{}
	}

	searchLower := strings.ToLower(query)
	var results []models.Node

	for _, node := range ds.data.Nodes {
		labelMatch := strings.Contains(strings.ToLower(node.Label), searchLower)
		idMatch := strings.Contains(strings.ToLower(node.ID), searchLower)
		typeMatch := node.Type != "" && strings.Contains(strings.ToLower(node.Type), searchLower)

		if labelMatch || idMatch || typeMatch {
			results = append(results, node)
		}
	}

	return results
}

func (ds *DataService) GetNeighbors(nodeID string) (models.NeighborResponse, bool) {
	ds.mutex.RLock()
	defer ds.mutex.RUnlock()

	_, exists := ds.GetNodeByID(nodeID)
	if !exists {
		return models.NeighborResponse{}, false
	}

	var neighborIDs []string
	var relatedEdges []models.Edge

	for _, edge := range ds.data.Edges {
		if edge.Source == nodeID {
			neighborIDs = append(neighborIDs, edge.Target)
			relatedEdges = append(relatedEdges, edge)
		} else if edge.Target == nodeID {
			neighborIDs = append(neighborIDs, edge.Source)
			relatedEdges = append(relatedEdges, edge)
		}
	}

	var neighborNodes []models.Node
	for _, node := range ds.data.Nodes {
		for _, neighborID := range neighborIDs {
			if node.ID == neighborID {
				neighborNodes = append(neighborNodes, node)
				break
			}
		}
	}

	return models.NeighborResponse{
		Nodes: neighborNodes,
		Edges: relatedEdges,
	}, true
}

func (ds *DataService) ReloadData() error {
	return ds.loadData()
}

func (ds *DataService) GetStats() map[string]int {
	ds.mutex.RLock()
	defer ds.mutex.RUnlock()

	typeCounts := make(map[string]int)
	for _, node := range ds.data.Nodes {
		if node.Type != "" {
			typeCounts[node.Type]++
		}
	}

	return map[string]int{
		"totalNodes":       len(ds.data.Nodes),
		"totalEdges":       len(ds.data.Edges),
		"conceptNodes":     typeCounts["概念"],
		"toolNodes":        typeCounts["工具"],
		"applicationNodes": typeCounts["应用"],
	}
}