qrate/bot.py

"""
Signal-Qobuz bot — main entry point.

Usage:
    python bot.py
"""

import asyncio
import logging

import config
import qobuz_client
from qobuz_client import AlbumItem
from session import PendingArtist, ResultEntry, State, UserSession
from signal_client import SignalClient

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

_sessions: dict[str, UserSession] = {}


def _get_session(number: str) -> UserSession:
    if number not in _sessions:
        _sessions[number] = UserSession()
    return _sessions[number]


def _results_message(entries: list[ResultEntry]) -> str:
    lines = ["Here's what I found:\n"]
    for i, e in enumerate(entries, 1):
        lines.append(f"[{i}] {e.display}")
    lines.append("\nReply with a number to download, or send a new search.")
    return "\n".join(lines)


def _artist_options_message(name: str, total: int) -> str:
    return (
        f"{name} — {total} album(s) on Qobuz.\n\n"
        f"[1] Full discography ({total} albums)\n"
        f"[2] Top {config.MAX_ARTIST_ALBUMS} most recent albums\n\n"
        "Reply with 1 or 2, or send a new search."
    )


async def handle_message(signal: SignalClient, sender: str, text: str) -> None:
    session = _get_session(sender)
    text = text.strip()

    if session.state == State.DOWNLOADING:
        await signal.send_message(sender, "A download is already in progress. Please wait.")
        return

    # ── ARTIST OPTIONS: waiting for "1" or "2" ──────────────────────────────
    if session.state == State.ARTIST_OPTIONS and session.pending_artist:
        if text == "1":
            pending = session.pending_artist
            session.state = State.DOWNLOADING
            await signal.send_message(
                sender, f"Downloading full discography: {pending.name}…"
            )
            await _download_album_list(signal, sender, session, pending.albums)
            return
        if text == "2":
            pending = session.pending_artist
            albums = pending.albums[:config.MAX_ARTIST_ALBUMS]
            session.state = State.DOWNLOADING
            await signal.send_message(
                sender,
                f"Downloading {len(albums)} most recent album(s) by {pending.name}…",
            )
            await _download_album_list(signal, sender, session, albums)
            return
        # Anything else → treat as a new search

    # ── RESULTS: waiting for a number ────────────────────────────────────────
    if session.state == State.RESULTS and text.isdigit():
        index = int(text) - 1
        if not (0 <= index < len(session.entries)):
            await signal.send_message(
                sender,
                f"Please pick a number between 1 and {len(session.entries)}.",
            )
            return

        entry = session.entries[index]

        if entry.kind == "album":
            session.state = State.DOWNLOADING
            await signal.send_message(sender, f"Downloading: {entry.display}…")
            try:
                await qobuz_client.download_url(entry.item.url)
                await signal.send_message(sender, "Download complete.")
            except Exception as exc:
                logger.error("Album download error: %s", exc)
                await signal.send_message(sender, f"Download failed: {exc}")
            session.reset()
            return

        if entry.kind == "artist":
            await signal.send_message(
                sender, f"Fetching discography for {entry.display}…"
            )
            try:
                albums = await qobuz_client.get_artist_albums(entry.item.url)
            except Exception as exc:
                logger.error("Artist albums error: %s", exc)
                await signal.send_message(sender, f"Error fetching albums: {exc}")
                session.reset()
                return

            if not albums:
                await signal.send_message(sender, "No albums found for this artist.")
                session.reset()
                return

            # Extract artist name from the display text (format: "Name - (N releases)")
            artist_name = entry.display.split(" - (")[0].strip()
            session.state = State.ARTIST_OPTIONS
            session.pending_artist = PendingArtist(name=artist_name, albums=albums)
            await signal.send_message(
                sender, _artist_options_message(artist_name, len(albums))
            )
            return

    # ── NEW SEARCH (IDLE or fallthrough) ─────────────────────────────────────
    session.reset()
    await signal.send_message(sender, f'Searching Qobuz for "{text}"…')

    try:
        results = await qobuz_client.search(text)
    except Exception as exc:
        logger.error("Search error: %s", exc)
        await signal.send_message(sender, f"Search failed: {exc}")
        return

    entries: list[ResultEntry] = []

    for item in results.artists:
        entries.append(ResultEntry(display=item.text, kind="artist", item=item))

    for item in results.albums:
        entries.append(ResultEntry(display=item.text, kind="album", item=item))

    if not entries:
        await signal.send_message(sender, "No results found. Try a different query.")
        return

    session.entries = entries
    session.state = State.RESULTS
    await signal.send_message(sender, _results_message(entries))


async def _download_album_list(
    signal: SignalClient,
    sender: str,
    session: UserSession,
    albums: list[AlbumItem],
) -> None:
    errors = []
    for album in albums:
        try:
            await qobuz_client.download_album_id(album.id)
            await signal.send_message(sender, f"✓ {album.title}")
        except Exception as exc:
            logger.error("Download error for album %s: %s", album.id, exc)
            errors.append(album.title)
            await signal.send_message(sender, f"✗ {album.title} (failed)")

    ok = len(albums) - len(errors)
    if errors:
        await signal.send_message(sender, f"Done — {ok}/{len(albums)} succeeded.")
    else:
        await signal.send_message(sender, "All downloads complete.")

    session.reset()


async def main() -> None:
    signal = SignalClient()
    config.DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)

    logger.info(
        "Bot starting — signal-cli at %s:%s, downloads → %s",
        config.SIGNAL_HOST,
        config.SIGNAL_PORT,
        config.DOWNLOAD_DIR,
    )

    async def dispatch(sender: str, text: str) -> None:
        try:
            await handle_message(signal, sender, text)
        except Exception as exc:
            logger.exception("Unhandled error for %s: %s", sender, exc)
            try:
                await signal.send_message(sender, "An unexpected error occurred. Please try again.")
            except Exception:
                pass

    await signal.listen(dispatch)


if __name__ == "__main__":
    asyncio.run(main())