package queue import ( "bytes" "context" "database/sql" "encoding/json" "floorvisualizer/internal/logger" "fmt" "io" "net/http" "os" "path/filepath" "strings" "time" "github.com/redis/go-redis/v9" "floorvisualizer/internal/handler" "floorvisualizer/internal/openrouter" ) // StartWorkers launches n goroutines that consume jobs from the Redis queue. func StartWorkers(ctx context.Context, n int, rdb *redis.Client, client *openrouter.Client, db *sql.DB) { for i := 0; i < n; i++ { go func(workerID int) { logger.Info("[worker %d] started", workerID) for { select { case <-ctx.Done(): logger.Info("[worker %d] shutting down", workerID) return default: } jobID, err := DequeueJob(ctx, rdb) if err != nil { if ctx.Err() != nil { return } logger.Error("[worker %d] dequeue error: %v", workerID, err) time.Sleep(time.Second) continue } logger.Info("[worker %d] picked job %s", workerID, jobID) processJob(ctx, rdb, client, db, jobID) } }(i + 1) } } func processJob(ctx context.Context, rdb *redis.Client, client *openrouter.Client, db *sql.DB, jobID string) { totalStart := time.Now() SetJobStatus(ctx, rdb, jobID, "processing") logger.Info("[%s] TIMER job start", jobID) // Load payload payload, err := GetJobPayload(ctx, rdb, jobID) if err != nil { SetJobError(ctx, rdb, jobID, "Failed to load job payload") return } inputPath := payload["input_path"] floorID := payload["floor_id"] patternCode := payload["pattern_code"] roomCode := payload["room_code"] // Look up floor and pattern from DB floor := handler.FindFloorBySKUPublic(db, floorID) pattern := handler.FindPatternByCodePublic(patternCode) room := handler.FindRoomByCodePublic(roomCode) groutLines := payloadBool(payload, "grout_lines") whiteBaseboard := payloadBool(payload, "white_baseboard") materialSource := payload["material_source"] // "image", "json", or "both" if materialSource == "" { materialSource = "image" } if floor == nil || pattern == nil { SetJobError(ctx, rdb, jobID, "Invalid floor or pattern") return } logger.Info("[%s] TIMER options floor_id=%s pattern_code=%s room_code=%s elapsed=%s", jobID, floorID, patternCode, roomCode, time.Since(totalStart).Round(time.Millisecond)) // Resize input to max 2048px to speed up AI processing workPath := filepath.Join("uploads", fmt.Sprintf("input_%s_resized.png", jobID)) if err := openrouter.ResizeImageToFit(inputPath, workPath); err != nil { logger.Warn("[%s] Failed to resize input, using original: %v", jobID, err) workPath = inputPath } outputPath := filepath.Join("outputs", fmt.Sprintf("result_%s.png", jobID)) jobCtx, cancel := context.WithTimeout(ctx, 600*time.Second) defer cancel() // Step 0: content check SetJobProgress(jobCtx, rdb, jobID, "check", "Checking content...") stepStart := time.Now() logger.Info("[%s] TIMER step=check start elapsed=%s", jobID, time.Since(totalStart).Round(time.Millisecond)) passed, score := validateContent(client, jobCtx, workPath) logStepDone(jobCtx, rdb, jobID, "check", stepStart, totalStart, "score=%d passed=%t", score, passed) if score == 0 { logger.Warn("[%s] Content check failed (score=0)", jobID) } if !passed { SetJobError(jobCtx, rdb, jobID, "Content check failed: Please upload an indoor room photo") return } // Step 1: generate floor (single combined prompt — AI identifies floor + replaces it) roomLabel := "" if room != nil { roomLabel = " · " + room.Name } SetJobProgress(jobCtx, rdb, jobID, "generate", fmt.Sprintf("Generating %s %s%s...", floor.Name, pattern.Name, roomLabel)) stepStart = time.Now() logger.Info("[%s] TIMER step=generate start floor=%s pattern=%s elapsed=%s", jobID, floor.Name, pattern.Name, time.Since(totalStart).Round(time.Millisecond)) // Material source: image / json / both if materialSource == "json" || materialSource == "both" { if desc := handler.LookupMaterialDesc(floor.SKU); desc != "" { if floor.Description != "" { floor.Description = desc + " | " + floor.Description } else { floor.Description = desc } } } refImagePath := "" if materialSource != "json" { refImagePath = downloadReferenceImage(jobCtx, floor.ImageURL, jobID) } imageSize := floorImageSize() logger.Info("[%s] TIMER generate request material=%s image_size=%s input=%s output=%s ref=%s", jobID, materialSource, imageSize, workPath, outputPath, refImagePath) floorPrompt := handler.BuildCombinedFloorPrompt(*floor, *pattern, room, groutLines, whiteBaseboard) _, err = client.GenerateImageToFile(jobCtx, openrouter.ImageGenerationRequest{ Prompt: floorPrompt, OutputPath: outputPath, InputImagePath: workPath, ImageSize: imageSize, ReferenceImagePath: refImagePath, }) if err != nil { logStepError(jobCtx, rdb, jobID, "generate", totalStart, stepStart, err) SetJobError(jobCtx, rdb, jobID, "Floor generation failed: "+err.Error()) return } if err := openrouter.RepairNearBlackBorder(outputPath); err != nil { logger.Warn("[%s] Failed to repair output border: %v", jobID, err) } logStepDone(jobCtx, rdb, jobID, "generate", stepStart, totalStart, "") // Clean up intermediate files stepStart = time.Now() logger.Info("[%s] TIMER step=cleanup start elapsed=%s", jobID, time.Since(totalStart).Round(time.Millisecond)) if err := os.Remove(inputPath); err != nil { logger.Warn("[%s] Failed to remove input: %v", jobID, err) } if workPath != inputPath { if err := os.Remove(workPath); err != nil { logger.Warn("[%s] Failed to remove resized input: %v", jobID, err) } } if refImagePath != "" { if err := os.Remove(refImagePath); err != nil { logger.Warn("[%s] Failed to remove reference image: %v", jobID, err) } } totalDuration := time.Since(totalStart) _ = SetJobTiming(jobCtx, rdb, jobID, "total", totalDuration, totalDuration) SetJobResult(jobCtx, rdb, jobID, "/"+filepath.ToSlash(outputPath)) logger.Info("[%s] TIMER job done total=%s result=%s", jobID, totalDuration.Round(time.Millisecond), "/"+filepath.ToSlash(outputPath)) } func logStepDone(ctx context.Context, rdb *redis.Client, jobID, step string, stepStart, totalStart time.Time, format string, args ...any) { duration := time.Since(stepStart) elapsed := time.Since(totalStart) _ = SetJobTiming(ctx, rdb, jobID, step, duration, elapsed) if format != "" { logger.Info("[%s] TIMER step=%s done duration=%s elapsed=%s %s", jobID, step, duration.Round(time.Millisecond), elapsed.Round(time.Millisecond), fmt.Sprintf(format, args...)) return } logger.Info("[%s] TIMER step=%s done duration=%s elapsed=%s", jobID, step, duration.Round(time.Millisecond), elapsed.Round(time.Millisecond)) } func logStepError(ctx context.Context, rdb *redis.Client, jobID, step string, totalStart, stepStart time.Time, err error) { duration := time.Since(stepStart) elapsed := time.Since(totalStart) _ = SetJobTiming(ctx, rdb, jobID, step, duration, elapsed) logger.Error("[%s] TIMER step=%s error duration=%s elapsed=%s err=%v", jobID, step, duration.Round(time.Millisecond), elapsed.Round(time.Millisecond), err) } // validateContent checks if the image is an indoor room photo. // Uses local CLIP ViT-B/32 first, falls back to Gemini 2.5 Flash if CLIP is unavailable. func validateContent(geminiClient *openrouter.Client, ctx context.Context, imagePath string) (bool, int) { imgData, err := os.ReadFile(imagePath) if err != nil { return false, 0 } // Try local CLIP server first if score, ok := clipCheck(imgData); ok { logger.Info("Content check via local CLIP: score=%d passed=%t", score, score >= 6) return score >= 6, score } // Fallback to Gemini checkCtx, cancel := context.WithTimeout(ctx, 20*time.Second) defer cancel() dataURI := "data:image/png;base64," + base64enc(imgData) resp, err := geminiClient.Chat(checkCtx, openrouter.ChatRequest{ Model: "google/gemini-2.5-flash", Messages: []openrouter.Message{{Role: "user", Content: []openrouter.ContentPart{ {Type: "image_url", ImageURL: &openrouter.ImageURL{URL: dataURI}}, {Type: "text", Text: `请按 1-10 分评估这张图片是否适合做室内地面替换。 8-10 分:清晰室内房间,地面可见。 5-7 分:室内图片,但角度或地面可见度一般。 1-4 分:室外、宠物特写、物体特写、没有可替换地面,或不适合做地面替换。 只返回一个数字(1-10)。`}, }}}, }) if err != nil || len(resp.Choices) == 0 { return true, 0 } score := parseScore(resp.Choices[0].Message.Content) return score >= 6, score } // clipCheck sends image bytes to the local CLIP server and returns (score, ok). func clipCheck(imgData []byte) (int, bool) { req, err := http.NewRequest(http.MethodPost, "http://127.0.0.1:5100/check", bytes.NewReader(imgData)) if err != nil { return 0, false } req.Header.Set("Content-Type", "application/octet-stream") resp, err := http.DefaultClient.Do(req) if err != nil { return 0, false } defer resp.Body.Close() if resp.StatusCode != 200 { return 0, false } var result struct { Score int `json:"score"` Passed bool `json:"passed"` } if json.NewDecoder(resp.Body).Decode(&result) != nil { return 0, false } return result.Score, true } func parseScore(raw string) int { for _, s := range []string{"10", "9", "8", "7", "6", "5", "4", "3", "2", "1"} { if len(raw) >= len(s) && raw[:len(s)] == s { v := 0 fmt.Sscanf(s, "%d", &v) return v } } return 0 } func base64enc(data []byte) string { const tbl = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/" var b []byte for i := 0; i < len(data); i += 3 { b0, b1, b2 := data[i], byte(0), byte(0) if i+1 < len(data) { b1 = data[i+1] } if i+2 < len(data) { b2 = data[i+2] } b = append(b, tbl[b0>>2]) b = append(b, tbl[((b0&3)<<4)|(b1>>4)]) if i+1 < len(data) { b = append(b, tbl[((b1&15)<<2)|(b2>>6)]) } else { b = append(b, '=') } if i+2 < len(data) { b = append(b, tbl[b2&63]) } else { b = append(b, '=') } } return string(b) } // downloadReferenceImage downloads the floor material image from the given URL // and saves it to a local temp file. Returns the local file path, or empty string on failure. func downloadReferenceImage(ctx context.Context, imageURL, jobID string) string { if imageURL == "" { return "" } refPath := filepath.Join("outputs", fmt.Sprintf("ref_%s.png", jobID)) req, err := http.NewRequestWithContext(ctx, http.MethodGet, imageURL, nil) if err != nil { logger.Warn("[%s] Failed to create reference image request: %v", jobID, err) return "" } resp, err := http.DefaultClient.Do(req) if err != nil { logger.Warn("[%s] Failed to download reference image from %s: %v", jobID, imageURL, err) return "" } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { logger.Warn("[%s] Reference image download returned HTTP %d from %s", jobID, resp.StatusCode, imageURL) return "" } f, err := os.Create(refPath) if err != nil { logger.Warn("[%s] Failed to create reference image file: %v", jobID, err) return "" } defer f.Close() if _, err := io.Copy(f, resp.Body); err != nil { logger.Warn("[%s] Failed to save reference image: %v", jobID, err) return "" } logger.Info("[%s] Downloaded reference image from %s", jobID, imageURL) return refPath } func payloadBool(payload map[string]string, key string) bool { switch payload[key] { case "1", "true", "yes", "y", "on": return true default: return false } } func floorImageSize() string { value := strings.TrimSpace(os.Getenv("FLOOR_IMAGE_SIZE")) if value == "" { return "2K" } return value }