RawMark/internal/api/storage.go

package api

import (
	"context"
	"errors"
	"io/fs"
	"net/http"
	"path/filepath"
	"rawmark/internal/i18n"
	"strings"
	"time"
)

type storageCategory struct {
	ID   string `json:"id"`
	Size int64  `json:"size"`
}

type storageUsage struct {
	TotalSize        int64             `json:"totalSize"`
	Categories       []storageCategory `json:"categories"`
	AssetsSize       int64             `json:"assetsSize"`
	TrashSize        int64             `json:"trashSize"`
	BuiltinFontsSize int64             `json:"builtinFontsSize"`
	CustomFontsSize  int64             `json:"customFontsSize"`
	CalculatedAt     int64             `json:"calculatedAt"`
}

func (s *Server) getStorage(w http.ResponseWriter, r *http.Request) {
	usage, err := scanStorage(r.Context(), s.DataDir, s.vault.Root())
	if err != nil {
		internalError(w, r, err)
		return
	}
	usage.BuiltinFontsSize = s.BuiltinFontsSize
	usage.CustomFontsSize = usage.Categories[2].Size
	usage.Categories[2].Size += s.BuiltinFontsSize
	usage.TotalSize += s.BuiltinFontsSize
	writeJSON(w, http.StatusOK, usage)
}

// Count logical bytes of regular files only, without following internal symlinks; scan external vaults separately.
// For nested roots, assign each file to the more specific root to avoid double counting.
func scanStorage(ctx context.Context, dataDir, vaultDir string) (storageUsage, error) {
	usage := storageUsage{Categories: []storageCategory{{ID: "vault"}, {ID: "databases"}, {ID: "fonts"}, {ID: "themes"}, {ID: "plugins"}, {ID: "other"}}}
	if dataDir == "" {
		return usage, i18n.New("application_data_directory_is_not_a1f183")
	}
	dataRoot, err := filepath.EvalSymlinks(dataDir)
	if err != nil {
		return usage, err
	}
	dataRoot, err = filepath.Abs(dataRoot)
	if err != nil {
		return usage, err
	}
	vaultRoot, err := filepath.EvalSymlinks(vaultDir)
	if err != nil {
		return usage, err
	}
	vaultRoot, err = filepath.Abs(vaultRoot)
	if err != nil {
		return usage, err
	}
	appCategory := func(rel string) int {
		top := strings.Split(filepath.ToSlash(rel), "/")[0]
		switch top {
		case "rawmark.db", "rawmark.db-wal", "rawmark.db-shm", "index.db", "index.db-wal", "index.db-shm":
			return 1
		case "fonts":
			return 2
		case "themes":
			return 3
		case "plugins":
			return 4
		default:
			return 5
		}
	}
	walk := func(root, skip string, isVault bool) error {
		return filepath.WalkDir(root, func(path string, entry fs.DirEntry, walkErr error) error {
			if err := ctx.Err(); err != nil {
				return err
			}
			if walkErr != nil {
				if path != root && errors.Is(walkErr, fs.ErrNotExist) {
					return nil // Files may move or disappear during the scan.
				}
				return walkErr
			}
			if path != root && path == skip && entry.IsDir() {
				return filepath.SkipDir
			}
			if entry.IsDir() || entry.Type()&fs.ModeSymlink != 0 {
				return nil
			}
			info, err := entry.Info()
			if errors.Is(err, fs.ErrNotExist) {
				return nil
			}
			if err != nil {
				return err
			}
			if !info.Mode().IsRegular() {
				return nil
			}
			rel, err := filepath.Rel(root, path)
			if err != nil {
				return err
			}
			category := 0
			if !isVault || dataRoot == vaultRoot {
				category = appCategory(rel)
				if isVault && category == 5 {
					category = 0
				}
			}
			size := info.Size()
			usage.TotalSize += size
			usage.Categories[category].Size += size
			if category == 0 {
				parts := strings.Split(filepath.ToSlash(rel), "/")
				if len(parts) >= 3 && parts[0] == ".rawmark" && parts[1] == "trash" {
					usage.TrashSize += size
				} else {
					for _, part := range parts[:len(parts)-1] {
						if part == "assets" || strings.HasSuffix(part, ".assets") {
							usage.AssetsSize += size
							break
						}
					}
				}
			}
			return nil
		})
	}
	if dataRoot != vaultRoot {
		if err := walk(dataRoot, vaultRoot, false); err != nil {
			return storageUsage{}, err
		}
	}
	if err := walk(vaultRoot, dataRoot, true); err != nil {
		return storageUsage{}, err
	}
	usage.CalculatedAt = time.Now().UnixMilli()
	return usage, nil
}