package handler import ( "context" "database/sql" "fmt" "io" "net/http" "os" "path/filepath" "strings" "time" "floorvisualizer/internal/openrouter" "floorvisualizer/internal/repository" ) // ── Floor options ────────────────────────────────────────── type FloorOption struct { ID string `json:"id"` Name string `json:"name"` Category string `json:"category"` Material string `json:"material"` ColorTone string `json:"color_tone"` Finish string `json:"finish"` SizeLabel string `json:"size_label"` ImageURL string `json:"image_url"` SKU string `json:"sku"` } type PatternOption struct { Code int `json:"code"` Name string `json:"name"` Pattern string `json:"pattern"` Category string `json:"category"` } type RoomOption struct { Code int `json:"code"` Name string `json:"name"` } var woodPatterns = []PatternOption{ {Code: 1, Name: "Straight Lay", Category: "wood", Pattern: "straight/parallel planks running from front to back with staggered joints"}, {Code: 2, Name: "Horizontal Lay", Category: "wood", Pattern: "horizontal planks running side to side across the width of the room"}, {Code: 3, Name: "Herringbone", Category: "wood", Pattern: "herringbone/zigzag pattern with rectangular planks arranged in a broken V shape"}, {Code: 4, Name: "Chevron", Category: "wood", Pattern: "chevron pattern with planks cut at an angle forming a continuous V shape"}, } var tilePatterns = []PatternOption{ {Code: 5, Name: "Straight Grid", Category: "tile", Pattern: "straight grid layout with tiles perfectly aligned in rows and columns, clean orthogonal grout lines"}, {Code: 6, Name: "Running Bond", Category: "tile", Pattern: "running bond/brick pattern with each row offset by 50%, staggered brickwork appearance"}, {Code: 7, Name: "1/3 Offset", Category: "tile", Pattern: "tiles offset by 1/3 of their length in each row, subtle stepped pattern"}, {Code: 8, Name: "Hexagonal", Category: "tile", Pattern: "hexagonal/honeycomb tiles fitted together in a seamless honeycomb mesh"}, {Code: 9, Name: "Diagonal (45deg)", Category: "tile", Pattern: "tiles laid at a 45-degree diagonal angle to the walls, diamond orientation"}, } var roomOptions = []RoomOption{ {Code: 0, Name: "Auto Detect"}, {Code: 1, Name: "Living Room"}, {Code: 2, Name: "Bedroom"}, {Code: 3, Name: "Kitchen"}, {Code: 4, Name: "Dining Room"}, {Code: 5, Name: "Bathroom"}, {Code: 6, Name: "Study Room"}, {Code: 7, Name: "Hallway"}, } func findPatternByCode(code int) *PatternOption { for _, p := range append(woodPatterns, tilePatterns...) { if p.Code == code { return &p } } return nil } func findRoomByCode(code int) *RoomOption { if code == 0 { return nil } for _, r := range roomOptions { if r.Code == code { return &r } } return nil } // loadFloorOptions loads floor options from the database, one per style_name per brand. func loadFloorOptions(db *sql.DB, category string) ([]FloorOption, error) { prods, _, err := repository.QueryFilteredProducts(db, repository.FilterParams{ Category: category, Limit: 200, }) if err != nil { return nil, err } // Dedup by brand+style_name, take first as representative seen := map[string]bool{} var out []FloorOption for _, p := range prods { key := p.Brand + "|" + p.StyleName if seen[key] { continue } seen[key] = true out = append(out, FloorOption{ ID: p.SKU, Name: p.StyleName, Category: p.Category, Material: p.Material, ColorTone: p.ColorTone, Finish: p.Finish, SizeLabel: p.SizeLabel, ImageURL: p.MainImageURL, SKU: p.SKU, }) } return out, nil } func findFloorBySKU(db *sql.DB, sku string) *FloorOption { p, err := repository.GetProductBySKU(db, sku) if err != nil { return nil } return &FloorOption{ ID: p.SKU, Name: p.StyleName, Category: p.Category, Material: p.Material, ColorTone: p.ColorTone, Finish: p.Finish, SizeLabel: p.SizeLabel, ImageURL: p.MainImageURL, SKU: p.SKU, } } // ── Public accessors for worker ──────────────────────────── func FindFloorBySKUPublic(db *sql.DB, sku string) *FloorOption { return findFloorBySKU(db, sku) } func FindPatternByCodePublic(code string) *PatternOption { c := 0 fmt.Sscanf(code, "%d", &c) return findPatternByCode(c) } func FindRoomByCodePublic(code string) *RoomOption { c := 0 fmt.Sscanf(code, "%d", &c) return findRoomByCode(c) } func BuildMaskPromptPublic(room *RoomOption) string { return buildMaskPrompt(room) } func BuildFloorPromptPublic(f FloorOption, p PatternOption, room *RoomOption) string { return buildFloorPrompt(f, p, room) } // ── Floor Generate Handler (Redis queue) ────────────────── func SetQueue(enqueue func(ctx context.Context, jobID string, payload map[string]string) error) { generateQueue = enqueue } var generateQueue func(ctx context.Context, jobID string, payload map[string]string) error var queryStatus func(ctx context.Context, jobID string) (*JobStatusResp, error) func SetStatusQuery(fn func(ctx context.Context, jobID string) (*JobStatusResp, error)) { queryStatus = fn } type JobStatusResp struct { JobID string `json:"job_id"` Status string `json:"status"` Step string `json:"step"` Message string `json:"message"` Result string `json:"result,omitempty"` Error string `json:"error,omitempty"` } func FloorGenerate(client *openrouter.Client, db *sql.DB) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { writeErr(w, 405, "Method not allowed") return } if err := r.ParseMultipartForm(20 << 20); err != nil { writeErr(w, 400, err.Error()) return } file, _, err := r.FormFile("image") if err != nil { writeErr(w, 400, "No image") return } defer file.Close() floor := findFloorBySKU(db, r.FormValue("floor_id")) patternCode := 0 fmt.Sscanf(r.FormValue("pattern_code"), "%d", &patternCode) pattern := findPatternByCode(patternCode) if floor == nil || pattern == nil { writeErr(w, 400, "Invalid options") return } os.MkdirAll("uploads", 0755) jobID := fmt.Sprintf("%d", time.Now().UnixNano()) inputPath := filepath.Join("uploads", fmt.Sprintf("input_%s.png", jobID)) dst, err := os.Create(inputPath) if err != nil { writeErr(w, 500, "Failed to save image") return } if _, err := io.Copy(dst, file); err != nil { dst.Close() writeErr(w, 500, "Failed to write image") return } dst.Close() os.MkdirAll("outputs", 0755) if generateQueue == nil { writeErr(w, 500, "Queue not initialized") return } // Store job payload then enqueue payload := map[string]string{ "input_path": inputPath, "floor_id": r.FormValue("floor_id"), "pattern_code": r.FormValue("pattern_code"), "room_code": r.FormValue("room_code"), } if err := generateQueue(r.Context(), jobID, payload); err != nil { writeErr(w, 500, "Failed to enqueue job") return } _ = payload // payload stored inside generateQueue wrapper writeJSON(w, 200, map[string]string{"job_id": jobID}) } } // ── Progress ─────────────────────────────────────────────── // ── Floor Status Query ───────────────────────────────── func FloorStatus(w http.ResponseWriter, r *http.Request) { jobID := r.URL.Query().Get("job_id") if jobID == "" { writeErr(w, 400, "job_id is required") return } if queryStatus == nil { writeErr(w, 500, "status query not initialized") return } js, err := queryStatus(r.Context(), jobID) if err != nil { writeJSON(w, 404, map[string]string{"error": "job not found"}) return } writeJSON(w, 200, js) } // ── Options ──────────────────────────────────────────────── func FloorOptions(db *sql.DB) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { woodFloors, _ := loadFloorOptions(db, "Hardwood") // non-critical: returns empty on error woodFloors2, _ := loadFloorOptions(db, "Engineered Wood") tileFloors, _ := loadFloorOptions(db, "Wood-Look Tile") vinylFloors, _ := loadFloorOptions(db, "SPC/LVP") laminateFloors, _ := loadFloorOptions(db, "Laminate") allWood := append(woodFloors, woodFloors2...) allTile := append(tileFloors, laminateFloors...) writeJSON(w, 200, map[string]any{ "wood_floors": allWood, "wood_patterns": woodPatterns, "tile_floors": allTile, "tile_patterns": tilePatterns, "vinyl_floors": vinylFloors, "rooms": roomOptions, }) } } // ── Prompts ──────────────────────────────────────────────── func buildMaskPrompt(room *RoomOption) string { base := `Create a pure black-and-white mask image for this room photo. WHITE (#FFFFFF) areas: ONLY the flat horizontal floor. BLACK (#000000) areas — everything else: stairs, baseboards, walls, furniture, doors, windows, ceiling, decorations. Output ONLY the mask image. No text.` if room != nil && room.Code != 0 { base += fmt.Sprintf("\nContext: This is a %s.", room.Name) } return base } func buildFloorPrompt(floor FloorOption, pattern PatternOption, room *RoomOption) string { var roomCtx string if room != nil && room.Code != 0 { roomCtx = fmt.Sprintf("\nRoom context: This is a %s.", room.Name) } matType := "wood plank" isTile := floor.Category == "Wood-Look Tile" || floor.Category == "Laminate" if isTile { matType = "tile" } materialDesc := fmt.Sprintf("%s material, %s tone, %s finish", floor.Material, floor.ColorTone, floor.Finish) if floor.SizeLabel != "" { materialDesc += fmt.Sprintf(", %s size", floor.SizeLabel) } return fmt.Sprintf( `Replace ONLY the white-masked floor area with %s (%s) %s flooring in a %s pattern.%s Material details: %s | Layout: %s Style: photorealistic, soft natural matte finish, minimal reflections, no mirror-like gloss, seamless blend with room lighting.`, floor.Name, matType, materialDesc, pattern.Name, roomCtx, materialDesc, pattern.Pattern, ) } // ── Content Check ───────────────────────────────────────── func validateContent(client *openrouter.Client, ctx context.Context, imagePath string) (bool, int) { imgData, err := os.ReadFile(imagePath) if err != nil { return false, 0 } dataURI := "data:image/png;base64," + b64enc(imgData) resp, err := client.Chat(ctx, 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: `Score this image 1-10 for indoor floor replacement suitability. 8-10: Clear indoor room, floor visible. 5-7: Indoor but poor angle. 1-4: Outdoor/pets/objects/no floor. Return ONLY a number (1-10).`}, }}}, }) if err != nil || len(resp.Choices) == 0 { return true, 0 } score := parseScore(resp.Choices[0].Message.Content) return score >= 6, score } 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 b64enc(data []byte) string { const tbl = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/" var b strings.Builder 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.WriteByte(tbl[b0>>2]) b.WriteByte(tbl[((b0&3)<<4)|(b1>>4)]) if i+1 < len(data) { b.WriteByte(tbl[((b1&15)<<2)|(b2>>6)]) } else { b.WriteByte('=') } if i+2 < len(data) { b.WriteByte(tbl[b2&63]) } else { b.WriteByte('=') } } return b.String() }