diff --git a/internal/dl/ytdlp.go b/internal/dl/ytdlp.go index ee309d1..da25a68 100644 --- a/internal/dl/ytdlp.go +++ b/internal/dl/ytdlp.go @@ -10,6 +10,7 @@ import ( "os/exec" "path/filepath" "strings" + "sync" "syscall" "time" @@ -20,6 +21,36 @@ import ( type Client struct { binaryPath string cookiesFile string + cache map[string]*cachedMediaInfo + cacheMu sync.Mutex +} + +type cachedMediaInfo struct { + info *domain.MediaInfo + ts time.Time +} + +const cacheTTL = 30 * time.Minute +const cacheMaxEntries = 200 + +func (c *Client) getCached(url string) *domain.MediaInfo { + c.cacheMu.Lock() + defer c.cacheMu.Unlock() + entry, ok := c.cache[url] + if !ok || time.Since(entry.ts) > cacheTTL { + delete(c.cache, url) + return nil + } + return entry.info +} + +func (c *Client) setCache(url string, info *domain.MediaInfo) { + c.cacheMu.Lock() + defer c.cacheMu.Unlock() + if len(c.cache) >= cacheMaxEntries { + c.cache = make(map[string]*cachedMediaInfo) + } + c.cache[url] = &cachedMediaInfo{info: info, ts: time.Now()} } // NewClient creates a new yt-dlp client. It checks for the binary. @@ -29,7 +60,11 @@ func NewClient(cookiesFile string) (*Client, error) { return nil, fmt.Errorf("yt-dlp not found in PATH: %w", err) } slog.Info("yt-dlp found", "path", path) - return &Client{binaryPath: path, cookiesFile: cookiesFile}, nil + return &Client{ + binaryPath: path, + cookiesFile: cookiesFile, + cache: make(map[string]*cachedMediaInfo), + }, nil } type ytdlpFormat struct { @@ -70,6 +105,11 @@ type ytdlpPlaylistEntry struct { // GetMediaInfo fetches format information for a given URL. func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) { + if cached := c.getCached(url); cached != nil { + slog.Debug("cache hit", "url", url) + return cached, nil + } + args := []string{ "-J", "--no-download", "--no-warnings", @@ -106,6 +146,7 @@ func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) { if mi.IsPlaylist { mi.PlaylistCount = info.PlaylistCount mi.PlaylistEntries = c.parsePlaylistEntries(info.Entries) + c.setCache(url, mi) return mi, nil } @@ -124,11 +165,17 @@ func (c *Client) GetMediaInfo(url string) (*domain.MediaInfo, error) { } mi.EstimatedSize = maxSize + c.setCache(url, mi) return mi, nil } // GetPlaylistInfo fetches a flat playlist listing (fast, no format detail). func (c *Client) GetPlaylistInfo(url string) (*domain.MediaInfo, error) { + if cached := c.getCached(url); cached != nil { + slog.Debug("cache hit", "url", url) + return cached, nil + } + args := []string{ "-J", "--flat-playlist", "--no-download", "--no-warnings", @@ -159,6 +206,7 @@ func (c *Client) GetPlaylistInfo(url string) (*domain.MediaInfo, error) { PlaylistEntries: c.parsePlaylistEntries(info.Entries), } + c.setCache(url, mi) return mi, nil } @@ -376,6 +424,36 @@ func (c *Client) SearchYoutube(query string) (url, title string, err error) { return info.Entries[0].URL, info.Entries[0].Title, nil } +// Search does a yt-dlp search and returns entries. +func (c *Client) Search(query, source string, limit int) ([]domain.PlaylistEntry, error) { + searchPrefix := "ytsearch" + if source == "ytmusic" { + searchPrefix = "ytmusicsearch" + } + + args := []string{ + "--flat-playlist", "-J", "--no-download", + "--no-warnings", + fmt.Sprintf("%s%d:%s", searchPrefix, limit, query), + } + if c.cookiesFile != "" { + args = append(args, "--cookies", c.cookiesFile) + } + + slog.Info("searching", "query", query, "source", source) + out, err := c.run(args) + if err != nil { + return nil, fmt.Errorf("search failed: %w", err) + } + + var info ytdlpInfo + if err := json.Unmarshal(out, &info); err != nil { + return nil, fmt.Errorf("parse search result: %w", err) + } + + return c.parsePlaylistEntries(info.Entries), nil +} + func (c *Client) run(args []string) ([]byte, error) { cmd := exec.Command(c.binaryPath, args...) var stderr bytes.Buffer diff --git a/internal/domain/types.go b/internal/domain/types.go index 66efc84..413a71c 100644 --- a/internal/domain/types.go +++ b/internal/domain/types.go @@ -177,3 +177,12 @@ type PlaylistState struct { PerPage int Selected map[string]bool } + +// SearchState tracks the user's current search results flow. +type SearchState struct { + Query string + Entries []PlaylistEntry + Page int + PerPage int + Source string // "youtube" or "ytmusic" +} diff --git a/internal/handler/download.go b/internal/handler/download.go index af68c2a..401b1c8 100644 --- a/internal/handler/download.go +++ b/internal/handler/download.go @@ -29,6 +29,7 @@ type stepState struct { Container string FormatIDs []string PlaylistState *domain.PlaylistState + SearchState *domain.SearchState } func (h *Handler) getUserStepState(chatID int64) *stepState { diff --git a/internal/handler/handler.go b/internal/handler/handler.go index a279411..69304a1 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -23,8 +23,9 @@ import ( type PendingKind string const ( - PendingURL PendingKind = "url" - PendingDRM PendingKind = "drm" + PendingURL PendingKind = "url" + PendingDRM PendingKind = "drm" + PendingSearch PendingKind = "search" ) type PendingInput struct { @@ -185,6 +186,10 @@ func (h *Handler) routeCommand(ctx context.Context, msg *models.Message, userID h.handleHistory(ctx, msg, userID) case "/download": h.handleDownloadCommand(ctx, msg, userID) + case "/search": + h.handleSearchCommand(ctx, msg, userID) + case "/music": + h.handleMusicSearchCommand(ctx, msg, userID) default: h.sendText(ctx, msg.Chat.ID, "Unknown command") } @@ -226,6 +231,9 @@ func (h *Handler) HandleCallback(ctx context.Context, cb *models.CallbackQuery) case data == "download": h.handleDownloadPrompt(ctx, chatID, msgID, cb.ID, userID) h.answerCb(ctx, cb.ID) + case data == "search": + h.promptForInput(ctx, chatID, msgID, userID, PendingSearch, "Enter a YouTube search query:", map[string]string{"source": "youtube"}) + h.answerCb(ctx, cb.ID) case data == "settings": h.settingsCallback(ctx, chatID, msgID, cb.ID, userID) case data == "help": @@ -280,6 +288,8 @@ func (h *Handler) routePrefixedCallback(ctx context.Context, chatID int64, msgID h.handleSettingsSet(ctx, chatID, msgID, userID, data[13:]) case strings.HasPrefix(data, "settings_"): h.handleSettingsSelection(ctx, chatID, msgID, userID, data[9:]) + case strings.HasPrefix(data, "search_"): + h.handleSearchCallback(ctx, chatID, msgID, userID, data[7:]) } } @@ -302,6 +312,12 @@ func (h *Handler) processPendingInput(ctx context.Context, msg *models.Message, h.handleURLInput(ctx, msg.Chat.ID, text, userID) case PendingDRM: h.handleDRMSearch(ctx, msg.Chat.ID, text, userID) + case PendingSearch: + source := "youtube" + if pi.Data != nil && pi.Data["source"] != "" { + source = pi.Data["source"] + } + h.handleSearchQuery(ctx, msg.Chat.ID, text, userID, source) default: h.sendText(ctx, msg.Chat.ID, "Unknown input type") } @@ -457,6 +473,8 @@ func (h *Handler) handleHelp(ctx context.Context, msg *models.Message) { /settings - Open settings /help - Show this help /history - View download history +/search - Search YouTube +/music - Search YouTube Music Just send me a URL to start downloading.` h.sendText(ctx, msg.Chat.ID, text) diff --git a/internal/handler/keyboard.go b/internal/handler/keyboard.go index 88342e8..a0f0303 100644 --- a/internal/handler/keyboard.go +++ b/internal/handler/keyboard.go @@ -14,6 +14,7 @@ func mainKeyboard() models.InlineKeyboardMarkup { InlineKeyboard: [][]models.InlineKeyboardButton{ { {Text: "Download", CallbackData: "download"}, + {Text: "Search", CallbackData: "search"}, {Text: "Settings", CallbackData: "settings"}, }, { diff --git a/internal/handler/search.go b/internal/handler/search.go new file mode 100644 index 0000000..02c1e2f --- /dev/null +++ b/internal/handler/search.go @@ -0,0 +1,136 @@ +package handler + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/go-telegram/bot/models" + + "uptodownBot/internal/domain" +) + +func (h *Handler) handleSearchCommand(ctx context.Context, msg *models.Message, userID int64) { + h.promptForInput(ctx, msg.Chat.ID, 0, userID, PendingSearch, + "Enter a YouTube search query:", map[string]string{"source": "youtube"}) +} + +func (h *Handler) handleMusicSearchCommand(ctx context.Context, msg *models.Message, userID int64) { + h.promptForInput(ctx, msg.Chat.ID, 0, userID, PendingSearch, + "Enter a YouTube Music search query:", map[string]string{"source": "ytmusic"}) +} + +func (h *Handler) handleSearchQuery(ctx context.Context, chatID int64, text string, userID int64, source string) { + results, err := h.YtDlp.Search(text, source, 10) + if err != nil { + h.sendText(ctx, chatID, fmt.Sprintf("Search failed: %s", err.Error())) + return + } + if len(results) == 0 { + h.sendText(ctx, chatID, "No results found.") + return + } + + state := &domain.SearchState{ + Query: text, + Entries: results, + Page: 0, + PerPage: 5, + Source: source, + } + ss := &stepState{SearchState: state} + h.setUserStepState(chatID, ss) + + kb := searchResultsKeyboard(state) + h.sendWithKB(ctx, chatID, fmt.Sprintf("Search results for '%s':", text), kb) +} + +func (h *Handler) handleSearchCallback(ctx context.Context, chatID int64, msgID int, userID int64, data string) { + switch { + case data == "refresh": + ss := h.getUserStepState(chatID) + if ss == nil || ss.SearchState == nil { + return + } + h.handleSearchQuery(ctx, chatID, ss.SearchState.Query, userID, ss.SearchState.Source) + case strings.HasPrefix(data, "page_"): + h.handleSearchPage(ctx, chatID, msgID, userID, data[5:]) + case strings.HasPrefix(data, "select_"): + h.handleSearchSelect(ctx, chatID, msgID, userID, data[7:]) + } +} + +func (h *Handler) handleSearchPage(ctx context.Context, chatID int64, msgID int, userID int64, pageStr string) { + ss := h.getUserStepState(chatID) + if ss == nil || ss.SearchState == nil { + return + } + page, err := strconv.Atoi(pageStr) + if err != nil { + return + } + totalPages := (len(ss.SearchState.Entries) + ss.SearchState.PerPage - 1) / ss.SearchState.PerPage + if page < 0 || page >= totalPages { + return + } + ss.SearchState.Page = page + h.setUserStepState(chatID, ss) + + kb := searchResultsKeyboard(ss.SearchState) + h.editText(ctx, chatID, msgID, fmt.Sprintf("Search results for '%s':", ss.SearchState.Query), &kb) +} + +func (h *Handler) handleSearchSelect(ctx context.Context, chatID int64, msgID int, userID int64, entryURL string) { + h.clearStepState(chatID) + h.editText(ctx, chatID, msgID, "Fetching media info...", nil) + h.handleURLInput(ctx, chatID, entryURL, userID) +} + +func searchResultsKeyboard(state *domain.SearchState) models.InlineKeyboardMarkup { + totalPages := (len(state.Entries) + state.PerPage - 1) / state.PerPage + if totalPages < 1 { + totalPages = 1 + } + + rows := [][]models.InlineKeyboardButton{} + start := state.Page * state.PerPage + end := start + state.PerPage + if end > len(state.Entries) { + end = len(state.Entries) + } + + for i := start; i < end; i++ { + e := state.Entries[i] + title := e.Title + if len(title) > 40 { + title = title[:37] + "..." + } + label := title + if e.Duration > 0 { + label += fmt.Sprintf(" (%s)", formatDuration(e.Duration)) + } + rows = append(rows, []models.InlineKeyboardButton{ + {Text: label, CallbackData: "search_select_" + e.URL}, + }) + } + + navRow := []models.InlineKeyboardButton{} + navRow = append(navRow, models.InlineKeyboardButton{ + Text: "<", CallbackData: fmt.Sprintf("search_page_%d", state.Page-1), + }) + navRow = append(navRow, models.InlineKeyboardButton{ + Text: fmt.Sprintf("Page %d/%d", state.Page+1, totalPages), CallbackData: "noop", + }) + navRow = append(navRow, models.InlineKeyboardButton{ + Text: ">", CallbackData: fmt.Sprintf("search_page_%d", state.Page+1), + }) + rows = append(rows, navRow) + + rows = append(rows, []models.InlineKeyboardButton{ + {Text: "New Search", CallbackData: "search_refresh"}, + {Text: "Cancel", CallbackData: "cancel_dl"}, + }) + + return models.InlineKeyboardMarkup{InlineKeyboard: rows} +}