streamrip-go/cmd/rip/search.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/AlecAivazis/survey/v2"

	"streamrip-go/internal/jsonutil"
)

type searchResult struct {
	ID         string
	Title      string
	Artist     string
	Album      string
	Date       string
	Releases   int
	TrackCount int
	Explicit   bool
}

type searchOptions struct {
	query      string
	limit      int
	ignoreDB   bool
	noDownload bool
	first      bool
	outputFile string
}

func parseSearchArgs(args []string, defaultLimit int) (searchOptions, error) {
	if defaultLimit <= 0 {
		defaultLimit = 20
	}
	limit := defaultLimit
	parts := make([]string, 0, len(args))
	ignoreDB := false
	noDownload := false
	first := false
	outputFile := ""
	for i := 0; i < len(args); i++ {
		if args[i] == "--" {
			if i+1 < len(args) {
				parts = append(parts, args[i+1:]...)
			}
			break
		}
		switch args[i] {
		case "--force", "--ignore-db":
			ignoreDB = true
			continue
		case "--no-download":
			noDownload = true
			continue
		case "--first":
			first = true
			continue
		case "--output-file":
			if i+1 >= len(args) {
				return searchOptions{}, fmt.Errorf("--output-file requires a path")
			}
			outputFile = strings.TrimSpace(args[i+1])
			if outputFile == "" {
				return searchOptions{}, fmt.Errorf("--output-file requires a non-empty path")
			}
			i++
			continue
		case "--num-results":
			if i+1 >= len(args) {
				return searchOptions{}, fmt.Errorf("--num-results requires a value")
			}
			v, err := strconv.Atoi(args[i+1])
			if err != nil || v <= 0 {
				return searchOptions{}, fmt.Errorf("invalid --num-results value %q", args[i+1])
			}
			limit = v
			i++
			continue
		}
		if args[i] == "--limit" {
			if i+1 >= len(args) {
				return searchOptions{}, fmt.Errorf("--limit requires a value")
			}
			v, err := strconv.Atoi(args[i+1])
			if err != nil || v <= 0 {
				return searchOptions{}, fmt.Errorf("invalid --limit value %q", args[i+1])
			}
			limit = v
			i++
			continue
		}
		if strings.HasPrefix(args[i], "-") {
			return searchOptions{}, fmt.Errorf("unknown option %q", args[i])
		}
		parts = append(parts, args[i])
	}
	return searchOptions{
		query:      strings.TrimSpace(strings.Join(parts, " ")),
		limit:      limit,
		ignoreDB:   ignoreDB,
		noDownload: noDownload,
		first:      first,
		outputFile: outputFile,
	}, nil
}

func promptSearchSelection(results []searchResult) ([]int, error) {
	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("Select results to download (e.g. 1,3-5; a=all; q=cancel): ")
		line, err := reader.ReadString('\n')
		if err != nil {
			return nil, err
		}
		line = strings.TrimSpace(line)
		if line == "" || strings.EqualFold(line, "q") || strings.EqualFold(line, "quit") {
			return nil, nil
		}
		if strings.EqualFold(line, "a") || strings.EqualFold(line, "all") {
			out := make([]int, 0, len(results))
			for i := range results {
				out = append(out, i)
			}
			return out, nil
		}

		selected := map[int]struct{}{}
		chunks := strings.Split(line, ",")
		ok := true
		for _, raw := range chunks {
			part := strings.TrimSpace(raw)
			if part == "" {
				continue
			}
			if strings.Contains(part, "-") {
				bounds := strings.SplitN(part, "-", 2)
				if len(bounds) != 2 {
					ok = false
					break
				}
				start, err1 := strconv.Atoi(strings.TrimSpace(bounds[0]))
				end, err2 := strconv.Atoi(strings.TrimSpace(bounds[1]))
				if err1 != nil || err2 != nil || start <= 0 || end <= 0 || start > end {
					ok = false
					break
				}
				for i := start; i <= end; i++ {
					if i > len(results) {
						ok = false
						break
					}
					selected[i-1] = struct{}{}
				}
				if !ok {
					break
				}
				continue
			}
			idx, err := strconv.Atoi(part)
			if err != nil || idx <= 0 || idx > len(results) {
				ok = false
				break
			}
			selected[idx-1] = struct{}{}
		}

		if !ok || len(selected) == 0 {
			fmt.Println("Invalid selection, try again.")
			continue
		}

		out := make([]int, 0, len(selected))
		for idx := range selected {
			out = append(out, idx)
		}
		for i := 1; i < len(out); i++ {
			for j := i; j > 0 && out[j] < out[j-1]; j-- {
				out[j], out[j-1] = out[j-1], out[j]
			}
		}
		return out, nil
	}
}

func promptSearchSelectionMenu(source, mediaType, query string, results []searchResult) ([]int, error) {
	if len(results) == 0 {
		return nil, nil
	}

	prevTemplate := survey.MultiSelectQuestionTemplate
	survey.MultiSelectQuestionTemplate = streamripLikeSearchTemplate
	defer func() {
		survey.MultiSelectQuestionTemplate = prevTemplate
	}()

	labels := make([]string, 0, len(results))
	labelToIndex := map[string]int{}
	for i, r := range results {
		label := formatSearchOptionLabel(i+1, mediaType, r)
		labels = append(labels, label)
		labelToIndex[label] = i
	}
	pageSize := len(labels)
	if pageSize < 15 {
		pageSize = 15
	}
	if pageSize > 30 {
		pageSize = 30
	}

	selected := []string{}
	prompt := &survey.MultiSelect{
		Message: fmt.Sprintf("Results for %s '%s' from %s", mediaType, query, jsonutil.TitleCase(source)),
		Help:    "SPACE - select, ENTER - download, ESC - exit",
		Options: labels,
		Description: func(value string, index int) string {
			resultIndex, ok := labelToIndex[value]
			if !ok || resultIndex < 0 || resultIndex >= len(results) {
				return ""
			}
			return formatSearchDetails(mediaType, results[resultIndex])
		},
		PageSize: pageSize,
	}
	if err := survey.AskOne(prompt, &selected); err != nil {
		if strings.Contains(strings.ToLower(err.Error()), "interrupt") {
			return nil, nil
		}
		return nil, err
	}

	if len(selected) == 0 {
		return nil, nil
	}
	out := make([]int, 0, len(selected))
	for _, label := range selected {
		if idx, ok := labelToIndex[label]; ok {
			out = append(out, idx)
		}
	}
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j] < out[j-1]; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out, nil
}

const streamripLikeSearchTemplate = `
{{ color "default+hb"}}{{ .Message }}{{color "reset"}}
{{- if .Help }}
{{ .Help }}
{{- end }}
{{- range $ix, $option := .PageEntries }}
{{ if eq $.SelectedIndex $ix }}>{{ else }} {{ end }} [{{ if index $.Checked $option.Index }}x{{ else }} {{ end }}] {{ $option.Value }}
{{- end }}

preview
{{- if gt (len .PageEntries) 0 }}
{{ $current := index .PageEntries .SelectedIndex }}
{{ $.GetDescription $current }}
{{- end }}
`

func writeSearchResultsToFile(source, mediaType string, results []searchResult, path string) error {
	type outItem struct {
		Source    string `json:"source"`
		MediaType string `json:"media_type"`
		ID        string `json:"id"`
		Title     string `json:"title"`
	}
	out := make([]outItem, 0, len(results))
	for _, r := range results {
		out = append(out, outItem{Source: source, MediaType: mediaType, ID: r.ID, Title: r.Title})
	}
	b, err := json.MarshalIndent(out, "", "  ")
	if err != nil {
		return err
	}
	dir := filepath.Dir(path)
	if dir != "" && dir != "." {
		if err = os.MkdirAll(dir, 0o755); err != nil {
			return err
		}
	}
	return os.WriteFile(path, b, 0o644)
}

func isAllowedSearchSource(source string) bool {
	return source == "qobuz" || source == "tidal" || source == "deezer" || source == "yandex" || source == "soundcloud"
}

func isAllowedMediaType(mediaType string) bool {
	switch mediaType {
	case "track", "album", "playlist", "artist", "label", "video":
		return true
	default:
		return false
	}
}

func promptSearchInteractive(defaultLimit int) (string, string, searchOptions, error) {
	reader := bufio.NewReader(os.Stdin)

	read := func(prompt string) (string, error) {
		fmt.Print(prompt)
		line, err := reader.ReadString('\n')
		if err != nil {
			return "", err
		}
		return strings.TrimSpace(line), nil
	}

	for {
		source, err := read("Source [qobuz/tidal/deezer/yandex/soundcloud]: ")
		if err != nil {
			return "", "", searchOptions{}, err
		}
		source = strings.ToLower(source)
		if !isAllowedSearchSource(source) {
			fmt.Println("Invalid source.")
			continue
		}

		mediaType, err := read("Type [track/album/playlist/artist/label/video]: ")
		if err != nil {
			return "", "", searchOptions{}, err
		}
		mediaType = strings.ToLower(mediaType)
		if !isAllowedMediaType(mediaType) {
			fmt.Println("Invalid media type.")
			continue
		}
		if source == "soundcloud" && mediaType != "track" && mediaType != "playlist" {
			fmt.Println("SoundCloud search supports track and playlist only.")
			continue
		}
		if source == "yandex" && mediaType != "track" && mediaType != "album" && mediaType != "playlist" && mediaType != "artist" {
			fmt.Println("Yandex search supports track, album, playlist, and artist only.")
			continue
		}

		query, err := read("Query: ")
		if err != nil {
			return "", "", searchOptions{}, err
		}
		if strings.TrimSpace(query) == "" {
			fmt.Println("Query cannot be empty.")
			continue
		}

		limitRaw, err := read(fmt.Sprintf("Limit [%d]: ", defaultLimit))
		if err != nil {
			return "", "", searchOptions{}, err
		}
		limit := defaultLimit
		if strings.TrimSpace(limitRaw) != "" {
			v, convErr := strconv.Atoi(limitRaw)
			if convErr != nil || v <= 0 {
				fmt.Println("Invalid limit.")
				continue
			}
			limit = v
		}

		return source, mediaType, searchOptions{query: query, limit: limit}, nil
	}
}

func normalizeSearchResults(source, mediaType string, pages []map[string]any) []searchResult {
	results := make([]searchResult, 0)
	seen := map[string]struct{}{}
	appendUnique := func(r searchResult) {
		if strings.TrimSpace(r.ID) == "" || strings.TrimSpace(r.Title) == "" {
			return
		}
		key := r.ID
		if _, ok := seen[key]; ok {
			return
		}
		seen[key] = struct{}{}
		results = append(results, r)
	}
	for _, page := range pages {
		switch source {
		case "qobuz":
			key := mediaType + "s"
			bucket, ok := page[key].(map[string]any)
			if !ok {
				continue
			}
			items, ok := bucket["items"].([]any)
			if !ok {
				continue
			}
			for _, raw := range items {
				itm, ok := raw.(map[string]any)
				if !ok {
					continue
				}
				id := asString(itm["id"])
				title := asString(itm["title"])
				if title == "" {
					title = asString(itm["name"])
				}
				if version := asString(itm["version"]); version != "" {
					title += " (" + version + ")"
				}
				artist := nestedSearchString(itm, "artist", "name")
				if artist == "" {
					artist = nestedSearchString(itm, "performer", "name")
				}
				album := nestedSearchString(itm, "album", "title")
				trackCount := searchInt(itm["tracks_count"])
				if trackCount == 0 {
					trackCount = searchInt(itm["track_count"])
				}
				explicit := searchBool(itm["parental_warning"])
				date := firstNonEmpty(
					asString(itm["release_date_original"]),
					asString(itm["release_date"]),
					asString(itm["releaseDate"]),
					asString(itm["streamStartDate"]),
					nestedSearchString(itm, "album", "release_date_original"),
					nestedSearchString(itm, "album", "release_date"),
				)
				releases := 0
				if mediaType == "artist" {
					releases = firstPositiveInt(
						searchInt(itm["albums_count"]),
						searchInt(itm["albumsCount"]),
						nestedSearchInt(itm, "albums", "total"),
					)
				}
				appendUnique(searchResult{ID: id, Title: title, Artist: artist, Album: album, Date: date, Releases: releases, TrackCount: trackCount, Explicit: explicit})
			}
		case "tidal":
			items, ok := page["items"].([]any)
			if !ok {
				continue
			}
			for _, raw := range items {
				itm, ok := raw.(map[string]any)
				if !ok {
					continue
				}
				if wrapped, ok := itm["item"].(map[string]any); ok {
					itm = wrapped
				}
				id := asString(itm["id"])
				title := asString(itm["title"])
				if title == "" {
					title = asString(itm["name"])
				}
				artist := nestedSearchString(itm, "artist", "name")
				if artist == "" {
					if artists, ok := itm["artists"].([]any); ok && len(artists) > 0 {
						if a0, ok := artists[0].(map[string]any); ok {
							artist = asString(a0["name"])
						}
					}
				}
				album := nestedSearchString(itm, "album", "title")
				trackCount := searchInt(itm["numberOfTracks"])
				if trackCount == 0 {
					trackCount = searchInt(itm["tracks_count"])
				}
				explicit := searchBool(itm["explicit"])
				date := firstNonEmpty(
					asString(itm["releaseDate"]),
					asString(itm["streamStartDate"]),
					asString(itm["release_date"]),
					nestedSearchString(itm, "album", "releaseDate"),
					nestedSearchString(itm, "album", "release_date"),
				)
				releases := 0
				if mediaType == "artist" {
					releases = firstPositiveInt(
						searchInt(itm["numberOfAlbums"]),
						searchInt(itm["albums_count"]),
						searchInt(itm["albumsCount"]),
						nestedSearchInt(itm, "albums", "total"),
					)
				}
				appendUnique(searchResult{ID: id, Title: title, Artist: artist, Album: album, Date: date, Releases: releases, TrackCount: trackCount, Explicit: explicit})
			}
		case "deezer":
			key := mediaType + "s"
			bucket, ok := page[key].(map[string]any)
			if !ok {
				continue
			}
			items, ok := bucket["items"].([]any)
			if !ok {
				continue
			}
			for _, raw := range items {
				itm, ok := raw.(map[string]any)
				if !ok {
					continue
				}
				id := asString(itm["id"])
				title := asString(itm["title"])
				if title == "" {
					title = asString(itm["name"])
				}
				artist := nestedSearchString(itm, "artist", "name")
				album := nestedSearchString(itm, "album", "title")
				trackCount := searchInt(itm["nb_tracks"])
				explicit := searchBool(itm["explicit_lyrics"])
				date := firstNonEmpty(
					asString(itm["release_date"]),
					nestedSearchString(itm, "album", "release_date"),
				)
				releases := 0
				if mediaType == "artist" {
					releases = firstPositiveInt(
						searchInt(itm["nb_album"]),
						searchInt(itm["albums_count"]),
						searchInt(itm["numberOfAlbums"]),
					)
				}
				appendUnique(searchResult{ID: id, Title: title, Artist: artist, Album: album, Date: date, Releases: releases, TrackCount: trackCount, Explicit: explicit})
			}
		case "soundcloud":
			items, ok := page["items"].([]any)
			if !ok {
				continue
			}
			for _, raw := range items {
				itm, ok := raw.(map[string]any)
				if !ok {
					continue
				}
				id := asString(itm["id"])
				title := asString(itm["title"])
				artist := nestedSearchString(itm, "artist", "name")
				trackCount := searchInt(itm["tracks_count"])
				date := firstNonEmpty(
					asString(itm["release_date"]),
					asString(itm["display_date"]),
				)
				appendUnique(searchResult{ID: id, Title: title, Artist: artist, Date: date, TrackCount: trackCount})
			}
		case "yandex":
			items, ok := page["items"].([]any)
			if !ok {
				continue
			}
			for _, raw := range items {
				itm, ok := raw.(map[string]any)
				if !ok {
					continue
				}
				id := asString(itm["id"])
				title := asString(itm["title"])
				if title == "" {
					title = asString(itm["name"])
				}
				artist := nestedSearchString(itm, "artist", "name")
				if artist == "" {
					artist = nestedSearchString(itm, "performer", "name")
				}
				album := nestedSearchString(itm, "album", "title")
				trackCount := firstPositiveInt(
					searchInt(itm["trackCount"]),
					searchInt(itm["track_count"]),
					searchInt(itm["tracks_count"]),
				)
				explicit := searchBool(itm["explicit"])
				date := firstNonEmpty(
					asString(itm["release_date"]),
					asString(itm["releaseDate"]),
					nestedSearchString(itm, "album", "release_date"),
					nestedSearchString(itm, "album", "releaseDate"),
				)
				releases := 0
				if mediaType == "artist" {
					releases = firstPositiveInt(
						searchInt(itm["albums_count"]),
						searchInt(itm["numberOfAlbums"]),
						nestedSearchInt(itm, "albums", "total"),
					)
				}
				appendUnique(searchResult{ID: id, Title: title, Artist: artist, Album: album, Date: date, Releases: releases, TrackCount: trackCount, Explicit: explicit})
			}
		}
	}
	return results
}

func formatSearchDetails(mediaType string, r searchResult) string {
	date := strings.TrimSpace(r.Date)
	if date == "" {
		date = "Unknown"
	}
	lines := []string{}
	switch mediaType {
	case "album":
		lines = append(lines, "Date released:", date)
		if r.TrackCount > 0 {
			lines = append(lines, "", fmt.Sprintf("%d Tracks", r.TrackCount))
		}
	case "artist":
		if strings.TrimSpace(r.Title) != "" {
			lines = append(lines, r.Title)
		}
		if r.Releases > 0 {
			lines = append(lines, "", fmt.Sprintf("%d Releases", r.Releases))
		}
	case "track":
		lines = append(lines, "Released on:", date)
	case "playlist":
		if r.TrackCount > 0 {
			lines = append(lines, fmt.Sprintf("%d Tracks", r.TrackCount))
		}
	default:
		if strings.TrimSpace(r.Title) != "" {
			lines = append(lines, r.Title)
		}
		if strings.TrimSpace(r.Artist) != "" {
			lines = append(lines, strings.TrimSpace(r.Artist))
		}
		if strings.TrimSpace(r.Album) != "" {
			lines = append(lines, strings.TrimSpace(r.Album))
		}
		if r.TrackCount > 0 {
			lines = append(lines, fmt.Sprintf("%d Tracks", r.TrackCount))
		}
		if r.Explicit {
			lines = append(lines, "Explicit")
		}
	}
	lines = append(lines, "", "ID: "+r.ID)
	return strings.Join(lines, "\n")
}

func formatSearchOptionLabel(index int, mediaType string, r searchResult) string {
	title := strings.TrimSpace(r.Title)
	if title == "" {
		title = "Unknown"
	}
	artist := strings.TrimSpace(r.Artist)
	switch mediaType {
	case "artist", "label":
		return fmt.Sprintf("%d. %s", index, title)
	default:
		if artist == "" {
			return fmt.Sprintf("%d. %s", index, title)
		}
		return fmt.Sprintf("%d. %s by %s", index, title, artist)
	}
}

func nestedSearchString(v map[string]any, keys ...string) string {
	cur := any(v)
	for _, key := range keys {
		m, ok := cur.(map[string]any)
		if !ok {
			return ""
		}
		cur = m[key]
	}
	return asString(cur)
}

func nestedSearchInt(v map[string]any, keys ...string) int {
	cur := any(v)
	for _, key := range keys {
		m, ok := cur.(map[string]any)
		if !ok {
			return 0
		}
		cur = m[key]
	}
	return searchInt(cur)
}

func firstPositiveInt(values ...int) int {
	for _, v := range values {
		if v > 0 {
			return v
		}
	}
	return 0
}

func searchInt(v any) int {
	switch t := v.(type) {
	case int:
		return t
	case int64:
		return int(t)
	case float64:
		return int(t)
	case string:
		i, _ := strconv.Atoi(t)
		return i
	default:
		return 0
	}
}

func searchBool(v any) bool {
	b, ok := v.(bool)
	return ok && b
}