main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/kluctl/go-embed-python/python"
)

// Configuration options
type Config struct {
	DataInterval    time.Duration
	BufferMaxSize   int
	PythonScriptDir string
	Debug           bool
}

// Simulates generating InfluxDB line protocol data
func generateLineProtocolData() []string {
	// Generate some sample measurements
	timestamp := time.Now().UnixNano()
	lines := []string{
		fmt.Sprintf("cpu,host=server01,region=us-west usage_user=0.64,usage_system=0.21 %d", timestamp),
		fmt.Sprintf("memory,host=server01,region=us-west used_percent=72.45,available=4321000000i %d", timestamp),
		fmt.Sprintf("disk,host=server01,region=us-west used_percent=65.23,free=2150000000i %d", timestamp),
	}
	return lines
}

// Converts line protocol data to a format that can be passed to Python
func formatDataForPython(lines []string) map[string]interface{} {
	// A simplified example - in reality, you'd parse the line protocol into proper batches
	measurements := make(map[string][]map[string]interface{})

	for _, line := range lines {
		parts := strings.Split(line, " ")
		if len(parts) < 2 {
			continue
		}
		
		measurement := strings.Split(parts[0], ",")[0]
		
		// Add to the appropriate batch based on measurement name
		row := map[string]interface{}{
			"line": line,
			"timestamp": parts[len(parts)-1],
		}
		
		measurements[measurement] = append(measurements[measurement], row)
	}

	// Construct table batches as expected by the Python function
	var batches []map[string]interface{}
	for tableName, rows := range measurements {
		batches = append(batches, map[string]interface{}{
			"table_name": tableName,
			"rows":      rows,
		})
	}

	return map[string]interface{}{
		"table_batches": batches,
	}
}

// A circular buffer implementation for storing line protocol data
type LineBuffer struct {
	data     []string
	maxSize  int
	mu       sync.Mutex
}

func NewLineBuffer(maxSize int) *LineBuffer {
	return &LineBuffer{
		data:    make([]string, 0, maxSize),
		maxSize: maxSize,
	}
}

func (b *LineBuffer) Add(lines []string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	
	b.data = append(b.data, lines...)
	
	// If we exceed max size, trim from the beginning
	if len(b.data) > b.maxSize {
		overflow := len(b.data) - b.maxSize
		b.data = b.data[overflow:]
	}
}

func (b *LineBuffer) GetAll() []string {
	b.mu.Lock()
	defer b.mu.Unlock()
	
	// Return a copy to avoid concurrent modification
	result := make([]string, len(b.data))
	copy(result, b.data)
	return result
}

func (b *LineBuffer) Clear() {
	b.mu.Lock()
	defer b.mu.Unlock()
	
	b.data = b.data[:0]
}

func (b *LineBuffer) Size() int {
	b.mu.Lock()
	defer b.mu.Unlock()
	
	return len(b.data)
}

func main() {
	// Parse command-line flags
	dataInterval := flag.Duration("interval", 5*time.Second, "Data generation interval")
	bufferSize := flag.Int("buffer-size", 1000, "Maximum buffer size")
	scriptDir := flag.String("script-dir", "scripts", "Directory for Python scripts")
	debug := flag.Bool("debug", false, "Enable debug output")
	
	flag.Parse()
	
	config := Config{
		DataInterval:    *dataInterval,
		BufferMaxSize:   *bufferSize,
		PythonScriptDir: *scriptDir,
		Debug:           *debug,
	}
	
	// Create script directory if it doesn't exist
	if err := os.MkdirAll(config.PythonScriptDir, 0755); err != nil {
		log.Fatalf("Failed to create script directory: %v", err)
	}
	
	// Create the processor scripts
	if err := createProcessorScripts(config.PythonScriptDir); err != nil {
		log.Fatalf("Failed to create processor scripts: %v", err)
	}
	
	// Initialize go-embed-python
	ep, err := python.NewEmbeddedPython("influxdb-pipeline")
	if err != nil {
		log.Fatalf("Failed to initialize embedded Python: %v", err)
	}
	
	// Create a buffer for line protocol data
	buffer := NewLineBuffer(config.BufferMaxSize)
	
	// Set up signal handling for graceful shutdown
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
	
	// Create a ticker for data generation
	ticker := time.NewTicker(config.DataInterval)
	defer ticker.Stop()
	
	log.Printf("Starting data generation (interval: %s, buffer size: %d)", config.DataInterval, config.BufferMaxSize)
	log.Println("Press Ctrl+C to stop")
	
	// Main loop
	running := true
	for running {
		select {
		case <-ticker.C:
			// Generate new data
			newData := generateLineProtocolData()
			buffer.Add(newData)
			
			bufferSize := buffer.Size()
			log.Printf("Generated %d new lines, buffer now has %d lines", len(newData), bufferSize)
			
			// Skip processing if buffer is empty
			if bufferSize == 0 {
				continue
			}
			
			// Get all data from buffer
			allData := buffer.GetAll()
			
			// Format data for Python
			dataForPython := formatDataForPython(allData)
			
			// Convert data to JSON
			dataJSON, err := json.Marshal(dataForPython)
			if err != nil {
				log.Printf("Error marshaling data to JSON: %v", err)
				continue
			}
			
			// Create a command to run the processor
			processorPath := fmt.Sprintf("%s/processor_server.py", config.PythonScriptDir)
			cmd, err := ep.PythonCmd("-u", processorPath)
			if err != nil {
				log.Printf("Error creating Python command: %v", err)
				continue
			}
			
			// Set up stdin/stdout pipes
			stdin, err := cmd.StdinPipe()
			if err != nil {
				log.Printf("Error getting stdin pipe: %v", err)
				continue
			}
			
			if config.Debug {
				cmd.Stderr = os.Stderr
			}
			
			// Start command
			var outBuf strings.Builder
			cmd.Stdout = &outBuf
			
			if err := cmd.Start(); err != nil {
				log.Printf("Error starting Python command: %v", err)
				continue
			}
			
			// Write data to stdin
			if _, err := stdin.Write(dataJSON); err != nil {
				log.Printf("Error writing to stdin: %v", err)
				stdin.Close()
				continue
			}
			
			// Close stdin to signal we're done writing
			stdin.Close()
			
			// Wait for command to complete
			if err := cmd.Wait(); err != nil {
				log.Printf("Error running Python command: %v", err)
				continue
			}
			
			// Parse the output
			outputStr := outBuf.String()
			if outputStr == "" {
				log.Println("No output from Python command")
				continue
			}
			
			// Parse the JSON result
			var result map[string]interface{}
			if err := json.Unmarshal([]byte(outputStr), &result); err != nil {
				log.Printf("Error parsing JSON result: %v\nOutput: %s", err, outputStr)
				continue
			}
			
			// Check for errors
			if errorMsg, ok := result["error"].(string); ok {
				log.Printf("Python error: %s", errorMsg)
				continue
			}

			// Print the result
			if config.Debug {
			    resultJSON, _ := json.MarshalIndent(result, "", "  ")
			    log.Printf("Python processing result: %s", resultJSON)
			} else {
			    processed := int(result["processed"].(float64))
			    writtenLines := result["written_lines"].([]interface{})
			    log.Printf("Processed %d table batches, generated %d derived metrics", 
			        processed, len(writtenLines))
			}
			
			// Clear the buffer
			buffer.Clear()
			
		case <-signals:
			log.Println("Shutting down...")
			running = false
		}
	}
	
	log.Println("Goodbye!")
}

// Create Python scripts for processing
func createProcessorScripts(scriptDir string) error {
	// Create the data processor module
	dataProcessorPath := fmt.Sprintf("%s/data_processor.py", scriptDir)
	dataProcessorContent := `
import sys

class LineBuilder:
    def __init__(self, measurement):
        self.measurement = measurement
        self.tags = {}
        self.fields = {}
        
    def tag(self, key, value):
        self.tags[key] = value
        return self
        
    def int64_field(self, key, value):
        self.fields[key] = f"{value}i"
        return self
        
    def to_line(self):
        tag_str = ",".join([f"{k}={v}" for k, v in self.tags.items()])
        field_str = ",".join([f"{k}={v}" for k, v in self.fields.items()])
        
        if tag_str:
            return f"{self.measurement},{tag_str} {field_str}"
        else:
            return f"{self.measurement} {field_str}"

class InfluxDB3Local:
    def __init__(self):
        self.lines = []
    
    def info(self, message):
        print(f"[INFO] {message}", file=sys.stderr)
    
    def write(self, line):
        if isinstance(line, LineBuilder):
            line_str = line.to_line()
        else:
            line_str = str(line)
        self.lines.append(line_str)
        print(f"[WRITE] {line_str}", file=sys.stderr)
        return True

def process_writes(influxdb3_local, table_batches, args=None):
    # Create InfluxDB mock if not provided
    if not isinstance(influxdb3_local, InfluxDB3Local):
        influxdb3_local = InfluxDB3Local()
    
    # Process data as it's written to the database
    for table_batch in table_batches:
        table_name = table_batch["table_name"]
        rows = table_batch["rows"]
        
        # Log information about the write
        influxdb3_local.info(f"Processing {len(rows)} rows from {table_name}")
        
        # Write derived data back to the database
        line = LineBuilder("processed_data")
        line.tag("source_table", table_name)
        line.int64_field("row_count", len(rows))
        influxdb3_local.write(line)
    
    return {
        "processed": len(table_batches),
        "written_lines": influxdb3_local.lines
    }
`
	
	if err := os.WriteFile(dataProcessorPath, []byte(dataProcessorContent), 0644); err != nil {
		return err
	}
	
	// Create the server script
	serverScriptPath := fmt.Sprintf("%s/processor_server.py", scriptDir)
	serverScriptContent := `
import json
import sys
import traceback
from data_processor import process_writes, InfluxDB3Local

# Create a persistent InfluxDB local instance
influxdb3_local = InfluxDB3Local()

try:
    # Read input data from stdin
    data = json.load(sys.stdin)
    table_batches = data.get("table_batches", [])
    
    # Process the data
    result = process_writes(influxdb3_local, table_batches)
    
    # Output the result as JSON
    print(json.dumps(result))
    
except Exception as e:
    # Print traceback to stderr for debugging
    traceback.print_exc(file=sys.stderr)
    
    # Return error information to Go
    error_result = {
        "error": str(e),
        "processed": 0,
        "written_lines": []
    }
    print(json.dumps(error_result))
`
	
	return os.WriteFile(serverScriptPath, []byte(serverScriptContent), 0644)
}