mcma-backend/app/infrastructure/db/repositories/track_repository.py

"""Track repository — adapter over ``AsyncSession``."""

import datetime as dt
import uuid

from sqlalchemy import case, func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.domain.entities.storage import FormatBreakdown, LibraryStats
from app.domain.entities.track import Track
from app.domain.errors import NotFoundError
from app.infrastructure.db.models.artist import ArtistModel
from app.infrastructure.db.models.enums import TrackAvailability
from app.infrastructure.db.models.track import TrackModel


def _to_entity(row: TrackModel) -> Track:
    return Track(
        id=row.id,
        title=row.title,
        artist_id=row.artist_id,
        album_id=row.album_id,
        storage_uri=row.storage_uri,
        file_format=row.file_format,
        file_size=row.file_size,
        source=row.source,
        source_id=row.source_id,
        duration_seconds=row.duration_seconds,
        genre=row.genre,
        year=row.year,
        track_number=row.track_number,
        metadata_status=row.metadata_status,
        metadata_error=row.metadata_error,
        enriched_at=row.enriched_at,
        availability=row.availability,
        created_at=row.created_at,
        updated_at=row.updated_at,
    )


class SqlAlchemyTrackRepository:
    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def get_by_id(self, track_id: uuid.UUID) -> Track | None:
        row = await self._session.get(TrackModel, track_id)
        return _to_entity(row) if row is not None else None

    async def get_many(self, ids: list[uuid.UUID]) -> list[Track]:
        if not ids:
            return []
        rows = (
            (await self._session.execute(select(TrackModel).where(TrackModel.id.in_(ids))))
            .scalars()
            .all()
        )
        return [_to_entity(r) for r in rows]

    async def get_by_source(self, source: str, source_id: str) -> Track | None:
        row = (
            await self._session.execute(
                select(TrackModel).where(
                    TrackModel.source == source,
                    TrackModel.source_id == source_id,
                )
            )
        ).scalar_one_or_none()
        return _to_entity(row) if row is not None else None

    async def add(
        self,
        *,
        id: uuid.UUID,
        title: str,
        artist_id: uuid.UUID,
        storage_uri: str | None,
        file_format: str | None,
        file_size: int | None,
        source: str,
        source_id: str,
        metadata_status: str,
        added_by: uuid.UUID | None,
        availability: str = TrackAvailability.LOCAL.value,
    ) -> Track:
        row = TrackModel(
            id=id,
            title=title,
            artist_id=artist_id,
            storage_uri=storage_uri,
            file_format=file_format,
            file_size=file_size,
            source=source,
            source_id=source_id,
            metadata_status=metadata_status,
            added_by=added_by,
            availability=availability,
        )
        self._session.add(row)
        await self._session.flush()
        await self._session.refresh(row)
        return _to_entity(row)

    async def materialize(
        self,
        track_id: uuid.UUID,
        *,
        storage_uri: str,
        file_format: str,
        file_size: int,
        bitrate: int | None,
    ) -> Track:
        """Fill in a remote placeholder's audio fields after a download (lazy
        materialization). ``track.id`` is unchanged, so likes/playlists/queue
        entries that already reference it keep working."""
        row = await self._session.get(TrackModel, track_id)
        if row is None:
            raise NotFoundError(f"Track {track_id} not found.")
        row.storage_uri = storage_uri
        row.file_format = file_format
        row.file_size = file_size
        if bitrate is not None:
            row.bitrate = bitrate
        row.availability = TrackAvailability.LOCAL.value
        await self._session.flush()
        await self._session.refresh(row)
        return _to_entity(row)

    async def delete(self, track_id: uuid.UUID) -> None:
        row = await self._session.get(TrackModel, track_id)
        if row is not None:
            await self._session.delete(row)
            await self._session.flush()

    async def genres(self) -> list[tuple[str, int]]:
        """Distinct non-null genres with their song counts, most common first.

        Defined before ``list`` — the method named ``list`` shadows the builtin
        in later annotations within the class body."""
        rows = (
            await self._session.execute(
                select(TrackModel.genre, func.count(TrackModel.id).label("cnt"))
                .where(TrackModel.genre.is_not(None))
                .group_by(TrackModel.genre)
                .order_by(func.count(TrackModel.id).desc())
            )
        ).all()
        return [(row.genre, row.cnt) for row in rows]

    async def list_similar(
        self,
        *,
        genre: str | None,
        artist_id: uuid.UUID,
        exclude_ids: list[uuid.UUID],
        limit: int,
    ) -> list[Track]:
        # Rank a same-genre hit above a same-artist hit; shuffle within a tier so
        # the mix varies. Only playable (locally-stored) tracks are candidates.
        if genre is not None:
            match = or_(TrackModel.genre == genre, TrackModel.artist_id == artist_id)
            score = case((TrackModel.genre == genre, 2), else_=0) + case(
                (TrackModel.artist_id == artist_id, 1), else_=0
            )
        else:
            match = TrackModel.artist_id == artist_id
            score = case((TrackModel.artist_id == artist_id, 1), else_=0)

        stmt = select(TrackModel).where(TrackModel.storage_uri.is_not(None), match)
        if exclude_ids:
            stmt = stmt.where(TrackModel.id.not_in(exclude_ids))
        stmt = stmt.order_by(score.desc(), func.random()).limit(limit)
        rows = (await self._session.execute(stmt)).scalars().all()
        return [_to_entity(r) for r in rows]

    async def sample_playable(
        self, *, exclude_ids: list[uuid.UUID], limit: int
    ) -> list[Track]:
        stmt = select(TrackModel).where(TrackModel.storage_uri.is_not(None))
        if exclude_ids:
            stmt = stmt.where(TrackModel.id.not_in(exclude_ids))
        stmt = stmt.order_by(func.random()).limit(limit)
        rows = (await self._session.execute(stmt)).scalars().all()
        return [_to_entity(r) for r in rows]

    async def library_stats(self) -> LibraryStats:
        """One-shot aggregate over the whole catalogue (no pagination). Defined
        before ``list`` for the same shadowing reason as ``genres``."""
        totals = (
            await self._session.execute(
                select(
                    func.count(TrackModel.id),
                    func.coalesce(func.sum(TrackModel.file_size), 0),
                    func.coalesce(func.sum(TrackModel.duration_seconds), 0),
                    func.coalesce(func.max(TrackModel.file_size), 0),
                    func.min(TrackModel.created_at),
                    func.max(TrackModel.created_at),
                )
            )
        ).one()

        fmt_rows = (
            await self._session.execute(
                select(
                    TrackModel.file_format,
                    func.count(TrackModel.id),
                    func.coalesce(func.sum(TrackModel.file_size), 0),
                )
                .where(TrackModel.file_format.is_not(None))
                .group_by(TrackModel.file_format)
                .order_by(func.sum(TrackModel.file_size).desc())
            )
        ).all()

        status_rows = (
            await self._session.execute(
                select(TrackModel.metadata_status, func.count(TrackModel.id)).group_by(
                    TrackModel.metadata_status
                )
            )
        ).all()

        source_rows = (
            await self._session.execute(
                select(TrackModel.source, func.count(TrackModel.id)).group_by(TrackModel.source)
            )
        ).all()

        return LibraryStats(
            total_tracks=totals[0],
            total_size=totals[1],
            total_duration_seconds=totals[2],
            largest_track_size=totals[3],
            earliest_added=totals[4],
            latest_added=totals[5],
            by_format=[
                FormatBreakdown(file_format=fmt, track_count=cnt, total_size=size)
                for fmt, cnt, size in fmt_rows
            ],
            by_metadata_status={status: cnt for status, cnt in status_rows},
            by_source={source: cnt for source, cnt in source_rows},
        )

    async def find_duplicate_groups(self) -> list[tuple[str, list[Track]]]:
        """Tracks that share an ``acoustid_fingerprint`` (the dedup key), grouped
        by it — only fingerprints with more than one track. Empty when clean.
        Defined before ``list`` for the same name-shadowing reason as ``genres``."""
        dup_fps = (
            select(TrackModel.acoustid_fingerprint)
            .where(TrackModel.acoustid_fingerprint.is_not(None))
            .group_by(TrackModel.acoustid_fingerprint)
            .having(func.count(TrackModel.id) > 1)
            .scalar_subquery()
        )
        rows = (
            (
                await self._session.execute(
                    select(TrackModel)
                    .where(TrackModel.acoustid_fingerprint.in_(dup_fps))
                    .order_by(TrackModel.acoustid_fingerprint, TrackModel.created_at)
                )
            )
            .scalars()
            .all()
        )
        groups: dict[str, list[Track]] = {}
        for row in rows:
            fingerprint = row.acoustid_fingerprint
            assert fingerprint is not None  # filtered to non-null above
            groups.setdefault(fingerprint, []).append(_to_entity(row))
        return list(groups.items())

    async def list_by_metadata_status(
        self, status: str, *, limit: int, offset: int
    ) -> list[Track]:
        """Tracks in a given ``metadata_status`` (e.g. ``pending``/``failed``),
        newest first. Defined before ``list`` (name-shadowing)."""
        rows = (
            (
                await self._session.execute(
                    select(TrackModel)
                    .where(TrackModel.metadata_status == status)
                    .order_by(TrackModel.created_at.desc())
                    .limit(limit)
                    .offset(offset)
                )
            )
            .scalars()
            .all()
        )
        return [_to_entity(r) for r in rows]

    async def all_storage_refs(self) -> list[tuple[uuid.UUID, str]]:
        """``(id, storage_uri)`` for every *local* track — for the cleanup
        worker's filesystem reconciliation. Remote placeholders have no local
        file (``availability != local``) and are skipped. No entity hydration."""
        rows = (
            await self._session.execute(
                select(TrackModel.id, TrackModel.storage_uri).where(
                    TrackModel.availability == TrackAvailability.LOCAL.value,
                    TrackModel.storage_uri.is_not(None),
                )
            )
        ).all()
        return [(row.id, row.storage_uri) for row in rows]

    async def count_by_metadata_status(self, status: str) -> int:
        return (
            await self._session.execute(
                select(func.count())
                .select_from(TrackModel)
                .where(TrackModel.metadata_status == status)
            )
        ).scalar_one()

    async def list_changed_since(
        self, *, since: dt.datetime | None, until: dt.datetime
    ) -> list[Track]:
        """Catalogue tracks changed in the window ``(since, until]`` (by
        ``updated_at``), oldest first — the delta a client caches for offline use.
        Defined before ``list`` (name-shadowing)."""
        stmt = select(TrackModel).where(TrackModel.updated_at <= until)
        if since is not None:
            stmt = stmt.where(TrackModel.updated_at > since)
        stmt = stmt.order_by(TrackModel.updated_at)
        rows = (await self._session.execute(stmt)).scalars().all()
        return [_to_entity(r) for r in rows]

    async def list(
        self,
        *,
        artist_id: uuid.UUID | None,
        album_id: uuid.UUID | None,
        q: str | None,
        source: str | None = None,
        sort_by: str = "created_at",
        order: str = "desc",
        limit: int = 50,
        offset: int = 0,
    ) -> list[Track]:
        stmt = select(TrackModel)
        if artist_id is not None:
            stmt = stmt.where(TrackModel.artist_id == artist_id)
        if album_id is not None:
            stmt = stmt.where(TrackModel.album_id == album_id)
        if source is not None:
            stmt = stmt.where(TrackModel.source == source)
        if q:
            stmt = stmt.where(TrackModel.title.ilike(f"%{q}%"))

        if sort_by == "artist":
            stmt = stmt.join(ArtistModel, TrackModel.artist_id == ArtistModel.id)
            col_artist = ArtistModel.name
            stmt = stmt.order_by(col_artist.asc() if order == "asc" else col_artist.desc())
        elif sort_by == "title":
            col_title = TrackModel.title
            stmt = stmt.order_by(col_title.asc() if order == "asc" else col_title.desc())
        else:
            stmt = stmt.order_by(
                TrackModel.created_at.asc() if order == "asc" else TrackModel.created_at.desc()
            )
        stmt = stmt.limit(limit).offset(offset)
        rows = (await self._session.execute(stmt)).scalars().all()
        return [_to_entity(r) for r in rows]

    async def count(
        self,
        *,
        artist_id: uuid.UUID | None,
        album_id: uuid.UUID | None,
        q: str | None,
        source: str | None = None,
    ) -> int:
        stmt = select(func.count()).select_from(TrackModel)
        if artist_id is not None:
            stmt = stmt.where(TrackModel.artist_id == artist_id)
        if album_id is not None:
            stmt = stmt.where(TrackModel.album_id == album_id)
        if source is not None:
            stmt = stmt.where(TrackModel.source == source)
        if q:
            stmt = stmt.where(TrackModel.title.ilike(f"%{q}%"))
        return (await self._session.execute(stmt)).scalar_one()

    async def update(
        self,
        track_id: uuid.UUID,
        *,
        title: str | None,
        genre: str | None,
        year: int | None,
        artist_id: uuid.UUID | None = None,
        album_id: uuid.UUID | None = None,
        track_number: int | None = None,
    ) -> Track:
        row = await self._session.get(TrackModel, track_id)
        if row is None:
            raise NotFoundError(f"Track {track_id} not found.")
        if title is not None:
            row.title = title
        if genre is not None:
            row.genre = genre
        if year is not None:
            row.year = year
        if artist_id is not None:
            row.artist_id = artist_id
        if album_id is not None:
            row.album_id = album_id
        if track_number is not None:
            row.track_number = track_number
        row.metadata_status = "manual"
        await self._session.flush()
        await self._session.refresh(row)
        return _to_entity(row)

    async def apply_enrichment(
        self,
        track_id: uuid.UUID,
        *,
        title: str,
        artist_id: uuid.UUID,
        album_id: uuid.UUID | None,
        genre: str | None,
        year: int | None,
        track_number: int | None,
        duration_seconds: int | None,
        bitrate: int | None,
        acoustid_fingerprint: str | None,
        musicbrainz_id: str | None,
        metadata_status: str,
        metadata_error: str | None = None,
    ) -> Track:
        row = await self._session.get(TrackModel, track_id)
        if row is None:
            raise NotFoundError(f"Track {track_id} not found.")
        # Identity + status are authoritative for an enrichment run.
        row.title = title
        row.artist_id = artist_id
        row.metadata_status = metadata_status
        # A finished run always stamps outcome: clear/set the reason and mark the
        # completion time so the UI can tell "still pending" from "done/failed".
        row.metadata_error = metadata_error
        row.enriched_at = dt.datetime.now(dt.UTC)
        # Nullable extras: fill gaps only — never erase data a prior run found.
        if album_id is not None:
            row.album_id = album_id
        if genre is not None:
            row.genre = genre
        if year is not None:
            row.year = year
        if track_number is not None:
            row.track_number = track_number
        if duration_seconds is not None:
            row.duration_seconds = duration_seconds
        if bitrate is not None:
            row.bitrate = bitrate
        if acoustid_fingerprint is not None:
            row.acoustid_fingerprint = acoustid_fingerprint
        if musicbrainz_id is not None:
            row.musicbrainz_id = musicbrainz_id
        await self._session.flush()
        await self._session.refresh(row)
        return _to_entity(row)

    async def mark_enrichment_failed(self, track_id: uuid.UUID, *, error: str) -> None:
        """Record that an enrichment run crashed (unexpected exception). Runs in
        its own session so the failure is persisted even though the run's own
        transaction rolled back. Never overwrites ``manual`` (a no-op then), and
        a missing track is a clean no-op."""
        row = await self._session.get(TrackModel, track_id)
        if row is None or row.metadata_status == "manual":
            return
        row.metadata_status = "failed"
        row.metadata_error = error
        row.enriched_at = dt.datetime.now(dt.UTC)
        await self._session.flush()