streamsql/stream/persistence.go

package stream

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/rulego/streamsql/logger"
)

// PersistenceManager 数据持久化管理器
type PersistenceManager struct {
	dataDir       string        // 持久化数据目录
	maxFileSize   int64         // 单个文件最大大小（字节）
	flushInterval time.Duration // 刷新间隔
	currentFile   *os.File      // 当前写入文件
	currentSize   int64         // 当前文件大小
	fileIndex     int           // 文件索引
	writeMutex    sync.Mutex    // 写入互斥锁
	flushTimer    *time.Timer   // 刷新定时器
	pendingData   []interface{} // 待写入数据
	pendingMutex  sync.Mutex    // 待写入数据互斥锁
	isRunning     bool          // 是否运行中
	runningMutex  sync.RWMutex  // 保护isRunning字段的读写锁
	stopChan      chan struct{} // 停止通道

	// 统计信息 (新增)
	totalPersisted int64
	totalLoaded    int64
	filesCreated   int64
}

// NewPersistenceManager 创建默认配置的持久化管理器
func NewPersistenceManager(dataDir string) *PersistenceManager {
	pm := &PersistenceManager{
		dataDir:       dataDir,
		maxFileSize:   10 * 1024 * 1024, // 10MB per file
		flushInterval: 5 * time.Second,  // 5秒刷新一次
		fileIndex:     0,
		pendingData:   make([]interface{}, 0),
		stopChan:      make(chan struct{}),
	}

	// 确保数据目录存在
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		logger.Error("Failed to create persistence directory: %v", err)
	}

	return pm
}

// NewPersistenceManagerWithConfig 创建自定义配置的持久化管理器
func NewPersistenceManagerWithConfig(dataDir string, maxFileSize int64, flushInterval time.Duration) *PersistenceManager {
	pm := &PersistenceManager{
		dataDir:       dataDir,
		maxFileSize:   maxFileSize,
		flushInterval: flushInterval,
		fileIndex:     0,
		pendingData:   make([]interface{}, 0),
		stopChan:      make(chan struct{}),
	}

	// 确保数据目录存在
	if err := os.MkdirAll(dataDir, 0755); err != nil {
		logger.Error("Failed to create persistence directory: %v", err)
	}

	return pm
}

// Start 启动持久化管理器
func (pm *PersistenceManager) Start() error {
	// 检查是否已经在运行
	pm.runningMutex.RLock()
	running := pm.isRunning
	pm.runningMutex.RUnlock()

	if running {
		return fmt.Errorf("persistence manager already running")
	}

	// 创建初始文件
	pm.writeMutex.Lock()
	if err := pm.createNewFile(); err != nil {
		pm.writeMutex.Unlock()
		return fmt.Errorf("failed to create initial file: %w", err)
	}
	pm.writeMutex.Unlock()

	// 设置运行状态
	pm.runningMutex.Lock()
	pm.isRunning = true
	pm.runningMutex.Unlock()

	// 启动定时刷新
	pm.startFlushTimer()

	// 启动后台处理协程
	go pm.backgroundProcessor()

	logger.Info("Persistence manager started successfully, data directory: %s", pm.dataDir)
	return nil
}

// Stop 停止持久化管理器
func (pm *PersistenceManager) Stop() error {
	// 检查是否正在运行
	pm.runningMutex.RLock()
	running := pm.isRunning
	pm.runningMutex.RUnlock()

	if !running {
		return nil
	}

	// 设置停止状态
	pm.runningMutex.Lock()
	pm.isRunning = false
	pm.runningMutex.Unlock()

	close(pm.stopChan)

	// 停止定时器
	pm.writeMutex.Lock()
	if pm.flushTimer != nil {
		pm.flushTimer.Stop()
	}
	pm.writeMutex.Unlock()

	// 刷新剩余数据
	pm.flushPendingData()

	// 关闭当前文件
	pm.writeMutex.Lock()
	if pm.currentFile != nil {
		pm.currentFile.Close()
		pm.currentFile = nil
	}
	pm.writeMutex.Unlock()

	logger.Info("Persistence manager stopped")
	return nil
}

// PersistData 持久化单条数据
func (pm *PersistenceManager) PersistData(data interface{}) error {
	// 检查是否正在运行
	pm.runningMutex.RLock()
	running := pm.isRunning
	pm.runningMutex.RUnlock()

	if !running {
		return fmt.Errorf("persistence manager not running")
	}

	pm.pendingMutex.Lock()
	pm.pendingData = append(pm.pendingData, data)
	pm.pendingMutex.Unlock()

	return nil
}

// LoadPersistedData 加载并删除持久化数据
func (pm *PersistenceManager) LoadPersistedData() ([]interface{}, error) {
	files, err := filepath.Glob(filepath.Join(pm.dataDir, "streamsql_overflow_*.log"))
	if err != nil {
		return nil, fmt.Errorf("failed to glob files: %w", err)
	}

	var allData []interface{}

	for _, filename := range files {
		data, err := pm.loadDataFromFile(filename)
		if err != nil {
			logger.Error("Failed to load file %s: %v", filename, err)
			continue
		}
		allData = append(allData, data...)

		// 在锁保护下更新统计信息
		pm.writeMutex.Lock()
		pm.totalLoaded += int64(len(data))
		pm.writeMutex.Unlock()

		// 加载后删除文件
		if err := os.Remove(filename); err != nil {
			logger.Error("Failed to delete loaded file %s: %v", filename, err)
		}
	}

	logger.Info("Loaded %d data records from persistence files", len(allData))
	return allData, nil
}

// GetStats 获取持久化统计信息
func (pm *PersistenceManager) GetStats() map[string]interface{} {
	pm.pendingMutex.Lock()
	pendingCount := len(pm.pendingData)
	pm.pendingMutex.Unlock()

	pm.writeMutex.Lock()
	currentFileSize := pm.currentSize
	fileIndex := pm.fileIndex
	totalPersisted := pm.totalPersisted
	totalLoaded := pm.totalLoaded
	filesCreated := pm.filesCreated
	pm.writeMutex.Unlock()

	// 安全地读取运行状态
	pm.runningMutex.RLock()
	running := pm.isRunning
	pm.runningMutex.RUnlock()

	return map[string]interface{}{
		"running":           running,
		"data_dir":          pm.dataDir,
		"pending_count":     pendingCount,
		"current_file_size": currentFileSize,
		"file_index":        fileIndex,
		"max_file_size":     pm.maxFileSize,
		"flush_interval":    pm.flushInterval.String(),
		"total_persisted":   totalPersisted,
		"total_loaded":      totalLoaded,
		"files_created":     filesCreated,
	}
}

// createNewFile 创建新的持久化文件
func (pm *PersistenceManager) createNewFile() error {
	// 关闭当前文件
	if pm.currentFile != nil {
		pm.currentFile.Close()
	}

	// 生成新文件名
	filename := fmt.Sprintf("streamsql_overflow_%d_%d.log",
		time.Now().Unix(), pm.fileIndex)
	filepath := filepath.Join(pm.dataDir, filename)

	// 创建新文件
	file, err := os.OpenFile(filepath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
	if err != nil {
		return fmt.Errorf("failed to create file %s: %w", filepath, err)
	}

	pm.currentFile = file
	pm.currentSize = 0
	pm.fileIndex++
	pm.filesCreated++

	// logger.Info("Created new persistence file: %s", filepath)
	return nil
}

// writeDataToFile 将数据写入文件
// 注意：此方法应该在writeMutex锁保护下调用
func (pm *PersistenceManager) writeDataToFile(data interface{}) error {
	if pm.currentFile == nil {
		return fmt.Errorf("no current file")
	}

	// 序列化数据
	jsonData, err := json.Marshal(map[string]interface{}{
		"timestamp": time.Now().Unix(),
		"data":      data,
	})
	if err != nil {
		return fmt.Errorf("failed to marshal data: %w", err)
	}

	// 添加换行符
	jsonData = append(jsonData, '\n')

	// 检查文件大小
	if pm.currentSize+int64(len(jsonData)) > pm.maxFileSize {
		if err := pm.createNewFile(); err != nil {
			return fmt.Errorf("failed to create new file: %w", err)
		}
	}

	// 写入数据
	n, err := pm.currentFile.Write(jsonData)
	if err != nil {
		return fmt.Errorf("failed to write data: %w", err)
	}

	pm.currentSize += int64(n)
	pm.totalPersisted++
	return nil
}

// flushPendingData 刷新待写入数据
func (pm *PersistenceManager) flushPendingData() {
	pm.pendingMutex.Lock()
	dataToWrite := make([]interface{}, len(pm.pendingData))
	copy(dataToWrite, pm.pendingData)
	pm.pendingData = pm.pendingData[:0] // 清空切片
	pm.pendingMutex.Unlock()

	if len(dataToWrite) == 0 {
		return
	}

	pm.writeMutex.Lock()
	defer pm.writeMutex.Unlock()

	// 批量写入数据
	for _, data := range dataToWrite {
		if err := pm.writeDataToFile(data); err != nil {
			logger.Error("Failed to write persistence data: %v", err)
		}
	}

	// 同步到磁盘
	if pm.currentFile != nil {
		_ = pm.currentFile.Sync()
	}

	// logger.Info("Flushed %d pending data records to disk", len(dataToWrite))
}

// startFlushTimer 启动刷新定时器
func (pm *PersistenceManager) startFlushTimer() {
	pm.writeMutex.Lock()
	pm.flushTimer = time.AfterFunc(pm.flushInterval, func() {
		// 安全地检查运行状态
		pm.runningMutex.RLock()
		running := pm.isRunning
		pm.runningMutex.RUnlock()

		if running {
			pm.flushPendingData()
			pm.startFlushTimer() // 重新启动定时器
		}
	})
	pm.writeMutex.Unlock()
}

// backgroundProcessor 后台处理协程
func (pm *PersistenceManager) backgroundProcessor() {
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			// 定期检查并处理
			pm.pendingMutex.Lock()
			pendingCount := len(pm.pendingData)
			pm.pendingMutex.Unlock()

			// 如果有大量待写入数据，立即刷新
			if pendingCount > 100 {
				pm.flushPendingData()
			}

		case <-pm.stopChan:
			return
		}
	}
}

// loadDataFromFile 从文件加载数据
func (pm *PersistenceManager) loadDataFromFile(filename string) ([]interface{}, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, fmt.Errorf("failed to open file %s: %w", filename, err)
	}
	defer file.Close()

	var data []interface{}
	scanner := bufio.NewScanner(file)

	for scanner.Scan() {
		line := scanner.Text()
		var record map[string]interface{}

		if err := json.Unmarshal([]byte(line), &record); err != nil {
			logger.Error("Failed to parse data line: %v", err)
			continue
		}

		// 提取实际数据
		if actualData, ok := record["data"]; ok {
			data = append(data, actualData)
		}
	}

	if err := scanner.Err(); err != nil {
		return nil, fmt.Errorf("failed to scan file: %w", err)
	}

	return data, nil
}