FloorVisualizer/main.go

package main

import (
	"context"
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"time"

	"github.com/joho/godotenv"

	"floorvisualizer/internal/api"
	"floorvisualizer/internal/handler"
	"floorvisualizer/internal/logger"
	"floorvisualizer/internal/middleware"
	"floorvisualizer/internal/model"
	"floorvisualizer/internal/openrouter"
	"floorvisualizer/internal/queue"
	"floorvisualizer/internal/repository"
	"floorvisualizer/internal/service"
)

const defaultAPIKey = "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"
const defaultJWTSecret = "floorviz-dev-secret-change-in-production"

var proxyURL = "http://127.0.0.1:7897"

func init() {
	if v := os.Getenv("PROXY_URL"); v != "" {
		proxyURL = v
	}
	if v := os.Getenv("DISABLE_PROXY"); v == "true" || v == "1" {
		proxyURL = ""
	}
}

func main() {
	godotenv.Load()

	// Load material description index
	if err := handler.LoadMaterialIndex("outputs/material_index.json"); err != nil {
		logger.Warn("Material index not loaded: %v", err)
	} else {
		logger.Info("Material index loaded")
	}

	// Start CLIP server for content check
	clipCmd := startCLIPServer()
	if clipCmd != nil {
		defer stopCLIPServer(clipCmd)
	}

	logDir := "logs"
	if v := os.Getenv("LOG_DIR"); v != "" {
		logDir = v
	}
	logLevel := "INFO"
	if v := os.Getenv("LOG_LEVEL"); v != "" {
		logLevel = v
	}
	logger.Init(logDir, logLevel)
	defer logger.Close()

	dsn := repository.DefaultDSN
	if v := os.Getenv("DATABASE_URL"); v != "" {
		dsn = v
	}
	db, err := repository.Connect(dsn)
	if err != nil {
		logger.Warn("Database not available: %v", err)
	} else {
		defer db.Close()
		logger.Info("Database connected (tables managed via schema.sql)")
		if err := repository.InitGORM(db); err != nil {
			logger.Warn("GORM init failed, falling back to raw SQL: %v", err)
		} else {
			logger.Info("GORM initialized")
		}
	}

	jwtSecret := defaultJWTSecret
	if s := os.Getenv("JWT_SECRET"); s != "" {
		jwtSecret = s
	}
	authSvc := service.NewAuthService(jwtSecret)

	apiKey := defaultAPIKey
	if k := os.Getenv("OPENROUTER_API_KEY"); k != "" {
		apiKey = k
	}
	client := openrouter.NewClient(apiKey)
	if proxyURL != "" {
		client = openrouter.NewClient(apiKey, openrouter.WithHTTPClient(handler.NewProxyClient(proxyURL)))
	}

	// ── Redis ──────────────────────────────────────────────
	redisAddr := "localhost:6379"
	if v := os.Getenv("REDIS_ADDR"); v != "" {
		redisAddr = v
	}
	var cancelWorkers context.CancelFunc
	rdb, err := queue.NewRedisClient(redisAddr, "")
	if err != nil {
		logger.Warn("Redis not available (%v), worker queue disabled", err)
	} else {
		defer rdb.Close()
		logger.Info("Redis connected")

		// Wire up the queue: generate handler enqueues to Redis
		handler.SetQueue(func(ctx context.Context, jobID string, payload map[string]string) error {
			if err := queue.StoreJobPayload(ctx, rdb, jobID, payload); err != nil {
				return err
			}
			return queue.EnqueueJob(ctx, rdb, jobID)
		})

		// Wire up status query
		handler.SetStatusQuery(func(ctx context.Context, jobID string) (*handler.JobStatusResp, error) {
			js, err := queue.GetJobStatus(ctx, rdb, jobID)
			if err != nil {
				return nil, err
			}
			return &handler.JobStatusResp{
				JobID: js.JobID, Status: js.Status, Step: js.Step,
				Message: js.Message, Result: js.Result, Error: js.Error,
			}, nil
		})

		// Start worker pool (default 5, override with WORKER_COUNT env)
		workerCount := 5
		if v := os.Getenv("WORKER_COUNT"); v != "" {
			if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 20 {
				workerCount = n
			}
		}
		workerCtx, cancel := context.WithCancel(context.Background())
		cancelWorkers = cancel
		defer cancelWorkers()
		queue.StartWorkers(workerCtx, workerCount, rdb, client, db)
		logger.Info("Worker pool started: %d workers", workerCount)
	}

	// ── Hot product cache ────────────────────────────────
	handler.SetProductsCache(
		// Getter: check Redis cache for hot products
		func(ctx context.Context, sku string) (*model.Product, bool) {
			val, err := rdb.Get(ctx, "product:"+sku).Result()
			if err != nil {
				return nil, false
			}
			var p model.Product
			if json.Unmarshal([]byte(val), &p) == nil {
				return &p, true
			}
			return nil, false
		},
		// Setter: only cache if this SKU is in the hot list (top 200)
		func(ctx context.Context, sku string, product *model.Product) error {
			rank, err := rdb.ZRevRank(ctx, "product:views", sku).Result()
			if err != nil || rank >= 200 {
				return nil // not hot enough, skip caching
			}
			data, _ := json.Marshal(product)
			ttl := 1*time.Hour + time.Duration(rand.Intn(600))*time.Second // ±10min jitter
			return rdb.Set(ctx, "product:"+sku, data, ttl).Err()
		},
		// View tracker: increment view count in sorted set
		func(ctx context.Context, sku string) {
			rdb.ZIncrBy(ctx, "product:views", 1, sku)
		},
	)

	mux := http.NewServeMux()
	api.RegisterRoutes(mux, client, db, authSvc)

	port := "8099"
	if p := os.Getenv("PORT"); p != "" {
		port = p
	}

	server := &http.Server{
		Addr:    ":" + port,
		Handler: middleware.Logging(mux),
	}

	// Graceful shutdown: Ctrl+C stops HTTP server + worker pool
	go func() {
		sig := make(chan os.Signal, 1)
		signal.Notify(sig, os.Interrupt)
		<-sig
		logger.Info("Shutting down...")
		if cancelWorkers != nil {
			cancelWorkers()
		}
		stopCLIPServer(clipCmd)
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		if err := server.Shutdown(shutdownCtx); err != nil {
			logger.Warn("HTTP server forced shutdown: %v", err)
		}
	}()

	logger.Info("FloorVisualizer -> http://localhost:%s", port)
	if err := server.ListenAndServe(); err != http.ErrServerClosed {
		log.Fatal(err)
	}
	logger.Info("Server stopped gracefully")
}

// startCLIPServer launches the Python CLIP server for content checking.
func startCLIPServer() *exec.Cmd {
	// Check if CLIP is already running
	if _, err := http.Get("http://127.0.0.1:5100/health"); err == nil {
		logger.Info("CLIP server already running on :5100")
		return nil
	}
	pythonPath := `C:\Users\10794\AppData\Local\Programs\Python\Python312\python.exe`
	cmd := exec.Command(pythonPath, "clip_server/server.py")
	cmd.Stdout = nil
	cmd.Stderr = nil
	if err := cmd.Start(); err != nil {
		logger.Warn("Failed to start CLIP server: %v (content check will use Gemini fallback)", err)
		return nil
	}
	logger.Info("CLIP server starting on :5100...")
	// Wait for it to be ready
	for i := 0; i < 30; i++ {
		time.Sleep(time.Second)
		if _, err := http.Get("http://127.0.0.1:5100/health"); err == nil {
			logger.Info("CLIP server ready")
			return cmd
		}
	}
	logger.Warn("CLIP server did not become ready in time")
	return cmd
}

// stopCLIPServer stops the CLIP Python server.
func stopCLIPServer(cmd *exec.Cmd) {
	if cmd == nil || cmd.Process == nil {
		return
	}
	logger.Info("Stopping CLIP server...")
	if err := cmd.Process.Kill(); err != nil {
		logger.Warn("Failed to stop CLIP server: %v", err)
	}
}