music-agregator/internal/database/download_repository.go

package database

import (
	"context"
	"fmt"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"
)

type Download struct {
	ID           string
	TorrentID    string
	AlbumID      string
	Format       string
	Quality      string
	State        string
	QbitHash     string
	SavePath     *string
	ErrorMessage *string
	QueuedAt     time.Time
	StartedAt    *time.Time
	CompletedAt  *time.Time
	CreatedAt    time.Time
	UpdatedAt    time.Time
}

type DownloadRepository struct {
	pool *pgxpool.Pool
}

func NewDownloadRepository(pool *pgxpool.Pool) *DownloadRepository {
	return &DownloadRepository{pool: pool}
}

func (r *DownloadRepository) Create(ctx context.Context, d *Download) error {
	err := r.pool.QueryRow(ctx,
		`INSERT INTO downloads (torrent_id, album_id, format, quality, state, qbit_hash, save_path)
		 VALUES ($1, $2, $3, $4, $5, $6, $7)
		 RETURNING id, queued_at, created_at, updated_at`,
		d.TorrentID, d.AlbumID, d.Format, d.Quality, d.State, d.QbitHash, d.SavePath,
	).Scan(&d.ID, &d.QueuedAt, &d.CreatedAt, &d.UpdatedAt)
	if err != nil {
		return fmt.Errorf("creating download: %w", err)
	}
	return nil
}

func (r *DownloadRepository) UpdateState(ctx context.Context, id string, state string) error {
	_, err := r.pool.Exec(ctx,
		`UPDATE downloads SET state = $1, updated_at = NOW() WHERE id = $2`, state, id,
	)
	if err != nil {
		return fmt.Errorf("updating download state: %w", err)
	}
	return nil
}

func (r *DownloadRepository) SetStarted(ctx context.Context, id string) error {
	_, err := r.pool.Exec(ctx,
		`UPDATE downloads SET state = 'downloading', started_at = NOW(), updated_at = NOW() WHERE id = $1`, id,
	)
	if err != nil {
		return fmt.Errorf("setting download started: %w", err)
	}
	return nil
}

func (r *DownloadRepository) SetCompleted(ctx context.Context, id string, savePath string) error {
	_, err := r.pool.Exec(ctx,
		`UPDATE downloads SET state = 'completed', save_path = $1, completed_at = NOW(), updated_at = NOW() WHERE id = $2`, savePath, id,
	)
	if err != nil {
		return fmt.Errorf("setting download completed: %w", err)
	}
	return nil
}

func (r *DownloadRepository) SetFailed(ctx context.Context, id string, errorMsg string) error {
	_, err := r.pool.Exec(ctx,
		`UPDATE downloads SET state = 'failed', error_message = $1, updated_at = NOW() WHERE id = $2`, errorMsg, id,
	)
	if err != nil {
		return fmt.Errorf("setting download failed: %w", err)
	}
	return nil
}

func (r *DownloadRepository) GetByAlbumID(ctx context.Context, albumID string) ([]*Download, error) {
	rows, err := r.pool.Query(ctx,
		`SELECT id, torrent_id, album_id, format, quality, state, qbit_hash, save_path, error_message, queued_at, started_at, completed_at, created_at, updated_at
		 FROM downloads WHERE album_id = $1 ORDER BY created_at DESC`, albumID,
	)
	if err != nil {
		return nil, fmt.Errorf("listing downloads: %w", err)
	}
	defer rows.Close()

	var downloads []*Download
	for rows.Next() {
		d := &Download{}
		if err := rows.Scan(&d.ID, &d.TorrentID, &d.AlbumID, &d.Format, &d.Quality, &d.State, &d.QbitHash, &d.SavePath, &d.ErrorMessage, &d.QueuedAt, &d.StartedAt, &d.CompletedAt, &d.CreatedAt, &d.UpdatedAt); err != nil {
			return nil, fmt.Errorf("scanning download: %w", err)
		}
		downloads = append(downloads, d)
	}
	return downloads, nil
}

func (r *DownloadRepository) GetActiveByTorrentID(ctx context.Context, torrentID string) (*Download, error) {
	d := &Download{}
	err := r.pool.QueryRow(ctx,
		`SELECT id, torrent_id, album_id, format, quality, state, qbit_hash, save_path, error_message, queued_at, started_at, completed_at, created_at, updated_at
		 FROM downloads WHERE torrent_id = $1 AND state NOT IN ('failed')
		 ORDER BY created_at DESC LIMIT 1`, torrentID,
	).Scan(&d.ID, &d.TorrentID, &d.AlbumID, &d.Format, &d.Quality, &d.State, &d.QbitHash, &d.SavePath, &d.ErrorMessage, &d.QueuedAt, &d.StartedAt, &d.CompletedAt, &d.CreatedAt, &d.UpdatedAt)
	if err != nil {
		return nil, fmt.Errorf("getting active download by torrent: %w", err)
	}
	return d, nil
}

func (r *DownloadRepository) GetActive(ctx context.Context) ([]*Download, error) {
	rows, err := r.pool.Query(ctx,
		`SELECT id, torrent_id, album_id, format, quality, state, qbit_hash, save_path, error_message, queued_at, started_at, completed_at, created_at, updated_at
		 FROM downloads WHERE state IN ('pending', 'downloading') ORDER BY created_at`,
	)
	if err != nil {
		return nil, fmt.Errorf("listing active downloads: %w", err)
	}
	defer rows.Close()

	var downloads []*Download
	for rows.Next() {
		d := &Download{}
		if err := rows.Scan(&d.ID, &d.TorrentID, &d.AlbumID, &d.Format, &d.Quality, &d.State, &d.QbitHash, &d.SavePath, &d.ErrorMessage, &d.QueuedAt, &d.StartedAt, &d.CompletedAt, &d.CreatedAt, &d.UpdatedAt); err != nil {
			return nil, fmt.Errorf("scanning download: %w", err)
		}
		downloads = append(downloads, d)
	}
	return downloads, nil
}

func (r *DownloadRepository) GetByID(ctx context.Context, id string) (*Download, error) {
	d := &Download{}
	err := r.pool.QueryRow(ctx,
		`SELECT id, torrent_id, album_id, format, quality, state, qbit_hash, save_path, error_message, queued_at, started_at, completed_at, created_at, updated_at
		 FROM downloads WHERE id = $1`, id,
	).Scan(&d.ID, &d.TorrentID, &d.AlbumID, &d.Format, &d.Quality, &d.State, &d.QbitHash, &d.SavePath, &d.ErrorMessage, &d.QueuedAt, &d.StartedAt, &d.CompletedAt, &d.CreatedAt, &d.UpdatedAt)
	if err != nil {
		return nil, fmt.Errorf("getting download by id: %w", err)
	}
	return d, nil
}

func (r *DownloadRepository) HasAlbumInQuality(ctx context.Context, albumID string, format string, quality string) (bool, error) {
	var exists bool
	err := r.pool.QueryRow(ctx,
		`SELECT EXISTS(
			SELECT 1 FROM downloads
			WHERE album_id = $1 AND format = $2 AND quality = $3 AND state IN ('completed', 'seeding')
		)`, albumID, format, quality,
	).Scan(&exists)
	if err != nil {
		return false, fmt.Errorf("checking album quality: %w", err)
	}
	return exists, nil
}