libnovel/backend/internal/backend/handlers_image.go

package backend

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/libnovel/backend/internal/cfai"
	"github.com/libnovel/backend/internal/domain"
)

// handleAdminImageGenModels handles GET /api/admin/image-gen/models.
// Returns the list of supported Cloudflare AI image generation models.
func (s *Server) handleAdminImageGenModels(w http.ResponseWriter, r *http.Request) {
	if s.deps.ImageGen == nil {
		jsonError(w, http.StatusServiceUnavailable, "image generation not configured (CFAI_ACCOUNT_ID/CFAI_API_TOKEN missing)")
		return
	}
	models := s.deps.ImageGen.Models()
	writeJSON(w, 0, map[string]any{"models": models})
}

// imageGenRequest is the JSON body for POST /api/admin/image-gen.
type imageGenRequest struct {
	// Prompt is the text description of the desired image.
	Prompt string `json:"prompt"`

	// Model is the CF Workers AI model ID (e.g. "@cf/black-forest-labs/flux-2-dev").
	// Defaults to the recommended model for the given type.
	Model string `json:"model"`

	// Type is either "cover" or "chapter".
	Type string `json:"type"`

	// Slug is the book slug. Required for cover; required for chapter.
	Slug string `json:"slug"`

	// Chapter number (1-based). Required when type == "chapter".
	Chapter int `json:"chapter"`

	// ReferenceImageB64 is an optional base64-encoded PNG/JPEG reference image.
	// When present the img2img path is used.
	ReferenceImageB64 string `json:"reference_image_b64"`

	// NumSteps overrides inference steps (default 20).
	NumSteps int `json:"num_steps"`

	// Width / Height override output dimensions (0 = model default).
	Width  int `json:"width"`
	Height int `json:"height"`

	// Guidance overrides prompt guidance scale (0 = model default).
	Guidance float64 `json:"guidance"`

	// Strength for img2img: 0.0–1.0, default 0.75.
	Strength float64 `json:"strength"`

	// SaveToCover when true stores the result as the book cover in MinIO
	// (overwriting any existing cover) and sets the book's cover URL.
	// Only valid when type == "cover".
	SaveToCover bool `json:"save_to_cover"`
}

// imageGenResponse is the JSON body returned by POST /api/admin/image-gen.
type imageGenResponse struct {
	// ImageB64 is the generated image as a base64-encoded PNG string.
	ImageB64 string `json:"image_b64"`
	// ContentType is "image/png" or "image/jpeg".
	ContentType string `json:"content_type"`
	// Saved indicates whether the image was persisted to MinIO.
	Saved bool `json:"saved"`
	// CoverURL is the URL the cover is now served from (only set when Saved==true).
	CoverURL string `json:"cover_url,omitempty"`
	// Model is the model that was used.
	Model string `json:"model"`
	// Bytes is the raw image size in bytes.
	Bytes int `json:"bytes"`
}

// handleAdminImageGen handles POST /api/admin/image-gen.
//
// Generates an image using Cloudflare Workers AI and optionally stores it.
// Multipart/form-data is also accepted so the reference image can be uploaded
// directly; otherwise the reference is expected as base64 JSON.
func (s *Server) handleAdminImageGen(w http.ResponseWriter, r *http.Request) {
	if s.deps.ImageGen == nil {
		jsonError(w, http.StatusServiceUnavailable, "image generation not configured (CFAI_ACCOUNT_ID/CFAI_API_TOKEN missing)")
		return
	}

	var req imageGenRequest
	var refImageData []byte

	ct := r.Header.Get("Content-Type")
	if strings.HasPrefix(ct, "multipart/form-data") {
		// Multipart: parse JSON fields from a "json" part + optional "reference" file part.
		if err := r.ParseMultipartForm(32 << 20); err != nil {
			jsonError(w, http.StatusBadRequest, "parse multipart: "+err.Error())
			return
		}
		if jsonPart := r.FormValue("json"); jsonPart != "" {
			if err := json.Unmarshal([]byte(jsonPart), &req); err != nil {
				jsonError(w, http.StatusBadRequest, "parse json field: "+err.Error())
				return
			}
		}
		if f, _, err := r.FormFile("reference"); err == nil {
			defer f.Close()
			refImageData, _ = io.ReadAll(f)
		}
	} else {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			jsonError(w, http.StatusBadRequest, "parse body: "+err.Error())
			return
		}
		if req.ReferenceImageB64 != "" {
			var decErr error
			refImageData, decErr = base64.StdEncoding.DecodeString(req.ReferenceImageB64)
			if decErr != nil {
				// Try std without padding
				refImageData, decErr = base64.RawStdEncoding.DecodeString(req.ReferenceImageB64)
				if decErr != nil {
					jsonError(w, http.StatusBadRequest, "decode reference_image_b64: "+decErr.Error())
					return
				}
			}
		}
	}

	if strings.TrimSpace(req.Prompt) == "" {
		jsonError(w, http.StatusBadRequest, "prompt is required")
		return
	}
	if req.Type != "cover" && req.Type != "chapter" {
		jsonError(w, http.StatusBadRequest, `type must be "cover" or "chapter"`)
		return
	}
	if req.Slug == "" {
		jsonError(w, http.StatusBadRequest, "slug is required")
		return
	}
	if req.Type == "chapter" && req.Chapter <= 0 {
		jsonError(w, http.StatusBadRequest, "chapter must be > 0 when type is chapter")
		return
	}

	// Resolve model
	model := cfai.ImageModel(req.Model)
	if model == "" {
		if req.Type == "cover" {
			model = cfai.DefaultImageModel
		} else {
			model = cfai.ImageModelFlux2Klein4B
		}
	}

	imgReq := cfai.ImageRequest{
		Prompt:   req.Prompt,
		Model:    model,
		NumSteps: req.NumSteps,
		Width:    req.Width,
		Height:   req.Height,
		Guidance: req.Guidance,
		Strength: req.Strength,
	}

	s.deps.Log.Info("admin: image gen requested",
		"type", req.Type, "slug", req.Slug, "chapter", req.Chapter,
		"model", model, "has_reference", len(refImageData) > 0)

	var imgData []byte
	var genErr error
	if len(refImageData) > 0 {
		imgData, genErr = s.deps.ImageGen.GenerateImageFromReference(r.Context(), imgReq, refImageData)
	} else {
		imgData, genErr = s.deps.ImageGen.GenerateImage(r.Context(), imgReq)
	}
	if genErr != nil {
		s.deps.Log.Error("admin: image gen failed", "err", genErr)
		jsonError(w, http.StatusBadGateway, "image generation failed: "+genErr.Error())
		return
	}

	contentType := sniffImageContentType(imgData)

	// ── Optional persistence ──────────────────────────────────────────────────
	var saved bool
	var coverURL string

	if req.SaveToCover && req.Type == "cover" && s.deps.CoverStore != nil {
		if err := s.deps.CoverStore.PutCover(r.Context(), req.Slug, imgData, contentType); err != nil {
			s.deps.Log.Error("admin: save generated cover failed", "slug", req.Slug, "err", err)
			// Non-fatal: still return the image
		} else {
			saved = true
			coverURL = fmt.Sprintf("/api/cover/novelfire.net/%s", req.Slug)
			s.deps.Log.Info("admin: generated cover saved", "slug", req.Slug, "bytes", len(imgData))
		}
	}

	// Encode result as base64
	b64 := base64.StdEncoding.EncodeToString(imgData)

	writeJSON(w, 0, imageGenResponse{
		ImageB64:    b64,
		ContentType: contentType,
		Saved:       saved,
		CoverURL:    coverURL,
		Model:       string(model),
		Bytes:       len(imgData),
	})
}

// saveCoverRequest is the JSON body for POST /api/admin/image-gen/save-cover.
type saveCoverRequest struct {
	// Slug is the book slug whose cover should be overwritten.
	Slug string `json:"slug"`
	// ImageB64 is the base64-encoded image bytes (PNG or JPEG).
	ImageB64 string `json:"image_b64"`
}

// handleAdminImageGenSaveCover handles POST /api/admin/image-gen/save-cover.
//
// Accepts a pre-generated image as base64 and stores it as the book cover in
// MinIO, replacing the existing one. Does not call Cloudflare AI at all.
func (s *Server) handleAdminImageGenSaveCover(w http.ResponseWriter, r *http.Request) {
	if s.deps.CoverStore == nil {
		jsonError(w, http.StatusServiceUnavailable, "cover store not configured")
		return
	}

	var req saveCoverRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		jsonError(w, http.StatusBadRequest, "parse body: "+err.Error())
		return
	}
	if req.Slug == "" {
		jsonError(w, http.StatusBadRequest, "slug is required")
		return
	}
	if req.ImageB64 == "" {
		jsonError(w, http.StatusBadRequest, "image_b64 is required")
		return
	}

	imgData, err := base64.StdEncoding.DecodeString(req.ImageB64)
	if err != nil {
		imgData, err = base64.RawStdEncoding.DecodeString(req.ImageB64)
		if err != nil {
			jsonError(w, http.StatusBadRequest, "decode image_b64: "+err.Error())
			return
		}
	}

	contentType := sniffImageContentType(imgData)
	if err := s.deps.CoverStore.PutCover(r.Context(), req.Slug, imgData, contentType); err != nil {
		s.deps.Log.Error("admin: save-cover failed", "slug", req.Slug, "err", err)
		jsonError(w, http.StatusInternalServerError, "save cover: "+err.Error())
		return
	}

	s.deps.Log.Info("admin: cover saved via image-gen", "slug", req.Slug, "bytes", len(imgData))
	writeJSON(w, 0, map[string]any{
		"saved":     true,
		"cover_url": fmt.Sprintf("/api/cover/novelfire.net/%s", req.Slug),
		"bytes":     len(imgData),
	})
}

// sniffImageContentType returns the MIME type of the image bytes.
func sniffImageContentType(data []byte) string {
	if len(data) >= 4 {
		// PNG: 0x89 P N G
		if data[0] == 0x89 && data[1] == 0x50 && data[2] == 0x4e && data[3] == 0x47 {
			return "image/png"
		}
		// JPEG: FF D8 FF
		if data[0] == 0xFF && data[1] == 0xD8 && data[2] == 0xFF {
			return "image/jpeg"
		}
		// WebP: RIFF....WEBP
		if len(data) >= 12 && data[0] == 'R' && data[1] == 'I' && data[2] == 'F' && data[3] == 'F' &&
			data[8] == 'W' && data[9] == 'E' && data[10] == 'B' && data[11] == 'P' {
			return "image/webp"
		}
	}
	return "image/png"
}

// saveChapterImageRequest is the JSON body for POST /api/admin/image-gen/save-chapter-image.
type saveChapterImageRequest struct {
	// Slug is the book slug.
	Slug string `json:"slug"`
	// Chapter is the 1-based chapter number.
	Chapter int `json:"chapter"`
	// ImageB64 is the base64-encoded image bytes (PNG or JPEG).
	ImageB64 string `json:"image_b64"`
}

// handleAdminImageGenSaveChapterImage handles POST /api/admin/image-gen/save-chapter-image.
//
// Accepts a pre-generated image as base64 and stores it as the chapter illustration
// in MinIO, replacing the existing one if present. Does not call Cloudflare AI.
func (s *Server) handleAdminImageGenSaveChapterImage(w http.ResponseWriter, r *http.Request) {
	if s.deps.ChapterImageStore == nil {
		jsonError(w, http.StatusServiceUnavailable, "chapter image store not configured")
		return
	}

	var req saveChapterImageRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		jsonError(w, http.StatusBadRequest, "parse body: "+err.Error())
		return
	}
	if req.Slug == "" {
		jsonError(w, http.StatusBadRequest, "slug is required")
		return
	}
	if req.Chapter <= 0 {
		jsonError(w, http.StatusBadRequest, "chapter must be > 0")
		return
	}
	if req.ImageB64 == "" {
		jsonError(w, http.StatusBadRequest, "image_b64 is required")
		return
	}

	imgData, err := base64.StdEncoding.DecodeString(req.ImageB64)
	if err != nil {
		imgData, err = base64.RawStdEncoding.DecodeString(req.ImageB64)
		if err != nil {
			jsonError(w, http.StatusBadRequest, "decode image_b64: "+err.Error())
			return
		}
	}

	contentType := sniffImageContentType(imgData)
	if err := s.deps.ChapterImageStore.PutChapterImage(r.Context(), req.Slug, req.Chapter, imgData, contentType); err != nil {
		s.deps.Log.Error("admin: save-chapter-image failed", "slug", req.Slug, "chapter", req.Chapter, "err", err)
		jsonError(w, http.StatusInternalServerError, "save chapter image: "+err.Error())
		return
	}

	s.deps.Log.Info("admin: chapter image saved", "slug", req.Slug, "chapter", req.Chapter, "bytes", len(imgData))
	writeJSON(w, 0, map[string]any{
		"saved":     true,
		"image_url": fmt.Sprintf("/api/chapter-image/novelfire.net/%s/%d", req.Slug, req.Chapter),
		"bytes":     len(imgData),
	})
}

// handleHeadChapterImage handles HEAD /api/chapter-image/{domain}/{slug}/{n}.
//
// Returns 200 when an image exists for this chapter, 404 otherwise.
// Used by the SSR loader to check existence without downloading the full image.
func (s *Server) handleHeadChapterImage(w http.ResponseWriter, r *http.Request) {
	if s.deps.ChapterImageStore == nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	slug := r.PathValue("slug")
	nStr := r.PathValue("n")
	n, err := strconv.Atoi(nStr)
	if err != nil || n <= 0 {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if s.deps.ChapterImageStore.ChapterImageExists(r.Context(), slug, n) {
		w.WriteHeader(http.StatusOK)
	} else {
		w.WriteHeader(http.StatusNotFound)
	}
}

// handleGetChapterImage handles GET /api/chapter-image/{domain}/{slug}/{n}.
//
// Serves the stored chapter illustration directly from MinIO.
// Returns 404 when no image has been saved for this chapter.
func (s *Server) handleGetChapterImage(w http.ResponseWriter, r *http.Request) {
	if s.deps.ChapterImageStore == nil {
		http.NotFound(w, r)
		return
	}

	slug := r.PathValue("slug")
	nStr := r.PathValue("n")
	n, err := strconv.Atoi(nStr)
	if err != nil || n <= 0 {
		jsonError(w, http.StatusBadRequest, "invalid chapter number")
		return
	}

	data, contentType, ok, err := s.deps.ChapterImageStore.GetChapterImage(r.Context(), slug, n)
	if err != nil {
		s.deps.Log.Error("chapter-image: get failed", "slug", slug, "n", n, "err", err)
		jsonError(w, http.StatusInternalServerError, "could not retrieve chapter image")
		return
	}
	if !ok {
		http.NotFound(w, r)
		return
	}

	w.Header().Set("Content-Type", contentType)
	w.Header().Set("Cache-Control", "public, max-age=31536000, immutable")
	w.Header().Set("Content-Length", fmt.Sprintf("%d", len(data)))
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(data)
}

// handleAdminImageGenAsync handles POST /api/admin/image-gen/async.
//
// Fire-and-forget variant: validates the request, creates an ai_job record of
// kind "image-gen", spawns a background goroutine, and returns HTTP 202 with
// {job_id} immediately. The goroutine calls Cloudflare AI, stores the result
// as base64 in the job payload, and marks the job done/failed when finished.
//
// The admin can then review the result via the ai-jobs page and approve
// (save as cover) or reject (discard) the image.
func (s *Server) handleAdminImageGenAsync(w http.ResponseWriter, r *http.Request) {
	if s.deps.ImageGen == nil {
		jsonError(w, http.StatusServiceUnavailable, "image generation not configured (CFAI_ACCOUNT_ID/CFAI_API_TOKEN missing)")
		return
	}
	if s.deps.AIJobStore == nil {
		jsonError(w, http.StatusServiceUnavailable, "ai job store not configured")
		return
	}

	var req imageGenRequest
	var refImageData []byte

	ct := r.Header.Get("Content-Type")
	if strings.HasPrefix(ct, "multipart/form-data") {
		if err := r.ParseMultipartForm(32 << 20); err != nil {
			jsonError(w, http.StatusBadRequest, "parse multipart: "+err.Error())
			return
		}
		if jsonPart := r.FormValue("json"); jsonPart != "" {
			if err := json.Unmarshal([]byte(jsonPart), &req); err != nil {
				jsonError(w, http.StatusBadRequest, "parse json field: "+err.Error())
				return
			}
		}
		if f, _, err := r.FormFile("reference"); err == nil {
			defer f.Close()
			refImageData, _ = io.ReadAll(f)
		}
	} else {
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			jsonError(w, http.StatusBadRequest, "parse body: "+err.Error())
			return
		}
		if req.ReferenceImageB64 != "" {
			var decErr error
			refImageData, decErr = base64.StdEncoding.DecodeString(req.ReferenceImageB64)
			if decErr != nil {
				refImageData, decErr = base64.RawStdEncoding.DecodeString(req.ReferenceImageB64)
				if decErr != nil {
					jsonError(w, http.StatusBadRequest, "decode reference_image_b64: "+decErr.Error())
					return
				}
			}
		}
	}

	if strings.TrimSpace(req.Prompt) == "" {
		jsonError(w, http.StatusBadRequest, "prompt is required")
		return
	}
	if req.Type != "cover" && req.Type != "chapter" {
		jsonError(w, http.StatusBadRequest, `type must be "cover" or "chapter"`)
		return
	}
	if req.Slug == "" {
		jsonError(w, http.StatusBadRequest, "slug is required")
		return
	}
	if req.Type == "chapter" && req.Chapter <= 0 {
		jsonError(w, http.StatusBadRequest, "chapter must be > 0 when type is chapter")
		return
	}

	// Resolve model.
	model := cfai.ImageModel(req.Model)
	if model == "" {
		if req.Type == "cover" {
			model = cfai.DefaultImageModel
		} else {
			model = cfai.ImageModelFlux2Klein4B
		}
	}

	// Encode request params as job payload so the UI can reconstruct context.
	type jobParams struct {
		Prompt    string `json:"prompt"`
		Type      string `json:"type"`
		Chapter   int    `json:"chapter,omitempty"`
		NumSteps  int    `json:"num_steps,omitempty"`
		Width     int    `json:"width,omitempty"`
		Height    int    `json:"height,omitempty"`
		Guidance  float64 `json:"guidance,omitempty"`
		Strength  float64 `json:"strength,omitempty"`
		HasRef    bool   `json:"has_ref,omitempty"`
	}
	paramsJSON, _ := json.Marshal(jobParams{
		Prompt:   req.Prompt,
		Type:     req.Type,
		Chapter:  req.Chapter,
		NumSteps: req.NumSteps,
		Width:    req.Width,
		Height:   req.Height,
		Guidance: req.Guidance,
		Strength: req.Strength,
		HasRef:   len(refImageData) > 0,
	})

	jobID, createErr := s.deps.AIJobStore.CreateAIJob(r.Context(), domain.AIJob{
		Kind:    "image-gen",
		Slug:    req.Slug,
		Status:  domain.TaskStatusPending,
		Model:   string(model),
		Payload: string(paramsJSON),
		Started: time.Now(),
	})
	if createErr != nil {
		jsonError(w, http.StatusInternalServerError, "create ai job: "+createErr.Error())
		return
	}

	jobCtx, jobCancel := context.WithCancel(context.Background())
	registerCancelJob(jobID, jobCancel)

	// Mark running before returning.
	_ = s.deps.AIJobStore.UpdateAIJob(r.Context(), jobID, map[string]any{
		"status": string(domain.TaskStatusRunning),
	})

	s.deps.Log.Info("admin: image-gen async started",
		"job_id", jobID, "slug", req.Slug, "type", req.Type, "model", model)

	// Capture locals for the goroutine.
	store := s.deps.AIJobStore
	imageGen := s.deps.ImageGen
	coverStore := s.deps.CoverStore
	logger := s.deps.Log
	capturedReq := req
	capturedModel := model
	capturedRefImage := refImageData

	go func() {
		defer deregisterCancelJob(jobID)
		defer jobCancel()

		if jobCtx.Err() != nil {
			_ = store.UpdateAIJob(context.Background(), jobID, map[string]any{
				"status":   string(domain.TaskStatusCancelled),
				"finished": time.Now().Format(time.RFC3339),
			})
			return
		}

		imgReq := cfai.ImageRequest{
			Prompt:   capturedReq.Prompt,
			Model:    capturedModel,
			NumSteps: capturedReq.NumSteps,
			Width:    capturedReq.Width,
			Height:   capturedReq.Height,
			Guidance: capturedReq.Guidance,
			Strength: capturedReq.Strength,
		}

		var imgData []byte
		var genErr error
		if len(capturedRefImage) > 0 {
			imgData, genErr = imageGen.GenerateImageFromReference(jobCtx, imgReq, capturedRefImage)
		} else {
			imgData, genErr = imageGen.GenerateImage(jobCtx, imgReq)
		}

		if genErr != nil {
			logger.Error("admin: image-gen async failed", "job_id", jobID, "err", genErr)
			_ = store.UpdateAIJob(context.Background(), jobID, map[string]any{
				"status":        string(domain.TaskStatusFailed),
				"error_message": genErr.Error(),
				"finished":      time.Now().Format(time.RFC3339),
			})
			return
		}

		contentType := sniffImageContentType(imgData)
		b64 := base64.StdEncoding.EncodeToString(imgData)

		// Build result payload: include the original params + the generated image.
		type resultPayload struct {
			Prompt      string  `json:"prompt"`
			Type        string  `json:"type"`
			Chapter     int     `json:"chapter,omitempty"`
			ContentType string  `json:"content_type"`
			ImageB64    string  `json:"image_b64"`
			Bytes       int     `json:"bytes"`
			NumSteps    int     `json:"num_steps,omitempty"`
			Width       int     `json:"width,omitempty"`
			Height      int     `json:"height,omitempty"`
			Guidance    float64 `json:"guidance,omitempty"`
		}
		resultJSON, _ := json.Marshal(resultPayload{
			Prompt:      capturedReq.Prompt,
			Type:        capturedReq.Type,
			Chapter:     capturedReq.Chapter,
			ContentType: contentType,
			ImageB64:    b64,
			Bytes:       len(imgData),
			NumSteps:    capturedReq.NumSteps,
			Width:       capturedReq.Width,
			Height:      capturedReq.Height,
			Guidance:    capturedReq.Guidance,
		})

		_ = store.UpdateAIJob(context.Background(), jobID, map[string]any{
			"status":     string(domain.TaskStatusDone),
			"items_done": 1,
			"items_total": 1,
			"payload":    string(resultJSON),
			"finished":   time.Now().Format(time.RFC3339),
		})

		logger.Info("admin: image-gen async done",
			"job_id", jobID, "slug", capturedReq.Slug,
			"bytes", len(imgData), "content_type", contentType)

		// Suppress unused variable warning for coverStore when SaveToCover is false.
		_ = coverStore
	}()

	writeJSON(w, http.StatusAccepted, map[string]any{"job_id": jobID})
}