From 7b6f11fa52caffb828a9e90310e615b27e5ba66d Mon Sep 17 00:00:00 2001
From: Noah Hsu <i@nn.ci>
Date: Fri, 10 Jun 2022 16:49:52 +0800
Subject: [PATCH] feat: get account by path

---
 internal/model/account.go      |  8 +++----
 internal/operations/account.go | 40 +++++++++++++++++++++++++++-------
 pkg/utils/balance.go           | 18 +++++++++++++++
 3 files changed, 54 insertions(+), 12 deletions(-)
 create mode 100644 pkg/utils/balance.go

diff --git a/internal/model/account.go b/internal/model/account.go
index ee47b9e5..fa1d4ad4 100644
--- a/internal/model/account.go
+++ b/internal/model/account.go
@@ -3,12 +3,12 @@ package model
 import "time"
 
 type Account struct {
-	ID          uint      `json:"id" gorm:"primaryKey"`
-	VirtualPath string    `json:"virtual_path" gorm:"unique" binding:"required"`
-	Index       int       `json:"index"`
+	ID          uint      `json:"id" gorm:"primaryKey"`                          // unique key
+	VirtualPath string    `json:"virtual_path" gorm:"unique" binding:"required"` // must be standardized
+	Index       int       `json:"index"`                                         // use to sort
 	Driver      string    `json:"driver"`
 	Status      string    `json:"status"`
-	Addition    string    `json:"addition"`
+	Addition    string    `json:"addition"` // Additional information, defined in the corresponding driver
 	Remark      string    `json:"remark"`
 	Modified    time.Time `json:"modified"`
 	Sort
diff --git a/internal/operations/account.go b/internal/operations/account.go
index 91d8dfb3..c64d9ffa 100644
--- a/internal/operations/account.go
+++ b/internal/operations/account.go
@@ -7,8 +7,10 @@ import (
 	"github.com/alist-org/alist/v3/internal/store"
 	"github.com/alist-org/alist/v3/pkg/utils"
 	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
 	"sort"
 	"strings"
+	"sync"
 )
 
 // Although the driver type is stored,
@@ -94,8 +96,6 @@ func SaveDriverAccount(driver driver.Driver) error {
 	return nil
 }
 
-var balance = ".balance"
-
 // GetAccountsByPath get account by longest match path, contains balance account.
 // for example, there is /a/b,/a/c,/a/d/e,/a/d/e.balance
 // GetAccountsByPath(/a/d/e/f) => /a/d/e,/a/d/e.balance
@@ -103,11 +103,7 @@ func GetAccountsByPath(path string) []driver.Driver {
 	accounts := make([]driver.Driver, 0)
 	curSlashCount := 0
 	for _, v := range accountsMap {
-		virtualPath := utils.StandardizationPath(v.GetAccount().VirtualPath)
-		bIndex := strings.LastIndex(virtualPath, balance)
-		if bIndex != -1 {
-			virtualPath = virtualPath[:bIndex]
-		}
+		virtualPath := utils.GetActualVirtualPath(v.GetAccount().VirtualPath)
 		if virtualPath == "/" {
 			virtualPath = ""
 		}
@@ -154,7 +150,7 @@ func GetAccountFilesByPath(prefix string) []driver.FileInfo {
 	set := make(map[string]interface{})
 	for _, v := range accounts {
 		// balance account
-		if strings.Contains(v.GetAccount().VirtualPath, balance) {
+		if utils.IsBalance(v.GetAccount().VirtualPath) {
 			continue
 		}
 		full := utils.StandardizationPath(v.GetAccount().VirtualPath)
@@ -178,3 +174,31 @@ func GetAccountFilesByPath(prefix string) []driver.FileInfo {
 	}
 	return files
 }
+
+var balanceMap sync.Map
+
+// GetBalancedAccount get account by path
+func GetBalancedAccount(path string) driver.Driver {
+	path = utils.StandardizationPath(path)
+	accounts := GetAccountsByPath(path)
+	accountNum := len(accounts)
+	switch accountNum {
+	case 0:
+		return nil
+	case 1:
+		return accounts[0]
+	default:
+		virtualPath := utils.GetActualVirtualPath(accounts[0].GetAccount().VirtualPath)
+		cur, ok := balanceMap.Load(virtualPath)
+		i := 0
+		if ok {
+			i = cur.(int)
+			i = (i + 1) % accountNum
+			balanceMap.Store(virtualPath, i)
+		} else {
+			balanceMap.Store(virtualPath, i)
+		}
+		log.Debugln("use: ", i)
+		return accounts[i]
+	}
+}
diff --git a/pkg/utils/balance.go b/pkg/utils/balance.go
new file mode 100644
index 00000000..265c45f0
--- /dev/null
+++ b/pkg/utils/balance.go
@@ -0,0 +1,18 @@
+package utils
+
+import "strings"
+
+var balance = ".balance"
+
+func IsBalance(str string) bool {
+	return strings.Contains(str, balance)
+}
+
+// GetActualVirtualPath remove balance suffix
+func GetActualVirtualPath(virtualPath string) string {
+	bIndex := strings.LastIndex(virtualPath, ".balance")
+	if bIndex != -1 {
+		virtualPath = virtualPath[:bIndex]
+	}
+	return virtualPath
+}