Haystack/backend/images/handler.go

package images

import (
	"bytes"
	"database/sql"
	"encoding/base64"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"screenmark/screenmark/.gen/haystack/haystack/model"
	"screenmark/screenmark/middleware"
	"screenmark/screenmark/models"

	"github.com/charmbracelet/log"
	"github.com/go-chi/chi/v5"
)

type ImageHandler struct {
	logger     *log.Logger
	imageModel models.ImageModel
	userModel  models.UserModel
}

type ImagesReturn struct {
	UserImages       []models.UserImageWithImage  `json:"userImages"`
	ProcessingImages []models.UserProcessingImage `json:"processingImages"`
	Lists            []models.ListsWithImages     `json:"lists"`
}

func (h *ImageHandler) serveImage(w http.ResponseWriter, r *http.Request) {
	imageId, err := middleware.GetPathParamID(h.logger, "id", w, r)
	if err != nil {
		return
	}

	image, err := h.imageModel.Get(r.Context(), imageId)
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "Could not get image")
		return
	}

	// TODO: this could be part of the db table
	extension := filepath.Ext(image.ImageName)
	if len(extension) == 0 {
		// Same hack
		extension = "png"
	}
	extension = extension[1:]

	w.Header().Add("Content-Type", "image/"+extension)
	w.Write(image.Image)
}

func (h *ImageHandler) listImages(w http.ResponseWriter, r *http.Request) {
	userId, err := middleware.GetUserID(r.Context(), h.logger, w)
	if err != nil {
		return
	}

	images, err := h.userModel.GetUserImages(r.Context(), userId)
	if err != nil {
		middleware.WriteErrorInternal(h.logger, "could not get user images", w)
		return
	}

	processingImages, err := h.imageModel.GetProcessing(r.Context(), userId)
	if err != nil {
		middleware.WriteErrorInternal(h.logger, "could not get processing images", w)
		return
	}

	listsWithImages, err := h.userModel.ListWithImages(r.Context(), userId)
	if err != nil {
		middleware.WriteErrorInternal(h.logger, "could not get lists with images", w)
		return
	}

	imagesReturn := ImagesReturn{
		UserImages:       images,
		ProcessingImages: processingImages,
		Lists:            listsWithImages,
	}

	middleware.WriteJsonOrError(h.logger, imagesReturn, w)
}

func (h *ImageHandler) uploadImage(w http.ResponseWriter, r *http.Request) {
	imageName := chi.URLParam(r, "name")
	if len(imageName) == 0 {
		middleware.WriteErrorBadRequest(h.logger, "you need to provide a name in the path", w)
		return
	}

	userId, err := middleware.GetUserID(r.Context(), h.logger, w)
	if err != nil {
		return
	}

	contentType := r.Header.Get("Content-Type")

	image := make([]byte, 0)
	switch contentType {
	case "application/base64":
		decoder := base64.NewDecoder(base64.StdEncoding, r.Body)
		buf := &bytes.Buffer{}

		_, err := io.Copy(buf, decoder)
		if err != nil {
			middleware.WriteErrorBadRequest(h.logger, "base64 decoding failed", w)
			return
		}

		image = buf.Bytes()
	case "application/oclet-stream", "image/png":
		bodyData, err := io.ReadAll(r.Body)
		if err != nil {
			middleware.WriteErrorBadRequest(h.logger, "binary data reading failed", w)
			return
		}
		// TODO: check headers

		image = bodyData
	default:
		middleware.WriteErrorBadRequest(h.logger, "unsupported content type, need octet-stream or base64", w)
		return
	}

	userImage, err := h.imageModel.Process(r.Context(), userId, model.Image{
		Image:     image,
		ImageName: imageName,
	})

	if err != nil {
		middleware.WriteErrorInternal(h.logger, "could not save image to DB", w)
		return
	}

	middleware.WriteJsonOrError(h.logger, userImage, w)
}

func (h *ImageHandler) CreateRoutes(r chi.Router) {
	h.logger.Info("Mounting image router")

	// Public route for serving images (not protected)
	r.Get("/{id}", h.serveImage)

	// Protected routes
	r.Group(func(r chi.Router) {
		r.Use(middleware.ProtectedRoute)
		r.Use(middleware.SetJson)

		r.Get("/", h.listImages)
		r.Post("/{name}", h.uploadImage)
	})
}

func CreateImageHandler(db *sql.DB) ImageHandler {
	imageModel := models.NewImageModel(db)
	userModel := models.NewUserModel(db)
	logger := log.New(os.Stdout).WithPrefix("Images")

	return ImageHandler{
		logger:     logger,
		imageModel: imageModel,
		userModel:  userModel,
	}
}