network-topology-discovery/internal/storage/topology_storage.go

package storage

import (
	"crypto/rand"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"network-topology-discovery/pkg/models"
)

// TopologyStorage 拓扑存储管理
type TopologyStorage struct {
	mu            sync.RWMutex
	dataDir       string
	topologies    map[string]*models.Topology
	currentTopoID string
}

// NewTopologyStorage 创建拓扑存储
func NewTopologyStorage(dataDir string) (*TopologyStorage, error) {
	s := &TopologyStorage{
		dataDir:    dataDir,
		topologies: make(map[string]*models.Topology),
	}

	// 确保数据目录存在
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		return nil, fmt.Errorf("failed to create data directory: %w", err)
	}

	// 加载拓扑元数据
	if err := s.loadMeta(); err != nil {
		if !os.IsNotExist(err) {
			return nil, fmt.Errorf("failed to load topology meta: %w", err)
		}
		log.Printf("Creating new topology storage at %s", dataDir)
	}

	return s, nil
}

// loadMeta 加载拓扑元数据
func (s *TopologyStorage) loadMeta() error {
	metaFile := filepath.Join(s.dataDir, "topologies.json")
	data, err := os.ReadFile(metaFile)
	if err != nil {
		return err
	}

	var topos []models.Topology
	if err := json.Unmarshal(data, &topos); err != nil {
		return fmt.Errorf("failed to parse topology meta: %w", err)
	}

	for i := range topos {
		s.topologies[topos[i].ID] = &topos[i]
	}

	log.Printf("Loaded %d topologies from meta", len(topos))
	return nil
}

// saveMeta 保存拓扑元数据
func (s *TopologyStorage) saveMeta() error {
	topos := make([]models.Topology, 0, len(s.topologies))
	for _, topo := range s.topologies {
		topos = append(topos, *topo)
	}

	data, err := json.MarshalIndent(topos, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal topologies: %w", err)
	}

	metaFile := filepath.Join(s.dataDir, "topologies.json")
	if err := os.WriteFile(metaFile, data, 0644); err != nil {
		return fmt.Errorf("failed to write topology meta: %w", err)
	}

	return nil
}

// CreateTopology 创建新拓扑
func (s *TopologyStorage) CreateTopology(name, description, scanRange string, sshPort int, username string) (*models.Topology, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	topo := &models.Topology{
		ID:          generateID(),
		Name:        name,
		Description: description,
		ScanRange:   scanRange,
		SSHPort:     sshPort,
		Username:    username,
		CreatedAt:   time.Now(),
		UpdatedAt:   time.Now(),
		DeviceCount: 0,
	}

	s.topologies[topo.ID] = topo

	// 创建该拓扑的设备文件
	deviceFile := filepath.Join(s.dataDir, topo.ID+"_devices.json")
	if _, err := os.Stat(deviceFile); os.IsNotExist(err) {
		if err := os.WriteFile(deviceFile, []byte("[]"), 0644); err != nil {
			return nil, fmt.Errorf("failed to create device file: %w", err)
		}
	}

	// 保存元数据
	if err := s.saveMeta(); err != nil {
		return nil, err
	}

	log.Printf("Topology created: %s (%s)", topo.Name, topo.ID)
	return topo, nil
}

// GetTopology 获取拓扑
func (s *TopologyStorage) GetTopology(id string) (*models.Topology, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	topo, exists := s.topologies[id]
	if !exists {
		return nil, fmt.Errorf("topology not found: %s", id)
	}

	return topo, nil
}

// GetAllTopologies 获取所有拓扑
func (s *TopologyStorage) GetAllTopologies() ([]models.Topology, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	topos := make([]models.Topology, 0, len(s.topologies))
	for _, topo := range s.topologies {
		topos = append(topos, *topo)
	}

	return topos, nil
}

// UpdateTopology 更新拓扑
func (s *TopologyStorage) UpdateTopology(id string, updates map[string]interface{}) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	topo, exists := s.topologies[id]
	if !exists {
		return fmt.Errorf("topology not found: %s", id)
	}

	if name, ok := updates["name"].(string); ok {
		topo.Name = name
	}
	if desc, ok := updates["description"].(string); ok {
		topo.Description = desc
	}
	if scanRange, ok := updates["scan_range"].(string); ok {
		topo.ScanRange = scanRange
	}
	if sshPort, ok := updates["ssh_port"].(float64); ok {
		topo.SSHPort = int(sshPort)
	}
	if username, ok := updates["username"].(string); ok {
		topo.Username = username
	}

	topo.UpdatedAt = time.Now()

	return s.saveMeta()
}

// DeleteTopology 删除拓扑及其所有数据
func (s *TopologyStorage) DeleteTopology(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	if _, exists := s.topologies[id]; !exists {
		return fmt.Errorf("topology not found: %s", id)
	}

	// 删除设备文件
	deviceFile := filepath.Join(s.dataDir, id+"_devices.json")
	if err := os.Remove(deviceFile); err != nil && !os.IsNotExist(err) {
		log.Printf("Warning: failed to delete device file for topology %s: %v", id, err)
	}

	delete(s.topologies, id)

	// 更新当前拓扑
	if s.currentTopoID == id {
		s.currentTopoID = ""
	}

	return s.saveMeta()
}

// SetCurrentTopology 设置当前拓扑
func (s *TopologyStorage) SetCurrentTopology(id string) error {
	if _, err := s.GetTopology(id); err != nil {
		return err
	}

	s.mu.Lock()
	defer s.mu.Unlock()

	s.currentTopoID = id
	log.Printf("Current topology set to: %s", id)
	return nil
}

// GetCurrentTopologyID 获取当前拓扑ID
func (s *TopologyStorage) GetCurrentTopologyID() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.currentTopoID
}

// GetDeviceFilePath 获取拓扑的设备文件路径
func (s *TopologyStorage) GetDeviceFilePath(topoID string) string {
	return filepath.Join(s.dataDir, topoID+"_devices.json")
}

// generateID 生成唯一ID
func generateID() string {
	b := make([]byte, 16)
	rand.Read(b)
	return fmt.Sprintf("%x-%x-%x-%x-%x", b[0:4], b[4:6], b[6:8], b[8:10], b[10:])
}