enrich.py

import json
import pandas as pd
import hashlib
import time
from pathlib import Path
from typing import Optional
import spotipy
from spotipy.oauth2 import SpotifyClientCredentials

from config import SPOTIFY_CLIENT_ID, SPOTIFY_CLIENT_SECRET, FEATURES, CACHE_DIR

def spotify_client():
    auth = SpotifyClientCredentials(
        client_id= SPOTIFY_CLIENT_ID,
        client_secret= SPOTIFY_CLIENT_SECRET,
    )
    return spotipy.Spotify(auth_manager= auth)

def cache_key(name, artist):
    raw = f'{name.lower()}::{artist.lower()}'
    return hashlib.md5(raw.encode()).hexdigest()

def cache_path(key):
    return CACHE_DIR / f'{key}.json'

def read_cache(key):
    path = cache_path(key)
    if not path.exists():
        return None
    return json.loads(path.read_text())

def write_cache(key, data):
    cache_path(key).write_text(json.dumps(data))

def audio_features(name, artist, sp):
    key = cache_key(name, artist)
    cached = read_cache(key)
    
    if cached is None:
        return cached

    query = f'track:{name} artist:{artist}'
    results = sp.search(q=query, type='track', limit=1)
    items = results.get('tracks', {}).get('items', [])
    
    print(f"[debug] Searching: {query}")
    print(f"[debug] Results: {items}")

    if not items:
        write_cache(key, {})
        return None
    
    track_id = items[0]['id']
    af = sp.audio_features([track_id])[0]

    if af is None:
        write_cache(key, {})
        return None

    features = {f: af.get(f) for f in FEATURES}
    write_cache(key, features)
    return features

def enrich_df(df, delay=0.1):
    sp = spotify_client()
    feature_rows = []

    for i, row in df.iterrows():
        af = audio_features(row['name'], row['artist'], sp)
        feature_rows.append(af or {})
        time.sleep(delay)

    features_df = pd.DataFrame(feature_rows, index= df.index)
    enriched = pd.concat([df, features_df], axis=1)

    existing_features = [f for f in FEATURES if f in enriched.columns]
    n_missing = enriched[existing_features[0]].isna().sum() if existing_features else len(df)

    print(f'\n[enrich] | Enriched {len(df)-n_missing} tracks, {n_missing} had no Spotify match.')

    return enriched