FloorVisualizer/internal/queue/worker.go

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)
	groutColor := strings.ToLower(strings.TrimSpace(payload["grout_lines"]))
	baseboardColor := strings.ToLower(strings.TrimSpace(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, groutColor, baseboardColor)
	_, 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
}