FloorVisualizer/internal/handler/floor_handler.go

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")
		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 CheckContent(client *openrouter.Client) 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 {
			writeJSON(w, 400, map[string]string{"error": err.Error()})
			return
		}
		file, _, _ := r.FormFile("image")
		defer file.Close()
		os.MkdirAll("uploads", 0755)
		tmpPath := filepath.Join("uploads", fmt.Sprintf("check_%d.png", time.Now().UnixNano()))
		dst, _ := os.Create(tmpPath)
		io.Copy(dst, file)
		dst.Close()
		ctx, cancel := context.WithTimeout(r.Context(), 15*time.Second)
		defer cancel()
		passed, score := validateContent(client, ctx, tmpPath)
		msg := fmt.Sprintf("✅ Indoor photo score: %d/10", score)
		if !passed {
			msg = fmt.Sprintf("Non-indoor photo  score: %d/10", score)
		}
		writeJSON(w, 200, map[string]any{"passed": passed, "score": score, "message": msg})
	}
}

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()
}