uptodownbot/internal/repo/store.go

package repo

import (
	"database/sql"
	"embed"
	"log/slog"

	_ "modernc.org/sqlite"

	"github.com/pressly/goose/v3"

	"uptodownBot/internal/domain"
)

//go:embed migrations/*.sql
var migrationFS embed.FS

type Store struct {
	db *sql.DB
}

func NewStore(path string) (*Store, error) {
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, err
	}
	if err := db.Ping(); err != nil {
		return nil, err
	}
	if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
		slog.Warn("failed to set WAL mode", "error", err)
	}
	if err := runMigrations(db); err != nil {
		return nil, err
	}
	return &Store{db: db}, nil
}

func runMigrations(db *sql.DB) error {
	goose.SetBaseFS(migrationFS)
	if err := goose.SetDialect("sqlite3"); err != nil {
		return err
	}
	if err := goose.Up(db, "migrations"); err != nil {
		return err
	}
	return nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

func (s *Store) GetOrCreateUser(chatID int64) (int64, error) {
	if _, err := s.db.Exec(
		"INSERT OR IGNORE INTO users (chat_id) VALUES (?)",
		chatID,
	); err != nil {
		return 0, err
	}
	var id int64
	err := s.db.QueryRow("SELECT id FROM users WHERE chat_id = ?", chatID).Scan(&id)
	return id, err
}

func (s *Store) GetOrCreatePreferences(userID int64) (*domain.UserPreferences, error) {
	if _, err := s.db.Exec(
		"INSERT OR IGNORE INTO user_preferences (user_id) VALUES (?)",
		userID,
	); err != nil {
		return nil, err
	}
	return s.getPreferences(userID)
}

func (s *Store) getPreferences(userID int64) (*domain.UserPreferences, error) {
	row := s.db.QueryRow(`
		SELECT default_media_type, default_quality, default_container, default_language
		FROM user_preferences WHERE user_id = ?
	`, userID)
	var p domain.UserPreferences
	var mediaType string
	if err := row.Scan(&mediaType, &p.DefaultQuality, &p.DefaultContainer, &p.DefaultLanguage); err != nil {
		return nil, err
	}
	switch mediaType {
	case "audio":
		p.DefaultMediaType = domain.MediaTypeAudio
	case "video":
		p.DefaultMediaType = domain.MediaTypeVideo
	default:
		p.DefaultMediaType = domain.MediaTypeVideoAudio
	}
	return &p, nil
}

func mediaTypeToString(mt domain.MediaType) string {
	switch mt {
	case domain.MediaTypeAudio:
		return "audio"
	case domain.MediaTypeVideo:
		return "video"
	default:
		return "video_audio"
	}
}

func (s *Store) UpdatePreferences(userID int64, prefs *domain.UserPreferences) error {
	_, err := s.db.Exec(`
		UPDATE user_preferences SET default_media_type=?, default_quality=?,
		       default_container=?, default_language=?
		WHERE user_id=?
	`, mediaTypeToString(prefs.DefaultMediaType), prefs.DefaultQuality,
		prefs.DefaultContainer, prefs.DefaultLanguage, userID)
	return err
}

func (s *Store) GetOrCreateQuotas(userID int64) (*domain.UserQuotas, error) {
	if _, err := s.db.Exec(
		"INSERT OR IGNORE INTO user_quotas (user_id) VALUES (?)",
		userID,
	); err != nil {
		return nil, err
	}
	return s.getQuotas(userID)
}

func (s *Store) getQuotas(userID int64) (*domain.UserQuotas, error) {
	row := s.db.QueryRow(`
		SELECT daily_limit_mb, daily_used_mb, daily_date,
		       weekly_limit_mb, weekly_used_mb, week_start,
		       monthly_limit_mb, monthly_used_mb, month_start
		FROM user_quotas WHERE user_id = ?
	`, userID)
	var q domain.UserQuotas
	err := row.Scan(
		&q.DailyLimitMB, &q.DailyUsedMB, &q.DailyDate,
		&q.WeeklyLimitMB, &q.WeeklyUsedMB, &q.WeekStart,
		&q.MonthlyLimitMB, &q.MonthlyUsedMB, &q.MonthStart,
	)
	return &q, err
}

func (s *Store) UpdateQuotas(userID int64, q *domain.UserQuotas) error {
	_, err := s.db.Exec(`
		UPDATE user_quotas SET daily_limit_mb=?, daily_used_mb=?, daily_date=?,
		       weekly_limit_mb=?, weekly_used_mb=?, week_start=?,
		       monthly_limit_mb=?, monthly_used_mb=?, month_start=?,
		       updated_at=datetime('now')
		WHERE user_id=?
	`, q.DailyLimitMB, q.DailyUsedMB, q.DailyDate,
		q.WeeklyLimitMB, q.WeeklyUsedMB, q.WeekStart,
		q.MonthlyLimitMB, q.MonthlyUsedMB, q.MonthStart,
		userID)
	return err
}

func (s *Store) AddHistory(userID int64, record *domain.DownloadRecord) error {
	_, err := s.db.Exec(`
		INSERT INTO download_history (user_id, url, title, media_type, format_id, container, file_size, status)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?)
	`, userID, record.URL, record.Title, record.MediaType, record.FormatID,
		record.Container, record.FileSize, record.Status)
	return err
}

func (s *Store) GetHistory(userID int64, limit, offset int) ([]domain.DownloadRecord, error) {
	rows, err := s.db.Query(`
		SELECT id, user_id, url, title, media_type, format_id, container, file_size, status, created_at
		FROM download_history
		WHERE user_id = ?
		ORDER BY created_at DESC
		LIMIT ? OFFSET ?
	`, userID, limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var records []domain.DownloadRecord
	for rows.Next() {
		var r domain.DownloadRecord
		if err := rows.Scan(
			&r.ID, &r.UserID, &r.URL, &r.Title, &r.MediaType,
			&r.FormatID, &r.Container, &r.FileSize, &r.Status, &r.CreatedAt,
		); err != nil {
			return nil, err
		}
		records = append(records, r)
	}
	return records, rows.Err()
}

func (s *Store) GetHistoryCount(userID int64) (int, error) {
	var count int
	err := s.db.QueryRow("SELECT COUNT(1) FROM download_history WHERE user_id=?", userID).Scan(&count)
	return count, err
}

func (s *Store) DeleteUserData(userID int64) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	if _, err := tx.Exec("DELETE FROM download_history WHERE user_id=?", userID); err != nil {
		return err
	}
	if _, err := tx.Exec("DELETE FROM user_quotas WHERE user_id=?", userID); err != nil {
		return err
	}
	if _, err := tx.Exec("DELETE FROM user_preferences WHERE user_id=?", userID); err != nil {
		return err
	}
	if _, err := tx.Exec("DELETE FROM users WHERE id=?", userID); err != nil {
		return err
	}
	return tx.Commit()
}