main.py

import calendar
import json
import logging
import os
import random
import shutil
from collections import Counter, defaultdict
from pathlib import Path

import ijson
import matplotlib.pyplot as plt
import qrcode
import qrcode.image.svg
import spotipy
import typst
from matplotlib.backends.backend_pdf import PdfPages
from spotipy import SpotifyClientCredentials

logging.basicConfig(level=logging.INFO)
random.seed("hitster")


def get_env_var(key):
    value = os.getenv(key)
    if not value:
        raise EnvironmentError(f"Environment variable {key} is required but not set.")
    return value


def resolve_date(date_str):
    date_parts = date_str.split("-")[::-1]
    parts = [""] * (3 - len(date_parts)) + date_parts

    day = f"{int(parts[0])}." if parts[0] else ""
    month = calendar.month_name[int(parts[1])] if parts[1] else ""
    year = parts[2]

    return day, month, year

def get_spotify_songs(sp):
    jsons = os.getenv("SPOTIFY_JSONS")
    if not jsons:
        logging.info("Getting Songs from Playlist")
        return get_playlist_songs(sp, get_env_var("PLAYLIST_ID"))
    else:
        logging.info("Getting Songs from Spotify Streaming History")
        return get_streaming_history_songs(sp, jsons)

def get_playlist_songs(sp, playlist_id):
    songs = []
    results = sp.playlist_tracks(playlist_id)

    while results:
        for item in results["items"]:
            track = item["track"]
            if track:
                day, month, year = resolve_date(track["album"]["release_date"])
                songs.append(
                    {
                        "name": track["name"],
                        "artists": [artist["name"] for artist in track["artists"]],
                        "day": day,
                        "month": month,
                        "year": year,
                        "release_date": track["album"]["release_date"],
                        "url": track["external_urls"]["spotify"],
                        "id": track["id"],
                    }
                )

        results = sp.next(results) if results["next"] else None

    random.shuffle(songs)
    return songs

def get_streaming_history_songs(sp, path):
    songs = []
    field = "spotify_track_uri"
    max_songs = int(os.getenv("SPOTIFY_HISTORY_MAX", 150))

    def _process_file(file_path: Path, counts: defaultdict) -> None:
        try:
            with file_path.open("rb") as f:
                for rec in ijson.items(f, "item"):
                    entry = rec.get(field)
                    if entry:
                        counts[entry] += 1
        except Exception as exc:
            print(f"[WARN] Could not process {file_path}: {exc}")    
    
    def _analyze_streaming_history():
        base_path = Path(path)
        counts = defaultdict(int)

        if base_path.is_file():
            _process_file(base_path, counts)

        elif base_path.is_dir():
            for root, _, files in os.walk(base_path):
                for fname in files:
                    if fname.lower().endswith(".json"):
                        file_path = Path(root) / fname
                        _process_file(file_path, counts)
        else:
            raise FileNotFoundError(f"The supplied path does not exist: {base_path}")

        sorted_counts = sorted(
            counts.items(),
            key=lambda kv: kv[1],
            reverse=True
        )
        return sorted_counts
    
    def _chunk(seq, size=50):
        for i in range(0, len(seq), size):
            yield seq[i:i + size]
            
    def get_tracks_infos(track_ids):
        tracks_infos = []
        for batch in _chunk(track_ids):
            tracks = sp.tracks([song[0] for song in batch])
            tracks_infos += tracks["tracks"]
            
        return tracks_infos
    
    sorted_history = _analyze_streaming_history()
    raw_songs = get_tracks_infos(sorted_history[:max_songs])
    
    for track in raw_songs:
        day, month, year = resolve_date(track["album"]["release_date"])
        songs.append(
            {
                "name": track["name"],
                "artists": [artist["name"] for artist in track["artists"]],
                "day": day,
                "month": month,
                "year": year,
                "release_date": track["album"]["release_date"],
                "url": track["external_urls"]["spotify"],
                "id": track["id"],
            }
        )
    random.shuffle(songs)
    return songs

def generate_qr_codes(songs):
    if os.path.isdir("qr_codes"):
        shutil.rmtree("qr_codes")
    os.mkdir("qr_codes")

    for song in songs:
        img = qrcode.make(song["url"], image_factory=qrcode.image.svg.SvgPathImage)
        img.save(f"qr_codes/{song["id"]}.svg")


def generate_overview_pdf(songs, output_pdf):
    year_counts = Counter(int(song["year"]) for song in songs if "year" in song)

    min_year = min(year_counts.keys())
    max_year = max(year_counts.keys())
    all_years = list(range(min_year, max_year + 1))
    counts = [year_counts.get(year, 0) for year in all_years]

    plt.figure()
    plt.bar(all_years, counts, color="black")
    plt.ylabel("number of songs released")
    plt.xticks()

    with PdfPages(output_pdf) as pdf:
        pdf.savefig()
        plt.close()


def main():
    sp = spotipy.Spotify(
        auth_manager=SpotifyClientCredentials(
            client_id=get_env_var("CLIENT_ID"),
            client_secret=get_env_var("CLIENT_SECRET"),
        )
    )

    logging.info("Starting Spotify song retrieval")
    songs = get_spotify_songs(sp)

    logging.info("Writing songs to file")
    with open("songs.json", "w") as file:
        json.dump(songs, file, indent=4)

    logging.info("Generating QR codes")
    generate_qr_codes(songs)

    logging.info("Compiling Cards PDF")
    typst.compile("hitster.typ", output="hitster.pdf")

    logging.info("Compiling Year Overview PDF")
    generate_overview_pdf(songs, "overview.pdf")

    logging.info("Done")


if __name__ == "__main__":
    main()