haunt-fm/src/haunt_fm/models/track.py

from datetime import datetime

from pgvector.sqlalchemy import Vector
from sqlalchemy import REAL, BigInteger, DateTime, ForeignKey, Index, Integer, Text, func
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship

from haunt_fm.models.base import Base


class Track(Base):
    __tablename__ = "tracks"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    title: Mapped[str] = mapped_column(Text, nullable=False)
    artist: Mapped[str] = mapped_column(Text, nullable=False)
    album: Mapped[str | None] = mapped_column(Text)
    fingerprint: Mapped[str] = mapped_column(Text, unique=True, nullable=False)
    lastfm_url: Mapped[str | None] = mapped_column(Text)
    itunes_track_id: Mapped[int | None] = mapped_column(BigInteger)
    itunes_preview_url: Mapped[str | None] = mapped_column(Text)
    apple_music_id: Mapped[str | None] = mapped_column(Text)
    duration_ms: Mapped[int | None] = mapped_column(Integer)
    genre: Mapped[str | None] = mapped_column(Text)
    embedding_status: Mapped[str] = mapped_column(Text, nullable=False, default="pending")
    embedding_error: Mapped[str | None] = mapped_column(Text)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), server_default=func.now(), onupdate=func.now()
    )

    listen_events: Mapped[list["ListenEvent"]] = relationship(back_populates="track")
    embedding: Mapped["TrackEmbedding | None"] = relationship(back_populates="track")


class Profile(Base):
    __tablename__ = "profiles"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    name: Mapped[str] = mapped_column(Text, unique=True, nullable=False)
    display_name: Mapped[str | None] = mapped_column(Text)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())


class SpeakerProfileMapping(Base):
    __tablename__ = "speaker_profile_mappings"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    speaker_name: Mapped[str] = mapped_column(Text, unique=True, nullable=False)
    profile_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("profiles.id"), nullable=False)


class ListenEvent(Base):
    __tablename__ = "listen_events"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    track_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("tracks.id"), nullable=False)
    profile_id: Mapped[int | None] = mapped_column(BigInteger, ForeignKey("profiles.id"))
    source: Mapped[str] = mapped_column(Text, nullable=False, default="music_assistant")
    speaker_name: Mapped[str | None] = mapped_column(Text)
    listened_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
    duration_played: Mapped[int | None] = mapped_column(Integer)
    raw_payload: Mapped[dict | None] = mapped_column(JSONB)

    track: Mapped[Track] = relationship(back_populates="listen_events")


class TrackEmbedding(Base):
    __tablename__ = "track_embeddings"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    track_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("tracks.id"), unique=True, nullable=False)
    embedding = mapped_column(Vector(512), nullable=False)
    model_version: Mapped[str] = mapped_column(Text, nullable=False, default="laion/larger_clap_music")
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    __table_args__ = (
        Index("ix_track_embeddings_embedding_hnsw", "embedding", postgresql_using="hnsw", postgresql_with={"m": 16, "ef_construction": 64}, postgresql_ops={"embedding": "vector_cosine_ops"}),
    )

    track: Mapped[Track] = relationship(back_populates="embedding")


class SimilarityLink(Base):
    __tablename__ = "similarity_links"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    source_track_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("tracks.id"), nullable=False)
    target_track_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("tracks.id"), nullable=False)
    lastfm_match: Mapped[float | None] = mapped_column(REAL)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    __table_args__ = (
        Index("uq_similarity_link", "source_track_id", "target_track_id", unique=True),
    )


class TasteProfile(Base):
    __tablename__ = "taste_profiles"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    name: Mapped[str] = mapped_column(Text, unique=True, nullable=False, default="default")
    profile_id: Mapped[int | None] = mapped_column(BigInteger, ForeignKey("profiles.id"), unique=True)
    embedding = mapped_column(Vector(512), nullable=False)
    track_count: Mapped[int] = mapped_column(Integer, nullable=False)
    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())


class Playlist(Base):
    __tablename__ = "playlists"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    name: Mapped[str | None] = mapped_column(Text)
    known_pct: Mapped[int] = mapped_column(Integer, nullable=False)
    total_tracks: Mapped[int] = mapped_column(Integer, nullable=False)
    vibe: Mapped[str | None] = mapped_column(Text)
    alpha: Mapped[float | None] = mapped_column(REAL)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())

    tracks: Mapped[list["PlaylistTrack"]] = relationship(back_populates="playlist", cascade="all, delete-orphan")


class PlaylistTrack(Base):
    __tablename__ = "playlist_tracks"

    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
    playlist_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("playlists.id", ondelete="CASCADE"), nullable=False)
    track_id: Mapped[int] = mapped_column(BigInteger, ForeignKey("tracks.id"), nullable=False)
    position: Mapped[int] = mapped_column(Integer, nullable=False)
    is_known: Mapped[bool] = mapped_column(nullable=False)
    similarity_score: Mapped[float | None] = mapped_column(REAL)

    playlist: Mapped[Playlist] = relationship(back_populates="tracks")
    track: Mapped[Track] = relationship()