psysonic/src-tauri/crates/psysonic-library/src/live_search.rs

//! Live Search dropdown (spec §5.9 / P24) — column-scoped FTS with LIMIT inside
//! the FTS subquery (bm25 on ≤N rowids), then a cheap join to `track`.
//! Avoids the SQLite pitfall where `JOIN track … ORDER BY bm25` on an OR
//! MATCH scans/ranks the whole hit set on 100k+ libraries (10–20s queries).

use std::collections::{HashMap, HashSet};

use rusqlite::params;

use crate::dto::{LibraryAlbumDto, LibraryArtistDto, LibraryLiveSearchResponse, LibraryTrackDto};
use crate::search::{fts_column_prefix_query, fts_query_meets_min_len, library_scope_equals_sql};
use crate::store::LibraryStore;

const SONG_FTS_COLUMNS: [&str; 4] = ["title", "artist", "album", "album_artist"];
const ALBUM_FTS_COLUMNS: [&str; 2] = ["album", "album_artist"];

struct LiveHit {
    track: LibraryTrackDto,
}

/// `library_live_search` — read connection, scoped FTS rowid picks + join.
pub fn run_live_search(
    store: &LibraryStore,
    server_id: &str,
    query: &str,
    library_scope: Option<&str>,
    artist_limit: u32,
    album_limit: u32,
    song_limit: u32,
) -> Result<LibraryLiveSearchResponse, String> {
    if !fts_query_meets_min_len(query) {
        return Ok(LibraryLiveSearchResponse {
            artists: Vec::new(),
            albums: Vec::new(),
            tracks: Vec::new(),
            source: "local".to_string(),
        });
    }

    store.with_read_conn(|conn| {
        let scope = trimmed_scope(library_scope);
        let songs = query_songs(conn, query, server_id, scope.as_deref(), song_limit)?;
        let artists = query_artists(conn, query, server_id, scope.as_deref(), artist_limit)?;
        let albums = query_albums(conn, query, server_id, scope.as_deref(), album_limit)?;
        Ok(LibraryLiveSearchResponse {
            artists,
            albums,
            tracks: songs,
            source: "local".to_string(),
        })
    })
}

/// Top FTS rowids for one or more column-scoped MATCH strings (deduped, ordered).
fn collect_fts_rowids(
    conn: &rusqlite::Connection,
    match_queries: &[String],
    per_query_limit: i64,
    total_limit: usize,
) -> rusqlite::Result<Vec<i64>> {
    let sql =
        "SELECT rowid FROM track_fts WHERE track_fts MATCH ?1 ORDER BY bm25(track_fts) LIMIT ?2";
    let mut stmt = conn.prepare(sql)?;
    let mut seen = HashSet::new();
    let mut rowids = Vec::new();
    for mq in match_queries {
        let rows = stmt.query_map(params![mq, per_query_limit], |r| r.get(0))?;
        for rowid in rows {
            let rowid = rowid?;
            if seen.insert(rowid) {
                rowids.push(rowid);
                if rowids.len() >= total_limit {
                    return Ok(rowids);
                }
            }
        }
    }
    Ok(rowids)
}

fn column_matches(query: &str, columns: &[&str]) -> Result<Vec<String>, String> {
    columns
        .iter()
        .map(|col| {
            fts_column_prefix_query(col, query).ok_or_else(|| "empty query".to_string())
        })
        .collect()
}

fn trimmed_scope(scope: Option<&str>) -> Option<String> {
    scope
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .map(str::to_string)
}

fn append_library_scope(
    sql: &mut String,
    params: &mut Vec<rusqlite::types::Value>,
    library_scope: Option<&str>,
) {
    if let Some(scope) = library_scope.filter(|s| !s.trim().is_empty()) {
        sql.push_str(" AND ");
        sql.push_str(&library_scope_equals_sql("t"));
        params.push(rusqlite::types::Value::Text(scope.to_string()));
    }
}

fn query_songs(
    conn: &rusqlite::Connection,
    query: &str,
    server_id: &str,
    library_scope: Option<&str>,
    limit: u32,
) -> rusqlite::Result<Vec<LibraryTrackDto>> {
    let matches = column_matches(query, &SONG_FTS_COLUMNS).map_err(|e| {
        rusqlite::Error::ToSqlConversionFailure(Box::new(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            e,
        )))
    })?;
    let per_col = i64::from(limit.max(4));
    let rowids = collect_fts_rowids(conn, &matches, per_col, limit as usize)?;
    if rowids.is_empty() {
        return Ok(Vec::new());
    }
    fetch_tracks_by_rowids(conn, &rowids, server_id, library_scope)
}

fn query_artists(
    conn: &rusqlite::Connection,
    query: &str,
    server_id: &str,
    library_scope: Option<&str>,
    limit: u32,
) -> rusqlite::Result<Vec<LibraryArtistDto>> {
    let Some(artist_fts) = fts_column_prefix_query("artist", query) else {
        return Ok(Vec::new());
    };
    let fetch = limit.saturating_mul(3).clamp(limit, 24);
    let rowids = collect_fts_rowids(conn, &[artist_fts], i64::from(fetch), fetch as usize)?;
    if rowids.is_empty() {
        return Ok(Vec::new());
    }
    let placeholders = rowid_placeholders(rowids.len());
    let sql = format!(
        "SELECT t.server_id, t.artist_id, t.artist, t.synced_at, t.rowid \
         FROM track t \
         WHERE t.rowid IN ({placeholders}) \
           AND t.server_id = ? \
           AND t.deleted = 0 \
           AND t.artist_id IS NOT NULL AND t.artist_id != ''"
    );
    let mut params: Vec<rusqlite::types::Value> = rowids
        .iter()
        .copied()
        .map(rusqlite::types::Value::Integer)
        .collect();
    params.push(rusqlite::types::Value::Text(server_id.to_string()));
    let mut sql = sql;
    append_library_scope(&mut sql, &mut params, library_scope);
    let rank: HashMap<i64, usize> = rowids
        .iter()
        .enumerate()
        .map(|(i, &rid)| (rid, i))
        .collect();
    let mut ranked: Vec<(usize, LibraryArtistDto)> = Vec::new();
    let mut stmt = conn.prepare(&sql)?;
    for row in stmt.query_map(rusqlite::params_from_iter(params.iter()), |r| {
        Ok((
            r.get::<_, i64>(4)?,
            r.get::<_, String>(0)?,
            r.get::<_, String>(1)?,
            r.get::<_, Option<String>>(2)?,
            r.get::<_, i64>(3)?,
        ))
    })? {
        let (rowid, server_id, artist_id, artist, synced_at) = row?;
        let fts_rank = rank.get(&rowid).copied().unwrap_or(usize::MAX);
        ranked.push((
            fts_rank,
            LibraryArtistDto {
                server_id,
                id: artist_id,
                name: artist.unwrap_or_default(),
                album_count: None,
                synced_at,
                raw_json: serde_json::Value::Null,
            },
        ));
    }
    ranked.sort_by_key(|(r, _)| *r);
    let mut out = Vec::new();
    let mut seen = HashSet::new();
    for (_, dto) in ranked {
        if !seen.insert(dto.id.clone()) {
            continue;
        }
        out.push(dto);
        if out.len() >= limit as usize {
            break;
        }
    }
    Ok(out)
}

fn query_albums(
    conn: &rusqlite::Connection,
    query: &str,
    server_id: &str,
    library_scope: Option<&str>,
    limit: u32,
) -> rusqlite::Result<Vec<LibraryAlbumDto>> {
    let matches = column_matches(query, &ALBUM_FTS_COLUMNS).map_err(|e| {
        rusqlite::Error::ToSqlConversionFailure(Box::new(std::io::Error::new(
            std::io::ErrorKind::InvalidInput,
            e,
        )))
    })?;
    let fetch = limit.saturating_mul(3).clamp(limit, 24);
    let rowids = collect_fts_rowids(conn, &matches, i64::from(fetch), fetch as usize)?;
    if rowids.is_empty() {
        return Ok(Vec::new());
    }
    let placeholders = rowid_placeholders(rowids.len());
    let sql = format!(
        "SELECT t.server_id, t.album_id, t.album, t.artist, t.artist_id, t.year, \
                t.genre, t.cover_art_id, t.starred_at, t.synced_at, t.rowid \
         FROM track t \
         WHERE t.rowid IN ({placeholders}) \
           AND t.server_id = ? \
           AND t.deleted = 0 \
           AND t.album_id IS NOT NULL AND t.album_id != ''"
    );
    let mut params: Vec<rusqlite::types::Value> = rowids
        .iter()
        .copied()
        .map(rusqlite::types::Value::Integer)
        .collect();
    params.push(rusqlite::types::Value::Text(server_id.to_string()));
    let mut sql = sql;
    append_library_scope(&mut sql, &mut params, library_scope);
    let rank: HashMap<i64, usize> = rowids
        .iter()
        .enumerate()
        .map(|(i, &rid)| (rid, i))
        .collect();
    let mut ranked: Vec<(usize, LibraryAlbumDto)> = Vec::new();
    let mut stmt = conn.prepare(&sql)?;
    for row in stmt.query_map(rusqlite::params_from_iter(params.iter()), |r| {
        Ok((
            r.get::<_, i64>(10)?,
            r.get::<_, String>(0)?,
            r.get::<_, String>(1)?,
            r.get::<_, String>(2)?,
            r.get::<_, Option<String>>(3)?,
            r.get::<_, Option<String>>(4)?,
            r.get::<_, Option<i64>>(5)?,
            r.get::<_, Option<String>>(6)?,
            r.get::<_, Option<String>>(7)?,
            r.get::<_, Option<i64>>(8)?,
            r.get::<_, i64>(9)?,
        ))
    })? {
        let (
            rowid,
            server_id,
            album_id,
            album,
            artist,
            artist_id,
            year,
            genre,
            cover_art_id,
            starred_at,
            synced_at,
        ) = row?;
        let fts_rank = rank.get(&rowid).copied().unwrap_or(usize::MAX);
        ranked.push((
            fts_rank,
            LibraryAlbumDto {
                server_id,
                id: album_id,
                name: album,
                artist,
                artist_id,
                song_count: None,
                duration_sec: None,
                year,
                genre,
                cover_art_id,
                starred_at,
                synced_at,
                raw_json: serde_json::Value::Null,
            },
        ));
    }
    ranked.sort_by_key(|(r, _)| *r);
    let mut out = Vec::new();
    let mut seen = HashSet::new();
    for (_, dto) in ranked {
        if !seen.insert(dto.id.clone()) {
            continue;
        }
        out.push(dto);
        if out.len() >= limit as usize {
            break;
        }
    }
    Ok(out)
}

fn rowid_placeholders(n: usize) -> String {
    (0..n).map(|_| "?").collect::<Vec<_>>().join(", ")
}

fn fetch_tracks_by_rowids(
    conn: &rusqlite::Connection,
    rowids: &[i64],
    server_id: &str,
    library_scope: Option<&str>,
) -> rusqlite::Result<Vec<LibraryTrackDto>> {
    let placeholders = rowid_placeholders(rowids.len());
    let sql = format!(
        "SELECT \
          t.rowid, \
          t.server_id, t.id, t.title, t.artist, t.artist_id, t.album, t.album_id, \
          t.album_artist, t.duration_sec, t.track_number, t.disc_number, t.year, \
          t.genre, t.suffix, t.bit_rate, t.size_bytes, t.cover_art_id, \
          t.starred_at, t.user_rating, t.play_count, t.bpm, t.synced_at \
         FROM track t \
         WHERE t.rowid IN ({placeholders}) \
           AND t.server_id = ? \
           AND t.deleted = 0"
    );
    let mut params: Vec<rusqlite::types::Value> = rowids
        .iter()
        .copied()
        .map(rusqlite::types::Value::Integer)
        .collect();
    params.push(rusqlite::types::Value::Text(server_id.to_string()));
    let mut sql = sql;
    append_library_scope(&mut sql, &mut params, library_scope);
    let mut stmt = conn.prepare(&sql)?;
    let mut by_rowid: HashMap<i64, LibraryTrackDto> = HashMap::new();
    for row in stmt.query_map(rusqlite::params_from_iter(params.iter()), |r| {
        let rowid: i64 = r.get(0)?;
        let hit = map_live_hit_row(r, 1)?;
        Ok((rowid, hit.track))
    })? {
        let (rowid, track) = row?;
        by_rowid.insert(rowid, track);
    }
    Ok(rowids
        .iter()
        .filter_map(|rid| by_rowid.get(rid).cloned())
        .collect())
}

fn map_live_hit_row(row: &rusqlite::Row<'_>, offset: usize) -> rusqlite::Result<LiveHit> {
    Ok(LiveHit {
        track: LibraryTrackDto {
            server_id: row.get(offset)?,
            id: row.get(offset + 1)?,
            content_hash: None,
            title: row.get(offset + 2)?,
            title_sort: None,
            artist: row.get(offset + 3)?,
            artist_id: row.get(offset + 4)?,
            album: row.get(offset + 5)?,
            album_id: row.get(offset + 6)?,
            album_artist: row.get(offset + 7)?,
            duration_sec: row.get(offset + 8)?,
            track_number: row.get(offset + 9)?,
            disc_number: row.get(offset + 10)?,
            year: row.get(offset + 11)?,
            genre: row.get(offset + 12)?,
            suffix: row.get(offset + 13)?,
            bit_rate: row.get(offset + 14)?,
            size_bytes: row.get(offset + 15)?,
            cover_art_id: row.get(offset + 16)?,
            starred_at: row.get(offset + 17)?,
            user_rating: row.get(offset + 18)?,
            play_count: row.get(offset + 19)?,
            bpm: row.get(offset + 20)?,
            bpm_source: None,
            played_at: None,
            server_path: None,
            library_id: None,
            isrc: None,
            mbid_recording: None,
            replay_gain_track_db: None,
            replay_gain_album_db: None,
            server_updated_at: None,
            server_created_at: None,
            synced_at: row.get(offset + 21)?,
            enrichment: None,
            raw_json: serde_json::Value::Null,
        },
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::repos::{TrackRepository, TrackRow};
    use serde_json::json;

    fn track(
        server: &str,
        id: &str,
        title: &str,
        artist: &str,
        album: &str,
        album_id: &str,
        artist_id: &str,
    ) -> TrackRow {
        TrackRow {
            server_id: server.into(),
            id: id.into(),
            title: title.into(),
            title_sort: None,
            artist: Some(artist.into()),
            artist_id: Some(artist_id.into()),
            album: album.into(),
            album_id: Some(album_id.into()),
            album_artist: Some(artist.into()),
            duration_sec: 200,
            track_number: Some(1),
            disc_number: Some(1),
            year: None,
            genre: None,
            suffix: None,
            bit_rate: None,
            size_bytes: None,
            cover_art_id: Some(format!("cv_{album_id}")),
            starred_at: None,
            user_rating: None,
            play_count: None,
            played_at: None,
            server_path: None,
            library_id: None,
            isrc: None,
            mbid_recording: None,
            bpm: None,
            replay_gain_track_db: None,
            replay_gain_album_db: None,
            content_hash: None,
            server_updated_at: None,
            server_created_at: None,
            deleted: false,
            synced_at: 1,
            raw_json: "{}".into(),
        }
    }

    #[test]
    fn live_search_prefix_matches_partial_artist_name() {
        let store = LibraryStore::open_in_memory();
        TrackRepository::new(&store)
            .upsert_batch(&[
                track(
                    "s1",
                    "t1",
                    "Enter Sandman",
                    "Metallica",
                    "Metallica",
                    "al1",
                    "ar_meta",
                ),
                track("s1", "t2", "Other", "Other Artist", "Other Album", "al2", "ar2"),
            ])
            .unwrap();
        let resp = run_live_search(&store, "s1", "metal", None, 5, 5, 10).unwrap();
        assert!(
            resp.artists.iter().any(|a| a.name == "Metallica"),
            "expected Metallica from prefix query metal"
        );
        assert!(resp.tracks.iter().any(|t| t.artist.as_deref() == Some("Metallica")));
    }

    #[test]
    fn live_search_returns_songs_albums_artists_from_scoped_fts() {
        let store = LibraryStore::open_in_memory();
        TrackRepository::new(&store)
            .upsert_batch(&[
                track("s1", "t1", "Aurora Song", "Aurora Quartet", "Aurora Nights", "al1", "ar1"),
                track("s1", "t2", "Other", "Other Artist", "Other Album", "al2", "ar2"),
            ])
            .unwrap();
        let resp = run_live_search(&store, "s1", "aurora", None, 5, 5, 10).unwrap();
        assert_eq!(resp.tracks.len(), 1);
        assert_eq!(resp.albums.len(), 1);
        assert_eq!(resp.albums[0].id, "al1");
        assert_eq!(resp.artists.len(), 1);
        assert_eq!(resp.artists[0].id, "ar1");
        assert!(resp.tracks[0].raw_json.is_null());
    }

    #[test]
    fn live_search_does_not_surface_artist_from_unrelated_track_hit() {
        let store = LibraryStore::open_in_memory();
        TrackRepository::new(&store)
            .upsert_batch(&[
                track(
                    "s1",
                    "t1",
                    "Battle Hymn",
                    "Arch Enemy",
                    "Manowar Covers Vol 1",
                    "al1",
                    "ar_arch",
                ),
                track(
                    "s1",
                    "t2",
                    "Heart Of Steel",
                    "Manowar",
                    "Fighting the World",
                    "al2",
                    "ar_mano",
                ),
            ])
            .unwrap();
        let resp = run_live_search(&store, "s1", "manowar", None, 5, 5, 10).unwrap();
        assert!(
            resp.artists.iter().any(|a| a.name == "Manowar"),
            "expected Manowar artist"
        );
        assert!(
            !resp.artists.iter().any(|a| a.name == "Arch Enemy"),
            "Arch Enemy must not appear when only the album title mentions Manowar"
        );
        assert!(resp.albums.iter().any(|a| a.name.contains("Manowar")));
    }

    #[test]
    fn live_search_short_query_returns_empty_without_scanning() {
        let store = LibraryStore::open_in_memory();
        TrackRepository::new(&store)
            .upsert_batch(&[track(
                "s1",
                "t1",
                "Аура",
                "Artist",
                "Album",
                "al1",
                "ar1",
            )])
            .unwrap();
        let resp = run_live_search(&store, "s1", "а", None, 5, 5, 10).unwrap();
        assert!(resp.tracks.is_empty());
        assert!(resp.artists.is_empty());
        assert!(resp.albums.is_empty());
    }

    #[test]
    fn live_search_library_scope_narrows_results() {
        let store = LibraryStore::open_in_memory();
        let mut in_lib = track(
            "s1",
            "t1",
            "Scoped Song",
            "Scoped Artist",
            "Scoped Album",
            "al1",
            "ar1",
        );
        in_lib.library_id = Some("lib1".into());
        let mut other = track(
            "s1",
            "t2",
            "Scoped Song",
            "Other Artist",
            "Other Album",
            "al2",
            "ar2",
        );
        other.library_id = Some("lib2".into());
        TrackRepository::new(&store)
            .upsert_batch(&[in_lib, other])
            .unwrap();
        let resp = run_live_search(&store, "s1", "scoped", Some("lib1"), 5, 5, 10).unwrap();
        assert_eq!(resp.tracks.len(), 1);
        assert_eq!(resp.tracks[0].id, "t1");
        assert_eq!(resp.artists.len(), 1);
        assert_eq!(resp.artists[0].name, "Scoped Artist");
        assert_eq!(resp.albums.len(), 1);
        assert_eq!(resp.albums[0].name, "Scoped Album");
    }

    #[test]
    fn live_search_library_scope_matches_raw_json_when_column_null() {
        let store = LibraryStore::open_in_memory();
        let mut row = track(
            "s1",
            "t1",
            "Scoped Song",
            "Scoped Artist",
            "Scoped Album",
            "al1",
            "ar1",
        );
        row.raw_json = json!({"libraryId": 3}).to_string();
        TrackRepository::new(&store).upsert_batch(&[row]).unwrap();
        let resp = run_live_search(&store, "s1", "scoped", Some("3"), 5, 5, 10).unwrap();
        assert_eq!(resp.tracks.len(), 1);
        assert_eq!(resp.tracks[0].id, "t1");
    }

    /// Manual: `cargo test -p psysonic-library bench_disk_live_search --release -- --ignored --nocapture`
    #[test]
    #[ignore]
    fn bench_disk_live_search() {
        use std::path::PathBuf;
        use std::time::Instant;

        let path: PathBuf = std::env::var("HOME")
            .map(|h| PathBuf::from(h).join(".local/share/dev.psysonic.player/library.sqlite"))
            .expect("HOME");
        if !path.exists() {
            eprintln!("skip: no db at {}", path.display());
            return;
        }
        let conn = rusqlite::Connection::open_with_flags(
            &path,
            rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY,
        )
        .expect("open db");
        conn.pragma_update(None, "cache_size", -64000).unwrap();

        let server_id = std::env::var("PSYSONIC_BENCH_SERVER_ID").unwrap_or_else(|_| {
            conn.query_row(
                "SELECT server_id FROM track WHERE deleted = 0 LIMIT 1",
                [],
                |r| r.get::<_, String>(0),
            )
            .expect("server_id")
        });

        for q in ["manowar", "metallica", "arch enemy", "metal", "meta"] {
            let t0 = Instant::now();
            let songs = query_songs(&conn, q, &server_id, None, 10).unwrap();
            let t1 = Instant::now();
            let artists = query_artists(&conn, q, &server_id, None, 5).unwrap();
            let t2 = Instant::now();
            let albums = query_albums(&conn, q, &server_id, None, 5).unwrap();
            let t3 = Instant::now();
            eprintln!(
                "{q:?}: songs={} ({:.1}ms) artists={} ({:.1}ms) albums={} ({:.1}ms) total={:.1}ms",
                songs.len(),
                t1.duration_since(t0).as_secs_f64() * 1000.0,
                artists.len(),
                t2.duration_since(t1).as_secs_f64() * 1000.0,
                albums.len(),
                t3.duration_since(t2).as_secs_f64() * 1000.0,
                t3.duration_since(t0).as_secs_f64() * 1000.0,
            );
        }
    }
}