music-agregator/internal/torrent/qbittorrent.go

package torrent

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mime/multipart"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"strings"
	"sync"
)

type QBittorrentClient struct {
	baseURL   string
	username  string
	password  string
	client    *http.Client
	connected bool
	mu        sync.RWMutex
}

type qbTorrent struct {
	Hash     string  `json:"hash"`
	Name     string  `json:"name"`
	Size     int64   `json:"size"`
	Progress float64 `json:"progress"`
	DLSpeed  int64   `json:"dlspeed"`
	UPSpeed  int64   `json:"upspeed"`
	State    string  `json:"state"`
	SavePath string  `json:"save_path"`
}

func NewQBittorrentClient(baseURL, username, password string) (*QBittorrentClient, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}

	return &QBittorrentClient{
		baseURL:  strings.TrimRight(baseURL, "/"),
		username: username,
		password: password,
		client:   &http.Client{Jar: jar},
	}, nil
}

func (c *QBittorrentClient) apiURL(path string) string {
	return fmt.Sprintf("%s/api/v2%s", c.baseURL, path)
}

func (c *QBittorrentClient) mapState(state string) TorrentState {
	switch state {
	case "downloading", "forcedDL", "metaDL", "allocating":
		return StateDownloading
	case "uploading", "forcedUP", "stalledUP":
		return StateSeeding
	case "pausedDL", "pausedUP":
		return StatePaused
	case "queuedDL", "queuedUP":
		return StateQueued
	case "checkingDL", "checkingUP", "checkingResumeData":
		return StateChecking
	case "error", "missingFiles":
		return StateError
	default:
		return StateUnknown
	}
}

func (c *QBittorrentClient) mapTorrent(t qbTorrent) TorrentInfo {
	size := uint64(0)
	if t.Size > 0 {
		size = uint64(t.Size)
	}
	dlSpeed := uint64(0)
	if t.DLSpeed > 0 {
		dlSpeed = uint64(t.DLSpeed)
	}
	upSpeed := uint64(0)
	if t.UPSpeed > 0 {
		upSpeed = uint64(t.UPSpeed)
	}

	return TorrentInfo{
		Hash:          t.Hash,
		Name:          t.Name,
		Size:          size,
		Progress:      t.Progress,
		DownloadSpeed: dlSpeed,
		UploadSpeed:   upSpeed,
		State:         c.mapState(t.State),
		SavePath:      t.SavePath,
	}
}

func (c *QBittorrentClient) ensureConnected() error {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if !c.connected {
		return ErrNotConnected
	}
	return nil
}

func (c *QBittorrentClient) Connect(ctx context.Context) error {
	data := url.Values{}
	data.Set("username", c.username)
	data.Set("password", c.password)

	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/auth/login"), strings.NewReader(data.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := c.client.Do(req)
	if err != nil {
		return fmt.Errorf("%w: %v", ErrConnectionFailed, err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if string(body) == "Ok." {
		c.mu.Lock()
		c.connected = true
		c.mu.Unlock()
		return nil
	}

	return ErrAuthFailed
}

func (c *QBittorrentClient) Disconnect(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/auth/logout"), nil)
	if err != nil {
		return err
	}

	c.client.Do(req)
	c.mu.Lock()
	c.connected = false
	c.mu.Unlock()
	return nil
}

func (c *QBittorrentClient) ListTorrents(ctx context.Context) ([]TorrentInfo, error) {
	if err := c.ensureConnected(); err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, "GET", c.apiURL("/torrents/info"), nil)
	if err != nil {
		return nil, err
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var torrents []qbTorrent
	if err := json.NewDecoder(resp.Body).Decode(&torrents); err != nil {
		return nil, err
	}

	result := make([]TorrentInfo, len(torrents))
	for i, t := range torrents {
		result[i] = c.mapTorrent(t)
	}
	return result, nil
}

func (c *QBittorrentClient) GetTorrent(ctx context.Context, hash string) (*TorrentInfo, error) {
	if err := c.ensureConnected(); err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, "GET", c.apiURL("/torrents/info")+"?hashes="+hash, nil)
	if err != nil {
		return nil, err
	}

	resp, err := c.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	var torrents []qbTorrent
	if err := json.NewDecoder(resp.Body).Decode(&torrents); err != nil {
		return nil, err
	}

	if len(torrents) == 0 {
		return nil, ErrTorrentNotFound
	}

	info := c.mapTorrent(torrents[0])
	return &info, nil
}

func (c *QBittorrentClient) AddTorrentURL(ctx context.Context, torrentURL string, savePath *string) error {
	if err := c.ensureConnected(); err != nil {
		return err
	}

	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)
	w.WriteField("urls", torrentURL)
	if savePath != nil {
		w.WriteField("savepath", *savePath)
	}
	w.Close()

	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/torrents/add"), &buf)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", w.FormDataContentType())

	resp, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if !statusOK(resp.StatusCode) {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("%w: %s", ErrInvalidRequest, string(body))
	}

	return nil
}

func (c *QBittorrentClient) AddTorrentFile(ctx context.Context, data []byte, savePath *string) error {
	if err := c.ensureConnected(); err != nil {
		return err
	}

	var buf bytes.Buffer
	w := multipart.NewWriter(&buf)

	part, err := w.CreateFormFile("torrents", "torrent.torrent")
	if err != nil {
		return err
	}
	part.Write(data)

	if savePath != nil {
		w.WriteField("savepath", *savePath)
	}
	w.Close()

	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/torrents/add"), &buf)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", w.FormDataContentType())

	resp, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if !statusOK(resp.StatusCode) {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("%w: %s", ErrInvalidRequest, string(body))
	}

	return nil
}

func (c *QBittorrentClient) RemoveTorrent(ctx context.Context, hash string, deleteFiles bool) error {
	if err := c.ensureConnected(); err != nil {
		return err
	}

	data := url.Values{}
	data.Set("hashes", hash)
	if deleteFiles {
		data.Set("deleteFiles", "true")
	} else {
		data.Set("deleteFiles", "false")
	}

	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/torrents/delete"), strings.NewReader(data.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := c.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if !statusOK(resp.StatusCode) {
		return ErrTorrentNotFound
	}

	return nil
}

func (c *QBittorrentClient) PauseTorrent(ctx context.Context, hash string) error {
	if err := c.ensureConnected(); err != nil {
		return err
	}

	data := url.Values{}
	data.Set("hashes", hash)

	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/torrents/pause"), strings.NewReader(data.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	c.client.Do(req)
	return nil
}

func (c *QBittorrentClient) ResumeTorrent(ctx context.Context, hash string) error {
	if err := c.ensureConnected(); err != nil {
		return err
	}

	data := url.Values{}
	data.Set("hashes", hash)

	req, err := http.NewRequestWithContext(ctx, "POST", c.apiURL("/torrents/resume"), strings.NewReader(data.Encode()))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	c.client.Do(req)
	return nil
}

func statusOK(code int) bool {
	return code >= 200 && code < 300
}