FloorVisualizer/internal/openrouter/image.go

package openrouter

import (
	"bytes"
	"context"
	"encoding/base64"
	stderrors "errors"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/image/draw"
	_ "golang.org/x/image/webp"
)

const Gemini3ProImagePreviewModel = "google/gemini-3-pro-image-preview"

const maskImageInstruction = "The next image is a mask. Non-black pixels mark regions to edit; black pixels mark protected regions that must remain unchanged."

const referenceImageInstruction = "The next image is a reference for the floor material style. Match its color, texture, pattern, and overall appearance exactly when generating the new floor. The generated floor should look like this material."

var supportedAspectRatios = []struct {
	value  string
	width  int
	height int
}{
	{value: "1:1", width: 1, height: 1},
	{value: "2:3", width: 2, height: 3},
	{value: "3:2", width: 3, height: 2},
	{value: "3:4", width: 3, height: 4},
	{value: "4:3", width: 4, height: 3},
	{value: "4:5", width: 4, height: 5},
	{value: "5:4", width: 5, height: 4},
	{value: "9:16", width: 9, height: 16},
	{value: "16:9", width: 16, height: 9},
	{value: "21:9", width: 21, height: 9},
}

type ImageGenerationRequest struct {
	Prompt             string
	OutputPath         string
	Model              string
	Modalities         []string
	Provider           *Provider
	InputImagePath     string
	InputMimeType      string
	MaskImagePath      string
	MaskMimeType       string
	ReferenceImagePath string // optional: reference image for floor material style
	ReferenceMimeType  string
	AspectRatio        string
	ImageSize          string
	TopP               *float64
	TargetWidth        int
	TargetHeight       int
	MatchSizeImagePath string
}

type ImageGenerationResult struct {
	Path        string
	MimeType    string
	Model       string
	TextContent string
	Width       int
	Height      int
}

type ImageRecognitionRequest struct {
	Prompt             string
	OutputPath         string
	Model              string
	Modalities         []string
	Provider           *Provider
	InputImagePath     string
	InputMimeType      string
	AspectRatio        string
	ImageSize          string
	TopP               *float64
	TargetWidth        int
	TargetHeight       int
	MatchSizeImagePath string
	ResizeOutput       bool
}

type ImageRecognitionResult = ImageGenerationResult

// RecognizeImage calls an OpenRouter image-capable chat model and saves the
// first returned image as a PNG file while returning the text content for JSON parsing.
func (c *Client) RecognizeImage(ctx context.Context, req ImageRecognitionRequest) (*ImageRecognitionResult, error) {
	if strings.TrimSpace(req.Prompt) == "" {
		return nil, stderrors.New("openrouter: recognition prompt is required")
	}
	if strings.TrimSpace(req.InputImagePath) == "" {
		return nil, stderrors.New("openrouter: recognition input image path is required")
	}
	if strings.TrimSpace(req.OutputPath) == "" {
		return nil, stderrors.New("openrouter: recognition output path is required")
	}

	model := req.Model
	if model == "" {
		model = Gemini3ProImagePreviewModel
	}
	targetWidth, targetHeight := 0, 0
	if req.ResizeOutput || req.AspectRatio == "" {
		if req.MatchSizeImagePath == "" {
			req.MatchSizeImagePath = req.InputImagePath
		}
		var err error
		targetWidth, targetHeight, err = resolveTargetSize(ImageGenerationRequest{
			TargetWidth:        req.TargetWidth,
			TargetHeight:       req.TargetHeight,
			MatchSizeImagePath: req.MatchSizeImagePath,
		})
		if err != nil {
			return nil, err
		}
	}

	content, err := buildImageRecognitionContent(req)
	if err != nil {
		return nil, err
	}

	chatReq := ChatRequest{
		Model:      model,
		Modalities: imageOutputModalities(model, req.Modalities),
		Provider:   req.Provider,
		TopP:       req.TopP,
		Messages: []Message{
			{
				Role:    "user",
				Content: content,
			},
		},
	}
	if req.AspectRatio == "" && targetWidth > 0 && targetHeight > 0 {
		req.AspectRatio = aspectRatioString(targetWidth, targetHeight)
	}
	if req.AspectRatio != "" || req.ImageSize != "" {
		chatReq.ImageConfig = &ImageConfig{
			AspectRatio: req.AspectRatio,
			ImageSize:   normalizeImageSize(req.ImageSize),
		}
	}

	resp, err := c.chatCompletion(ctx, chatReq)
	if err != nil {
		return nil, err
	}
	outputWidth, outputHeight := 0, 0
	if req.ResizeOutput {
		outputWidth, outputHeight = targetWidth, targetHeight
	}
	return saveFirstResponseImage(resp.Choices, req.OutputPath, model, outputWidth, outputHeight)
}

// GenerateImageToFile calls an OpenRouter image-capable chat model and saves the
// first returned image as a PNG file.
func (c *Client) GenerateImageToFile(ctx context.Context, req ImageGenerationRequest) (*ImageGenerationResult, error) {
	if strings.TrimSpace(req.Prompt) == "" {
		return nil, stderrors.New("openrouter: image prompt is required")
	}
	if strings.TrimSpace(req.OutputPath) == "" {
		return nil, stderrors.New("openrouter: image output path is required")
	}

	model := req.Model
	if model == "" {
		model = Gemini3ProImagePreviewModel
	}
	if req.MatchSizeImagePath == "" {
		req.MatchSizeImagePath = req.InputImagePath
	}

	targetWidth, targetHeight, err := resolveTargetSize(req)
	if err != nil {
		return nil, err
	}

	content, err := buildImageGenerationContent(req)
	if err != nil {
		return nil, err
	}

	chatReq := ChatRequest{
		Model:      model,
		Modalities: imageOutputModalities(model, req.Modalities),
		Provider:   req.Provider,
		TopP:       req.TopP,
		Messages: []Message{
			{
				Role:    "user",
				Content: content,
			},
		},
	}
	if req.AspectRatio == "" && targetWidth > 0 && targetHeight > 0 {
		req.AspectRatio = aspectRatioString(targetWidth, targetHeight)
	}
	if req.AspectRatio != "" || req.ImageSize != "" {
		chatReq.ImageConfig = &ImageConfig{
			AspectRatio: req.AspectRatio,
			ImageSize:   normalizeImageSize(req.ImageSize),
		}
	}

	resp, err := c.chatCompletion(ctx, chatReq)
	if err != nil {
		return nil, err
	}
	return saveFirstResponseImage(resp.Choices, req.OutputPath, model, targetWidth, targetHeight)
}

func saveFirstResponseImage(choices []Choice, outputPath string, model string, targetWidth, targetHeight int) (*ImageGenerationResult, error) {
	outputPath = normalizePNGOutputPath(outputPath)
	if len(choices) == 0 {
		return nil, stderrors.New("openrouter: image response has no choices")
	}

	message := choices[0].Message
	if len(message.Images) == 0 {
		return nil, fmt.Errorf("openrouter: image response has no images: %s", message.Content)
	}

	_, data, err := decodeDataURL(message.Images[0].ImageURL.URL)
	if err != nil {
		return nil, err
	}
	data, width, height, err := encodePNGImageData(data, targetWidth, targetHeight)
	if err != nil {
		return nil, err
	}

	if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
		return nil, fmt.Errorf("openrouter: create output directory: %w", err)
	}
	if err := os.WriteFile(outputPath, data, 0644); err != nil {
		return nil, fmt.Errorf("openrouter: write image file: %w", err)
	}

	return &ImageGenerationResult{
		Path:        outputPath,
		MimeType:    "image/png",
		Model:       model,
		TextContent: message.Content,
		Width:       width,
		Height:      height,
	}, nil
}

func resolveTargetSize(req ImageGenerationRequest) (int, int, error) {
	if req.TargetWidth < 0 || req.TargetHeight < 0 {
		return 0, 0, stderrors.New("openrouter: target image size must be positive")
	}
	if req.TargetWidth > 0 || req.TargetHeight > 0 {
		if req.TargetWidth == 0 || req.TargetHeight == 0 {
			return 0, 0, stderrors.New("openrouter: target width and height must be set together")
		}
		return req.TargetWidth, req.TargetHeight, nil
	}
	if req.MatchSizeImagePath == "" {
		return 0, 0, nil
	}

	file, err := os.Open(req.MatchSizeImagePath)
	if err != nil {
		return 0, 0, fmt.Errorf("openrouter: open match size image: %w", err)
	}
	defer file.Close()

	config, _, err := image.DecodeConfig(file)
	if err != nil {
		return 0, 0, fmt.Errorf("openrouter: decode match size image config: %w", err)
	}
	return config.Width, config.Height, nil
}

func buildImageGenerationContent(req ImageGenerationRequest) ([]ContentPart, error) {
	if req.MaskImagePath != "" && req.InputImagePath == "" {
		return nil, stderrors.New("openrouter: mask image requires input image")
	}

	content := make([]ContentPart, 0, 6)
	if req.InputImagePath != "" {
		dataURI, err := imageFileDataURI(req.InputImagePath, req.InputMimeType)
		if err != nil {
			return nil, err
		}
		content = append(content, ContentPart{
			Type:     "image_url",
			ImageURL: &ImageURL{URL: dataURI},
		})
	}
	if req.MaskImagePath != "" {
		dataURI, err := imageFileDataURI(req.MaskImagePath, req.MaskMimeType)
		if err != nil {
			return nil, err
		}
		content = append(content, ContentPart{Type: "text", Text: maskImageInstruction})
		content = append(content, ContentPart{
			Type:     "image_url",
			ImageURL: &ImageURL{URL: dataURI},
		})
	}
	if req.ReferenceImagePath != "" {
		dataURI, err := imageFileDataURI(req.ReferenceImagePath, req.ReferenceMimeType)
		if err != nil {
			return nil, err
		}
		content = append(content, ContentPart{
			Type:     "image_url",
			ImageURL: &ImageURL{URL: dataURI},
		})
		content = append(content, ContentPart{Type: "text", Text: referenceImageInstruction})
	}
	content = append(content, ContentPart{Type: "text", Text: req.Prompt})
	return content, nil
}

func buildImageRecognitionContent(req ImageRecognitionRequest) ([]ContentPart, error) {
	dataURI, err := imageFileDataURI(req.InputImagePath, req.InputMimeType)
	if err != nil {
		return nil, err
	}
	return []ContentPart{
		{
			Type:     "image_url",
			ImageURL: &ImageURL{URL: dataURI},
		},
		{Type: "text", Text: req.Prompt},
	}, nil
}

func imageFileDataURI(path string, mimeType string) (string, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("openrouter: read input image: %w", err)
	}
	if mimeType == "" {
		mimeType = http.DetectContentType(data)
	}
	return fmt.Sprintf("data:%s;base64,%s", mimeType, base64.StdEncoding.EncodeToString(data)), nil
}

func normalizePNGOutputPath(outputPath string) string {
	ext := filepath.Ext(outputPath)
	if strings.EqualFold(ext, ".png") {
		return strings.TrimSuffix(outputPath, ext) + ".png"
	}
	if ext == "" {
		return outputPath + ".png"
	}
	return strings.TrimSuffix(outputPath, ext) + ".png"
}

func normalizeImageSize(imageSize string) string {
	imageSize = strings.TrimSpace(imageSize)
	if strings.HasSuffix(strings.ToLower(imageSize), "k") {
		return strings.TrimSuffix(imageSize, imageSize[len(imageSize)-1:]) + "K"
	}
	return imageSize
}

func imageOutputModalities(model string, override []string) []string {
	if len(override) > 0 {
		return normalizeModalities(override)
	}
	if isImageOnlyModel(model) {
		return []string{"image"}
	}
	return []string{"image", "text"}
}

func normalizeModalities(modalities []string) []string {
	normalized := make([]string, 0, len(modalities))
	seen := make(map[string]struct{}, len(modalities))
	for _, modality := range modalities {
		modality = strings.ToLower(strings.TrimSpace(modality))
		if modality == "" {
			continue
		}
		if _, ok := seen[modality]; ok {
			continue
		}
		seen[modality] = struct{}{}
		normalized = append(normalized, modality)
	}
	return normalized
}

func isImageOnlyModel(model string) bool {
	model = strings.ToLower(strings.TrimSpace(model))
	imageOnlyPrefixes := []string{
		"black-forest-labs/",
		"sourceful/",
		"bytedance-seed/",
	}
	for _, prefix := range imageOnlyPrefixes {
		if strings.HasPrefix(model, prefix) {
			return true
		}
	}

	imageOnlyModels := []string{
		"x-ai/grok-imagine-image-quality",
	}
	for _, candidate := range imageOnlyModels {
		if model == candidate {
			return true
		}
	}
	return false
}

func encodePNGImageData(data []byte, targetWidth, targetHeight int) ([]byte, int, int, error) {
	src, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, 0, 0, fmt.Errorf("openrouter: decode generated image: %w", err)
	}

	dst := src
	width := src.Bounds().Dx()
	height := src.Bounds().Dy()
	if targetWidth > 0 && targetHeight > 0 {
		resized := image.NewRGBA(image.Rect(0, 0, targetWidth, targetHeight))
		draw.ApproxBiLinear.Scale(resized, resized.Bounds(), src, src.Bounds(), draw.Over, nil)
		dst = resized
		width = targetWidth
		height = targetHeight
	}

	var buf bytes.Buffer
	if err := png.Encode(&buf, dst); err != nil {
		return nil, 0, 0, fmt.Errorf("openrouter: encode png image: %w", err)
	}
	return buf.Bytes(), width, height, nil
}

func aspectRatioString(width, height int) string {
	if width <= 0 || height <= 0 {
		return ""
	}

	sourceRatio := math.Log(float64(width) / float64(height))
	best := supportedAspectRatios[0]
	bestDistance := math.Abs(sourceRatio - math.Log(float64(best.width)/float64(best.height)))
	for _, candidate := range supportedAspectRatios[1:] {
		distance := math.Abs(sourceRatio - math.Log(float64(candidate.width)/float64(candidate.height)))
		if distance < bestDistance {
			best = candidate
			bestDistance = distance
		}
	}
	return best.value
}

func decodeDataURL(dataURL string) (string, []byte, error) {
	header, payload, ok := strings.Cut(dataURL, ",")
	if !ok {
		return "", nil, stderrors.New("openrouter: invalid image data url")
	}
	if !strings.HasPrefix(header, "data:") || !strings.Contains(header, ";base64") {
		return "", nil, stderrors.New("openrouter: unsupported image data url")
	}

	mimeType := strings.TrimPrefix(strings.Split(header, ";")[0], "data:")
	data, err := base64.StdEncoding.DecodeString(payload)
	if err != nil {
		return "", nil, fmt.Errorf("openrouter: decode image data: %w", err)
	}

	return mimeType, data, nil
}

// MaxImageDimension is the maximum width or height for images sent to the AI API.
const MaxImageDimension = 2048

// ResizeImageToFit reads an image, scales it so the longest side ≤ maxDim pixels,
// and writes the result as PNG.
func ResizeImageToFit(inputPath, outputPath string) error {
	data, err := os.ReadFile(inputPath)
	if err != nil {
		return fmt.Errorf("resize: read input: %w", err)
	}
	src, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return fmt.Errorf("resize: decode: %w", err)
	}
	b := src.Bounds()
	w, h := b.Dx(), b.Dy()
	longest := w
	if h > longest {
		longest = h
	}
	dst := src
	if longest > MaxImageDimension {
		scale := float64(MaxImageDimension) / float64(longest)
		newW := int(float64(w)*scale + 0.5)
		newH := int(float64(h)*scale + 0.5)
		resized := image.NewRGBA(image.Rect(0, 0, newW, newH))
		draw.ApproxBiLinear.Scale(resized, resized.Bounds(), src, src.Bounds(), draw.Over, nil)
		dst = resized
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, dst); err != nil {
		return fmt.Errorf("resize: encode png: %w", err)
	}
	if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
		return err
	}
	return os.WriteFile(outputPath, buf.Bytes(), 0644)
}