FloorVisualizer/main.go
dindang 5e35532dbb feat(floor): 优化地板样式数据结构和AI换地板功能
- 添加 FloorOption 结构体字段支持尺寸、描述和变体信息
- 优化 loadFloorOptions 函数,实现按品牌和花色名称去重并附带多尺寸变体
- 区分木材、瓷砖、乙烯基和层压板产品,并分类返回
- 更新室内铺装纹理选项,完善图案描述和名称本地化
- 丰富Floor Options API响应数据,包含尺寸、变体、描述等字段
- 增强AI换地板功能,支持尺寸变体SKU,自动调整物理尺寸比例
- 完善地板替换的AI提示词,增加材质锁定、尺寸说明和纹理一致性要求
- 改进地板识别蒙版生成逻辑,确保精准分割地板区域
- Redis任务状态查询接口增加耗时统计字段,提供更细粒度进度信息
- 更新go.mod依赖,新增redis和gorm相关包支持
2026-07-22 15:23:40 +08:00

199 lines
5.4 KiB
Go

package main
import (
"context"
"encoding/json"
"log"
"math/rand"
"net/http"
"os"
"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()
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()
}
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")
}