am-downloader/handlers.go

package main

import (
	"io"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
)

func (a *App) search(c *gin.Context) {
	keyword := strings.TrimSpace(c.Query("keyword"))
	if keyword == "" {
		c.JSON(400, gin.H{"error": "keyword required"})
		return
	}
	source := c.DefaultQuery("source", "kw")
	page, _ := strconv.Atoi(c.DefaultQuery("page", "1"))
	if page < 1 {
		page = 1
	}
	limit, _ := strconv.Atoi(c.DefaultQuery("limit", "20"))
	if limit < 1 {
		limit = 20
	}
	resp, err := a.am.Search(keyword, source, page, limit)
	if err != nil {
		c.JSON(502, gin.H{"error": err.Error()})
		return
	}
	defer resp.Body.Close()
	c.Status(resp.StatusCode)
	_, _ = io.Copy(c.Writer, resp.Body)
}

func (a *App) listTasks(c *gin.Context) {
	var ts []Task
	user := c.MustGet("user").(*User)
	query := a.db.Where("user_id = ?", user.ID)
	if user.Role == RoleAdmin && c.Query("all") == "1" {
		query = a.db
	}
	if err := query.Order("id desc").Limit(100).Find(&ts).Error; err != nil {
		c.JSON(500, gin.H{"error": err.Error()})
		return
	}
	c.JSON(200, gin.H{"data": ts})
}

func (a *App) addTask(c *gin.Context) {
	var in SongRequest
	if c.BindJSON(&in) != nil || in.RID == "" || in.Title == "" {
		c.JSON(400, gin.H{"error": "source/rid/title required"})
		return
	}
	if in.Source == "" {
		in.Source = "kw"
	}
	if in.Quality == "" {
		in.Quality = a.cfg.DefaultQuality
	}
	user := c.MustGet("user").(*User)

	var existing Task
	err := a.db.Where("user_id = ? AND source = ? AND r_id = ? AND quality = ?",
		user.ID, in.Source, in.RID, in.Quality).Order("id desc").First(&existing).Error
	if err == nil {
		existing.Title, existing.Artist, existing.Album, existing.Pic = in.Title, in.Artist, in.Album, in.Pic
		switch existing.Status {
		case "pending", "resolving", "downloading", "uploading", "upload_pending":
			c.JSON(200, existing)
			return
		case "done":
			c.JSON(200, existing)
			return
		case "error":
			if err := a.prepareRetry(&existing); err != nil {
				if err == errInsufficientCredits {
					c.JSON(402, gin.H{"error": err.Error()})
				} else {
					c.JSON(409, gin.H{"error": err.Error()})
				}
				return
			}
			a.enqueueTask(existing.ID)
			c.JSON(202, existing)
			return
		}
	}

	t := Task{
		UserID: user.ID, Source: in.Source, RID: in.RID,
		Title: in.Title, Artist: in.Artist, Album: in.Album,
		Pic: in.Pic, Quality: in.Quality, Status: "pending",
	}
	var reservation *CreditReservation
	err = a.transactionRetry(func(tx *gorm.DB) error {
		if err := tx.Create(&t).Error; err != nil {
			return err
		}
		var reserveErr error
		reservation, reserveErr = reserveCreditTx(tx, user.ID, t.ID, 1)
		if reserveErr != nil {
			return reserveErr
		}
		t.CreditReservationID = reservation.ID
		return tx.Model(&t).Update("credit_reservation_id", reservation.ID).Error
	})
	if err != nil {
		if a.db.Where("user_id = ? AND source = ? AND r_id = ? AND quality = ?",
			user.ID, in.Source, in.RID, in.Quality).First(&existing).Error == nil {
			c.JSON(200, existing)
			return
		}
		if err == errInsufficientCredits {
			c.JSON(402, gin.H{"error": err.Error()})
			return
		}
		c.JSON(500, gin.H{"error": err.Error()})
		return
	}

	a.enqueueTask(t.ID)
	c.JSON(202, t)
}

func (a *App) creditInfo(c *gin.Context) {
	user := c.MustGet("user").(*User)
	account, err := a.creditBalance(user.ID)
	if err != nil {
		c.JSON(500, gin.H{"error": err.Error()})
		return
	}
	c.JSON(200, gin.H{"balance": account.Balance, "reserved": account.Reserved, "available": account.Balance - account.Reserved})
}

func (a *App) retryTask(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.Status(400)
		return
	}
	user := c.MustGet("user").(*User)
	var t Task
	if a.db.Where("id = ? AND user_id = ?", id, user.ID).First(&t).Error != nil {
		c.Status(404)
		return
	}
	switch t.Status {
	case "done":
		c.JSON(409, gin.H{"error": "task is already completed"})
		return
	case "pending", "resolving", "downloading", "uploading":
		c.JSON(409, gin.H{"error": "task is already running"})
		return
	case "upload_pending", "error":
		if t.Status == "error" {
			if err := a.prepareRetry(&t); err != nil {
				if err == errInsufficientCredits {
					c.JSON(402, gin.H{"error": err.Error()})
				} else {
					c.JSON(409, gin.H{"error": err.Error()})
				}
				return
			}
		}
	default:
		c.JSON(409, gin.H{"error": "task cannot be retried from current state"})
		return
	}
	a.enqueueTask(t.ID)
	c.JSON(202, t)
}

func (a *App) deleteTask(c *gin.Context) {
	id, err := strconv.ParseUint(c.Param("id"), 10, 64)
	if err != nil {
		c.Status(400)
		return
	}
	user := c.MustGet("user").(*User)
	var t Task
	if err := a.db.Where("id = ? AND user_id = ?", id, user.ID).First(&t).Error; err != nil {
		c.Status(404)
		return
	}
	if t.CreditReservationID != 0 {
		if err := a.refundCredit(&t); err != nil {
			c.JSON(409, gin.H{"error": "cannot release task credit: " + err.Error()})
			return
		}
	}
	if err := a.db.Delete(&t).Error; err != nil {
		c.JSON(500, gin.H{"error": err.Error()})
		return
	}
	c.Status(204)
}