main.go

// main.go - Entry point for the Memory Engine application
package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"syscall"
	"time"

	"memory-engine/pkg/api"
	"memory-engine/pkg/config"
	"memory-engine/pkg/embeddings"
	"memory-engine/pkg/engine"
	"memory-engine/pkg/memory"
)

const (
	AppName    = "Memory Engine"
	AppVersion = "1.0.0"
	AppDesc    = "Semantic Memory Storage & Retrieval Engine with B+Tree indexing"
)

func main() {
	// Print banner
	printBanner()

	// Initialize configuration manager
	// Check for config path in environment variable or use default
	configPath := os.Getenv("MEMORY_ENGINE_CONFIG_PATH")
	if configPath == "" {
		configPath = "./configs/development.yaml"
	}
	configManager := config.NewConfigManager(configPath)

	// Load configuration
	cfg, err := configManager.Load()
	if err != nil {
		log.Fatalf("Failed to load configuration: %v", err)
	}

	// Initialize logger
	logger := initializeLogger(cfg)
	logger.Info("Starting Memory Engine", "version", AppVersion, "environment", cfg.App.Environment)

	// Debug: Log Gemini configuration
	logger.Debug("Gemini configuration", "api_key_length", len(cfg.Gemini.APIKey), "model", cfg.Gemini.Model)
	if cfg.Gemini.APIKey != "" {
		logger.Debug("Gemini API key loaded successfully")
	} else {
		logger.Debug("No Gemini API key found in configuration")
	}

	// Validate environment
	if err := validateEnvironment(cfg, logger); err != nil {
		logger.Error("Environment validation failed", "error", err)
		os.Exit(1)
	}

	// Initialize for environment
	initializeForEnvironment(cfg, logger)

	// Create application context
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Start resource monitoring
	startResourceMonitor(ctx, logger)

	// Initialize memory engine
	logger.Info("Initializing embedding service")
	var embeddingService engine.EmbeddingService

	// Initialize Gemini embedding service if available
	if cfg.Gemini.APIKey != "" {
		logger.Info("Creating Gemini embedding service")
		geminiService := embeddings.NewGeminiEmbeddingService(cfg.Gemini.APIKey)

		// Test the service
		testCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
		defer cancel()

		_, err := geminiService.GenerateEmbedding(testCtx, "test")
		if err != nil {
			logger.Error("Gemini embedding service test failed", "error", err)
			logger.Info("Continuing with fallback embeddings")
			// Create hybrid service as fallback
			embeddingService, err = embeddings.NewHybridEmbeddingService(nil, nil)
			if err != nil {
				logger.Error("Failed to create hybrid embedding service", "error", err)
				os.Exit(1)
			}
		} else {
			logger.Info("Gemini embedding service test successful")
			embeddingService = geminiService
		}
	} else {
		logger.Info("No Gemini API key found, using hybrid embedding service")
		embeddingService, err = embeddings.NewHybridEmbeddingService(nil, nil)
		if err != nil {
			logger.Error("Failed to create hybrid embedding service", "error", err)
			os.Exit(1)
		}
	}

	logger.Info("Initializing memory engine")
	engineConfig := convertToEngineConfig(cfg.ToEngineConfig())
	memoryEngine, err := engine.NewMemoryEngine(engineConfig, embeddingService)
	if err != nil {
		logger.Error("Failed to create memory engine", "error", err)
		os.Exit(1)
	}

	// Start memory engine
	logger.Info("Starting memory engine")
	if err := memoryEngine.Start(ctx); err != nil {
		logger.Error("Failed to start memory engine", "error", err)
		os.Exit(1)
	}

	// Initialize API server if enabled
	var apiServer *api.APIServer
	if cfg.API.Port > 0 {
		logger.Info("Initializing API server", "host", cfg.API.Host, "port", cfg.API.Port)
		apiServer = api.NewAPIServer(memoryEngine, &api.APIConfig{
			Host:           cfg.API.Host,
			Port:           cfg.API.Port,
			EnableCORS:     cfg.API.EnableCORS,
			EnableAuth:     cfg.API.EnableAuth,
			AuthToken:      cfg.API.AuthToken,
			RequestTimeout: cfg.API.RequestTimeout,
			MaxRequestSize: cfg.API.MaxRequestSize,
			EnableMetrics:  cfg.API.EnableMetrics,
			EnableLogging:  cfg.API.EnableLogging,
			EnableTLS:      cfg.API.EnableTLS,
			TLSCertFile:    cfg.API.TLSCertFile,
			TLSKeyFile:     cfg.API.TLSKeyFile,
			ReadTimeout:    cfg.API.ReadTimeout,
			WriteTimeout:   cfg.API.WriteTimeout,
			IdleTimeout:    cfg.API.IdleTimeout,
		})

		// Start API server in goroutine
		go func() {
			logger.Info("Starting API server")
			if err := apiServer.Start(); err != nil {
				logger.Error("API server error", "error", err)
			}
		}()
	}

	// Start metrics server if enabled
	if cfg.Metrics.Enabled {
		go startMetricsServer(cfg, logger)
	}

	// Run demo if in development mode
	if isDevelopment() {
		go runDemo(memoryEngine, logger)
	}

	// Setup graceful shutdown
	setupGracefulShutdown(ctx, cancel, memoryEngine, apiServer, logger)

	// Wait for shutdown signal
	<-ctx.Done()
	logger.Info("Memory Engine stopped")
}

func printBanner() {
	banner := `
███╗   ███╗███████╗███╗   ███╗ ██████╗ ██████╗ ██╗   ██╗
████╗ ████║██╔════╝████╗ ████║██╔═══██╗██╔══██╗╚██╗ ██╔╝
██╔████╔██║█████╗  ██╔████╔██║██║   ██║██████╔╝ ╚████╔╝ 
██║╚██╔╝██║██╔══╝  ██║╚██╔╝██║██║   ██║██╔══██╗  ╚██╔╝  
██║ ╚═╝ ██║███████╗██║ ╚═╝ ██║╚██████╔╝██║  ██║   ██║   
╚═╝     ╚═╝╚══════╝╚═╝     ╚═╝ ╚═════╝ ╚═╝  ╚═╝   ╚═╝   
                                                        
███████╗███╗   ██╗ ██████╗ ██╗███╗   ██╗███████╗       
██╔════╝████╗  ██║██╔════╝ ██║████╗  ██║██╔════╝       
█████╗  ██╔██╗ ██║██║  ███╗██║██╔██╗ ██║█████╗         
██╔══╝  ██║╚██╗██║██║   ██║██║██║╚██╗██║██╔══╝         
███████╗██║ ╚████║╚██████╔╝██║██║ ╚████║███████╗       
╚══════╝╚═╝  ╚═══╝ ╚═════╝ ╚═╝╚═╝  ╚═══╝╚══════╝       
`
	fmt.Print(banner)
	fmt.Printf("\n%s v%s\n", AppName, AppVersion)
	fmt.Printf("%s\n\n", AppDesc)
}

func initializeLogger(cfg *config.Config) Logger {
	// Initialize structured logger based on configuration
	return NewLogger(LoggerConfig{
		Level:          cfg.Logging.Level,
		Format:         cfg.Logging.Format,
		Output:         cfg.Logging.Output,
		EnableFile:     cfg.Logging.EnableFile,
		FilePath:       cfg.Logging.FilePath,
		MaxSize:        cfg.Logging.MaxSize,
		MaxBackups:     cfg.Logging.MaxBackups,
		MaxAge:         cfg.Logging.MaxAge,
		EnableRotation: cfg.Logging.EnableRotation,
		EnableJSON:     cfg.Logging.EnableJSON,
	})
}

func startMetricsServer(cfg *config.Config, logger Logger) {
	logger.Info("Starting metrics server", "port", cfg.Metrics.Port, "path", cfg.Metrics.Path)

	// Implementation would start a metrics server (Prometheus, etc.)
	// For now, just log that it would start
	logger.Info("Metrics server started", "endpoint", fmt.Sprintf("http://localhost:%d%s", cfg.Metrics.Port, cfg.Metrics.Path))
}

func setupGracefulShutdown(ctx context.Context, cancel context.CancelFunc, memoryEngine *engine.MemoryEngine, apiServer *api.APIServer, logger Logger) {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		sig := <-sigChan
		logger.Info("Received shutdown signal", "signal", sig.String())

		// Create shutdown context with timeout
		shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer shutdownCancel()

		// Stop API server first
		if apiServer != nil {
			logger.Info("Stopping API server")
			if err := apiServer.Stop(shutdownCtx); err != nil {
				logger.Error("Error stopping API server", "error", err)
			}
		}

		// Stop memory engine
		logger.Info("Stopping memory engine")
		if err := memoryEngine.Stop(shutdownCtx); err != nil {
			logger.Error("Error stopping memory engine", "error", err)
		}

		// Cancel main context
		cancel()
	}()
}

// Logger interface for structured logging
type Logger interface {
	Info(msg string, keysAndValues ...interface{})
	Error(msg string, keysAndValues ...interface{})
	Warn(msg string, keysAndValues ...interface{})
	Debug(msg string, keysAndValues ...interface{})
}

// LoggerConfig configures the logger
type LoggerConfig struct {
	Level          string
	Format         string
	Output         string
	EnableFile     bool
	FilePath       string
	MaxSize        int
	MaxBackups     int
	MaxAge         int
	EnableRotation bool
	EnableJSON     bool
}

// SimpleLogger implements basic logging
type SimpleLogger struct {
	config LoggerConfig
}

func NewLogger(config LoggerConfig) Logger {
	return &SimpleLogger{config: config}
}

func (l *SimpleLogger) Info(msg string, keysAndValues ...interface{}) {
	l.log("INFO", msg, keysAndValues...)
}

func (l *SimpleLogger) Error(msg string, keysAndValues ...interface{}) {
	l.log("ERROR", msg, keysAndValues...)
}

func (l *SimpleLogger) Warn(msg string, keysAndValues ...interface{}) {
	l.log("WARN", msg, keysAndValues...)
}

func (l *SimpleLogger) Debug(msg string, keysAndValues ...interface{}) {
	if l.config.Level == "debug" {
		l.log("DEBUG", msg, keysAndValues...)
	}
}

func (l *SimpleLogger) log(level, msg string, keysAndValues ...interface{}) {
	timestamp := time.Now().Format("2006-01-02 15:04:05")

	if l.config.EnableJSON {
		// JSON logging format
		logData := map[string]interface{}{
			"timestamp": timestamp,
			"level":     level,
			"message":   msg,
		}

		// Add key-value pairs
		for i := 0; i < len(keysAndValues); i += 2 {
			if i+1 < len(keysAndValues) {
				key := fmt.Sprintf("%v", keysAndValues[i])
				value := keysAndValues[i+1]
				logData[key] = value
			}
		}

		// Simple JSON output (in production, use a proper JSON logger)
		fmt.Printf("{\"timestamp\":\"%s\",\"level\":\"%s\",\"message\":\"%s\"", timestamp, level, msg)
		for i := 0; i < len(keysAndValues); i += 2 {
			if i+1 < len(keysAndValues) {
				key := fmt.Sprintf("%v", keysAndValues[i])
				value := keysAndValues[i+1]
				fmt.Printf(",\"%s\":\"%v\"", key, value)
			}
		}
		fmt.Println("}")
	} else {
		// Human-readable format
		fmt.Printf("[%s] %s: %s", timestamp, level, msg)

		// Add key-value pairs
		for i := 0; i < len(keysAndValues); i += 2 {
			if i+1 < len(keysAndValues) {
				key := keysAndValues[i]
				value := keysAndValues[i+1]
				fmt.Printf(" %v=%v", key, value)
			}
		}
		fmt.Println()
	}
}

// Demo function to show basic usage
func runDemo(memoryEngine *engine.MemoryEngine, logger Logger) {
	logger.Info("Running Memory Engine demo")

	ctx := context.Background()

	// Example 1: Store some sample content
	sampleChunk := &memory.MemoryChunk{
		ID:      "demo_chunk_1",
		Content: "The Memory Engine is a high-performance semantic memory system with B+tree indexing and intelligent chunking capabilities.",
		Metadata: map[string]interface{}{
			"source":   "demo",
			"category": "technical",
			"author":   "system",
		},
		Timestamp:  time.Now(),
		SourceType: memory.ContentTypeText,
		Language:   "en",
		Quality:    0.9,
	}

	logger.Info("Storing sample chunk", "id", sampleChunk.ID)
	if err := memoryEngine.Store(ctx, sampleChunk); err != nil {
		logger.Error("Failed to store chunk", "error", err)
		return
	}

	// Example 2: Search for content
	searchQuery := &memory.SearchQuery{
		Query:        "semantic memory system",
		MaxResults:   5,
		MinRelevance: 0.3,
	}

	logger.Info("Searching for content", "query", searchQuery.Query)
	results, err := memoryEngine.Retrieve(ctx, searchQuery)
	if err != nil {
		logger.Error("Search failed", "error", err)
		return
	}

	logger.Info("Search completed", "results_count", len(results))
	for i, result := range results {
		logger.Info("Search result",
			"index", i+1,
			"chunk_id", result.Chunk.ID,
			"score", result.CombinedScore,
			"content_preview", truncateString(result.Chunk.Content, 100))
	}

	// Example 3: Build context from results
	if len(results) > 0 {
		chunks := make([]*memory.MemoryChunk, len(results))
		for i, result := range results {
			chunks[i] = result.Chunk
		}

		context := memoryEngine.BuildContext(ctx, chunks, 1000)
		logger.Info("Built context", "context", context)
	}

	// Example 4: Show engine statistics
	stats := memoryEngine.GetStats()
	logger.Info("Engine statistics",
		"total_operations", stats.TotalOperations,
		"successful_ops", stats.SuccessfulOps,
		"average_latency", stats.AverageLatency,
		"uptime", stats.Uptime)

	logger.Info("Demo completed successfully")
}

func truncateString(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}

// Health check endpoint for container orchestration
func healthCheck(memoryEngine *engine.MemoryEngine) error {
	status := memoryEngine.GetStatus()
	if status != engine.StatusRunning {
		return fmt.Errorf("engine status: %s", status.String())
	}

	// Additional health checks could go here
	return nil
}

// Development mode helpers
func isDevelopment() bool {
	return os.Getenv("MEMORY_ENGINE_ENVIRONMENT") == "development" ||
		os.Getenv("ENVIRONMENT") == "development" ||
		os.Getenv("ENV") == "dev"
}

func isProduction() bool {
	return os.Getenv("MEMORY_ENGINE_ENVIRONMENT") == "production" ||
		os.Getenv("ENVIRONMENT") == "production" ||
		os.Getenv("ENV") == "prod"
}

// Environment-specific initialization
func initializeForEnvironment(cfg *config.Config, logger Logger) {
	switch cfg.App.Environment {
	case "development":
		logger.Info("Running in development mode")
		// Enable debug features, additional logging, etc.

	case "staging":
		logger.Info("Running in staging mode")
		// Enable some debug features but with production-like settings

	case "production":
		logger.Info("Running in production mode")
		// Optimize for performance, disable debug features

	default:
		logger.Warn("Unknown environment, defaulting to development", "environment", cfg.App.Environment)
	}
}

// Version information
func printVersion() {
	fmt.Printf("%s version %s\n", AppName, AppVersion)
	fmt.Printf("Build info: Go %s\n", "1.21+")
	fmt.Printf("Features: B+Tree indexing, Python bridge, Langfuse evaluation\n")
}

// Configuration validation
func validateEnvironment(cfg *config.Config, logger Logger) error {
	// Check required directories exist
	dirs := []string{cfg.App.DataPath, cfg.App.TempPath, cfg.App.LogPath}
	for _, dir := range dirs {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return fmt.Errorf("failed to create directory %s: %w", dir, err)
		}
	}

	// Validate Python environment if Python bridge is enabled
	if cfg.Engine.EnablePythonBridge {
		logger.Info("Validating Python environment", "python_path", cfg.Python.PythonPath)
		// Add Python validation logic here
	}

	// Validate Langfuse configuration if evaluation is enabled
	if cfg.Engine.EnableEvaluation && cfg.Evaluation.LangfuseAPIKey != "" {
		logger.Info("Langfuse integration enabled", "host", cfg.Evaluation.LangfuseHost)
	}

	return nil
}

// Resource monitoring
func startResourceMonitor(ctx context.Context, logger Logger) {
	ticker := time.NewTicker(time.Minute * 5)
	defer ticker.Stop()

	go func() {
		for {
			select {
			case <-ticker.C:
				// Monitor system resources
				// This would typically use a proper monitoring library
				logger.Debug("Resource monitor tick")

			case <-ctx.Done():
				return
			}
		}
	}()
}

// convertToEngineConfig converts CompleteEngineConfig to engine.EngineConfig
func convertToEngineConfig(completeConfig *config.CompleteEngineConfig) *engine.EngineConfig {
	return &engine.EngineConfig{
		Name:                completeConfig.Name,
		Version:             completeConfig.Version,
		DataPath:            completeConfig.DataPath,
		LogLevel:            completeConfig.LogLevel,
		EnableMetrics:       completeConfig.EnableMetrics,
		MetricsPort:         completeConfig.MetricsPort,
		MaxConcurrentOps:    completeConfig.MaxConcurrentOps,
		OperationTimeout:    completeConfig.OperationTimeout,
		HealthCheckInterval: completeConfig.HealthCheckInterval,
		EnableEvaluation:    completeConfig.EnableEvaluation,
		EnablePythonBridge:  completeConfig.EnablePythonBridge,
		EnableAutoOptimize:  completeConfig.EnableAutoOptimize,
		EnableHealthCheck:   completeConfig.EnableHealthCheck,
		EnablePlugins:       completeConfig.EnablePlugins,
		PluginDirectory:     completeConfig.PluginDirectory,
		StorageConfig:       completeConfig.StorageConfig,
		IngestionConfig:     completeConfig.IngestionConfig,
		RetrievalConfig:     completeConfig.RetrievalConfig,
		PythonConfig:        completeConfig.PythonConfig,
		EvaluationConfig:    completeConfig.EvaluationConfig,
	}
}