hiffi-artist-intelligence/collectors/riaa_collector.py

"""
collectors/riaa_collector.py
Queries the official RIAA Gold & Platinum searchable database to collect
verified certifications for artists.
"""

from datetime import datetime, timezone
import json
from pathlib import Path
import re
from typing import Any, Dict
import urllib.parse
from bs4 import BeautifulSoup
import requests

RIAA_ENDPOINT = "https://www.riaa.com/wp-admin/admin-ajax.php"
HEADERS = {
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
        "AppleWebKit/537.36 (KHTML, like Gecko) "
        "Chrome/120.0.0.0 Safari/537.36"
    ),
    "Referer": "https://www.riaa.com/gold-platinum/",
    "X-Requested-With": "XMLHttpRequest",
}


class RIAACollector:
    """Collects and parses release-level RIAA certification awards."""

    def __init__(self, cache_dir: Path = Path("data/raw/riaa_cache")):
        self.cache_dir = cache_dir
        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _award_multiplier(self, award_text: str) -> int:
        """Parses RIAA award strings into certified units."""
        txt = award_text.upper()
        if "DIAMOND" in txt:
            return 10_000_000
        multi_match = re.search(r"(\d+)\s*X\s*PLATINUM", txt)
        if multi_match:
            return int(multi_match.group(1)) * 1_000_000
        if "PLATINUM" in txt:
            return 1_000_000
        if "GOLD" in txt:
            return 500_000
        return 0

    def fetch_artist_awards(self, artist_name: str) -> Dict[str, Any]:
        """
        Queries RIAA search for an artist and returns aggregated certified units.
        Separates album awards from digital singles.
        """
        clean_name = artist_name.strip()
        safe_filename = re.sub(r"[^a-zA-Z0-9_]", "_", clean_name)
        cache_file = self.cache_dir / f"{safe_filename}.json"

        if cache_file.exists():
            try:
                with open(cache_file, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception:
                pass

        payload = {
            "action": "gold_platinum_search",
            "search_term": clean_name,
            "tab_active": "default-award",
            "page": 1,
        }

        try:
            resp = requests.post(RIAA_ENDPOINT, data=payload, headers=HEADERS, timeout=12)
            if resp.status_code != 200:
                return self._empty_result(clean_name, "REQUEST_FAILED")

            soup = BeautifulSoup(resp.text, "html.parser")
            rows = soup.select("tr.table_row")

            if not rows:
                result = self._empty_result(clean_name, "NOT_FOUND_RIAA")
                with open(cache_file, "w", encoding="utf-8") as f:
                    json.dump(result, f, indent=2)
                return result

            album_units = 0
            singles_units = 0
            awards_count = 0

            for r in rows:
                artist_col = r.select_one("td.artists_sec")
                format_col = r.select_one("td.format_sec")
                award_col = r.select_one("td.award_sec")

                if not artist_col or not award_col:
                    continue

                row_artist = artist_col.get_text(strip=True).upper()
                if clean_name.upper() not in row_artist:
                    continue

                fmt = format_col.get_text(strip=True).upper() if format_col else "ALBUM"
                award_txt = award_col.get_text(strip=True)
                units = self._award_multiplier(award_txt)

                if "SINGLE" in fmt:
                    singles_units += units
                else:
                    album_units += units
                awards_count += 1

            if awards_count == 0:
                result = self._empty_result(clean_name, "NOT_FOUND_RIAA")
            else:
                result = {
                    "artist_name": clean_name,
                    "riaa_certified_album_units": album_units,
                    "riaa_certified_singles_units": singles_units,
                    "riaa_total_certified_units": album_units + singles_units,
                    "riaa_awards_count": awards_count,
                    "sales_data_status": "VERIFIED_RIAA",
                    "sales_source_url": (
                        f"https://www.riaa.com/gold-platinum/?tab_active=default-award"
                        f"&se={urllib.parse.quote(clean_name)}#search_section"
                    ),
                    "collected_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
                }

            with open(cache_file, "w", encoding="utf-8") as f:
                json.dump(result, f, indent=2)
            return result

        except Exception as e:
            return self._empty_result(clean_name, f"ERROR_{type(e).__name__}")

    def _empty_result(self, clean_name: str, status: str) -> Dict[str, Any]:
        return {
            "artist_name": clean_name,
            "riaa_certified_album_units": None,
            "riaa_certified_singles_units": None,
            "riaa_total_certified_units": None,
            "riaa_awards_count": 0,
            "sales_data_status": status,
            "sales_source_url": (
                f"https://www.riaa.com/gold-platinum/?tab_active=default-award"
                f"&se={urllib.parse.quote(clean_name)}#search_section"
            ),
            "collected_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        }