More cleanup, more refractor, more energy, more passion, more footwork

pkg/qbit/storage.go

@@ -0,0 +1,151 @@
+package qbit
+
+import (
+	"encoding/json"
+	"os"
+	"sync"
+)
+
+type TorrentStorage struct {
+	torrents map[string]*Torrent
+	mu       sync.RWMutex
+	order    []string
+	filename string // Added to store the filename for persistence
+}
+
+func loadTorrentsFromJSON(filename string) (map[string]*Torrent, error) {
+	data, err := os.ReadFile(filename)
+	if err != nil {
+		return nil, err
+	}
+	torrents := make(map[string]*Torrent)
+	if err := json.Unmarshal(data, &torrents); err != nil {
+		return nil, err
+	}
+	return torrents, nil
+}
+
+func NewTorrentStorage(filename string) *TorrentStorage {
+	// Open the JSON file and read the data
+	torrents, err := loadTorrentsFromJSON(filename)
+	if err != nil {
+		torrents = make(map[string]*Torrent)
+	}
+	order := make([]string, 0, len(torrents))
+	for id := range torrents {
+		order = append(order, id)
+	}
+	// Create a new TorrentStorage
+	return &TorrentStorage{
+		torrents: torrents,
+		order:    order,
+		filename: filename,
+	}
+}
+
+func (ts *TorrentStorage) Add(torrent *Torrent) {
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	ts.torrents[torrent.Hash] = torrent
+	ts.order = append(ts.order, torrent.Hash)
+	_ = ts.saveToFile()
+}
+
+func (ts *TorrentStorage) AddOrUpdate(torrent *Torrent) {
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	if _, exists := ts.torrents[torrent.Hash]; !exists {
+		ts.order = append(ts.order, torrent.Hash)
+	}
+	ts.torrents[torrent.Hash] = torrent
+	_ = ts.saveToFile()
+}
+
+func (ts *TorrentStorage) GetByID(id string) *Torrent {
+	ts.mu.RLock()
+	defer ts.mu.RUnlock()
+	for _, torrent := range ts.torrents {
+		if torrent.ID == id {
+			return torrent
+		}
+	}
+	return nil
+}
+
+func (ts *TorrentStorage) Get(hash string) *Torrent {
+	ts.mu.RLock()
+	defer ts.mu.RUnlock()
+	return ts.torrents[hash]
+}
+
+func (ts *TorrentStorage) GetAll(category string, filter string, hashes []string) []*Torrent {
+	ts.mu.RLock()
+	defer ts.mu.RUnlock()
+	torrents := make([]*Torrent, 0)
+	for _, id := range ts.order {
+		torrent := ts.torrents[id]
+		if category != "" && torrent.Category != category {
+			continue
+		}
+		if filter != "" && torrent.State != filter {
+			continue
+		}
+		torrents = append(torrents, torrent)
+	}
+	if len(hashes) > 0 {
+		filtered := make([]*Torrent, 0, len(torrents))
+		for _, hash := range hashes {
+			if torrent := ts.torrents[hash]; torrent != nil {
+				filtered = append(filtered, torrent)
+			}
+		}
+		torrents = filtered
+	}
+	return torrents
+}
+
+func (ts *TorrentStorage) Update(torrent *Torrent) {
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	ts.torrents[torrent.Hash] = torrent
+	_ = ts.saveToFile()
+}
+
+func (ts *TorrentStorage) Delete(hash string) {
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	torrent, exists := ts.torrents[hash]
+	if !exists {
+		return
+	}
+	delete(ts.torrents, hash)
+	for i, id := range ts.order {
+		if id == hash {
+			ts.order = append(ts.order[:i], ts.order[i+1:]...)
+			break
+		}
+	}
+	// Delete the torrent folder
+	if torrent.ContentPath != "" {
+		err := os.RemoveAll(torrent.ContentPath)
+		if err != nil {
+			return
+		}
+	}
+	_ = ts.saveToFile()
+}
+
+func (ts *TorrentStorage) Save() error {
+	ts.mu.RLock()
+	defer ts.mu.RUnlock()
+	return ts.saveToFile()
+}
+
+// saveToFile is a helper function to write the current state to the JSON file
+func (ts *TorrentStorage) saveToFile() error {
+	data, err := json.MarshalIndent(ts.torrents, "", "  ")
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(ts.filename, data, 0644)
+}