iamge-labeling-api/main.go

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"github.com/joho/godotenv"
	_ "github.com/mattn/go-sqlite3"
	"github.com/google/uuid"
)

type Config struct {
	Port      string
	URLPrefix string
}

type Image struct {
	ID        int
	Name      string
	Delivered bool
	Score     int
	SecretKey string
	User      string
}

var db *sql.DB
var config Config

func main() {
	// Load environment variables
	err := godotenv.Load()
	if err != nil {
		log.Fatal("Error loading .env file")
	}

	config = Config{
		Port:      os.Getenv("PORT"),
		URLPrefix: os.Getenv("URL_PREFIX"),
	}

	if len(os.Args) < 2 {
		log.Fatal("Usage: ./image_manager [save_name|serve]")
	}

	command := os.Args[1]

	// Initialize database
	initDB()

	switch command {
	case "save_name":
		saveImageNames()
	case "serve":
		serveAPI()
	default:
		log.Fatal("Unknown command. Use 'save_name' or 'serve'")
	}
}

func initDB() {
	var err error
	db, err = sql.Open("sqlite3", "./image_asset.sqlite")
	if err != nil {
		log.Fatal(err)
	}

	// Create table if not exists
	_, err = db.Exec(`
		CREATE TABLE IF NOT EXISTS images (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			name TEXT NOT NULL UNIQUE,
			delivered INTEGER DEFAULT 0,
			score INTEGER DEFAULT 0,
			secret_key TEXT,
			user TEXT
		)
	`)
	if err != nil {
		log.Fatal(err)
	}
}

func saveImageNames() {
	// Get all files in image_asset directory
	files, err := os.ReadDir("./image_asset")
	if err != nil {
		log.Fatal(err)
	}

	// Insert image names into database
	for _, file := range files {
		if !file.IsDir() {
			// Check if the file is an image (simple extension check)
			ext := strings.ToLower(filepath.Ext(file.Name()))
			if ext == ".jpg" || ext == ".jpeg" || ext == ".png" || ext == ".gif" {
				_, err := db.Exec("INSERT OR IGNORE INTO images (name) VALUES (?)", file.Name())
				if err != nil {
					log.Printf("Error inserting %s: %v", file.Name(), err)
				} else {
					log.Printf("Saved: %s", file.Name())
				}
			}
		}
	}
}

func serveAPI() {
	// Create a new mux router
	router := http.NewServeMux()
	router.HandleFunc("/get-image", getImageHandler)
	router.HandleFunc("/save-score", saveScoreHandler)

	log.Printf("Server starting on port %s", config.Port)
	log.Fatal(http.ListenAndServe(":"+config.Port, router))
}

func getImageHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// Get a random undelivered image
	var img Image
	err := db.QueryRow(`
		SELECT id, name FROM images 
		WHERE delivered = 0 
		ORDER BY RANDOM() 
		LIMIT 1
	`).Scan(&img.ID, &img.Name)
	if err != nil {
		if err == sql.ErrNoRows {
			http.Error(w, "No images available", http.StatusNotFound)
			return
		}
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Generate secret key
	secretKey := uuid.New().String()

	// Update the image with delivered status and secret key
	_, err = db.Exec(`
		UPDATE images 
		SET delivered = 1, secret_key = ?
		WHERE id = ?
	`, secretKey, img.ID)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Prepare response
	response := map[string]interface{}{
		"url":       fmt.Sprintf("%s/%s", config.URLPrefix, img.Name),
		"secretKey": secretKey,
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(response)
}

func saveScoreHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	type Request struct {
		FileName  string `json:"fileName"`
		Score     int    `json:"score"`
		User      string `json:"user,omitempty"`
		SecretKey string `json:"secretKey"`
	}

	var req Request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	// Validate request
	if req.FileName == "" || req.SecretKey == "" {
		http.Error(w, "fileName and secretKey are required", http.StatusBadRequest)
		return
	}

	// Verify secret key and update score
	var img Image
	err = db.QueryRow(`
		SELECT id FROM images 
		WHERE name = ? AND secret_key = ?
	`, req.FileName, req.SecretKey).Scan(&img.ID)
	if err != nil {
		if err == sql.ErrNoRows {
			http.Error(w, "Invalid fileName or secretKey", http.StatusUnauthorized)
			return
		}
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Update score and user
	_, err = db.Exec(`
		UPDATE images 
		SET score = ?, user = ?
		WHERE id = ?
	`, req.Score, req.User, img.ID)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Score saved successfully"))
}