"""
test_riaa_probe.py
Targeted extraction script for Cam'ron and Rich Homie Quan.
Queries RIAA gold-platinum database via GET and parses id='search-award-table'.
"""

import re
import requests
from bs4 import BeautifulSoup

session = requests.Session()
session.headers.update({
    "User-Agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
        "AppleWebKit/537.36 (KHTML, like Gecko) "
        "Chrome/124.0.0.0 Safari/537.36"
    ),
    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
})


def query_artist_riaa(artist_name: str):
    """Executes a GET query to the verified search-award-table endpoint."""
    params = {
        "tab_active": "default-award",
        "se": artist_name,
    }
    resp = session.get("https://www.riaa.com/gold-platinum/", params=params, timeout=20)
    if resp.status_code != 200:
        raise ConnectionError(f"HTTP {resp.status_code} received from RIAA.")
    return resp.text


def parse_riaa_table(html_content: str, artist_name: str):
    """Extracts awards from table id='search-award-table' and separates Album vs Single units."""
    soup = BeautifulSoup(html_content, "html.parser")
    table = soup.find("table", {"id": "search-award-table"})

    if not table:
        print(f"No search table found for {artist_name}.")
        return {
            "artist": artist_name,
            "album_units": 0,
            "singles_units": 0,
            "total_certified_units": 0,
            "awards": [],
        }

    album_units = 0
    singles_units = 0
    parsed_awards = []

    rows = table.find_all("tr")
    for r in rows:
        cells = [td.get_text(" ", strip=True) for td in r.find_all(["td", "th"])]
        if not cells or "Format" in cells or "TITLE" in [c.upper() for c in cells]:
            continue

        row_text = " | ".join(cells)

        # Detect format
        is_album = bool(re.search(r"\balbum\b", row_text, re.I))
        is_single = bool(re.search(r"\bsingle\b", row_text, re.I))

        # Detect certification multiplier
        # e.g., "2x Multi-Platinum", "Platinum", "Gold"
        multiplier = 1
        mult_match = re.search(r"(\d+)\s*[xX]\s*multi[- ]platinum", row_text, re.I)
        if mult_match:
            multiplier = int(mult_match.group(1))

        if re.search(r"platinum", row_text, re.I):
            award_units = multiplier * 1_000_000
            award_type = f"{multiplier}x Platinum" if multiplier > 1 else "Platinum"
        elif re.search(r"gold", row_text, re.I):
            award_units = 500_000
            award_type = "Gold"
        else:
            award_units = 0
            award_type = "Certified"

        if is_album:
            album_units += award_units
            parsed_awards.append({"type": "ALBUM", "award": award_type, "units": award_units, "raw": row_text})
        elif is_single:
            singles_units += award_units
            parsed_awards.append({"type": "SINGLE", "award": award_type, "units": award_units, "raw": row_text})

    print(f"\n==========================================")
    print(f"Audit Results for: {artist_name}")
    print(f"==========================================")
    print(f"Certified Album Units:  {album_units:,}")
    print(f"Certified Singles Units: {singles_units:,}")
    print(f"Total Discrete Awards:  {len(parsed_awards)}")
    for a in parsed_awards:
        print(f"  [{a['type']}] {a['award']} ({a['units']:,} units) -> {a['raw'][:90]}")

    return {
        "artist": artist_name,
        "album_units": album_units,
        "singles_units": singles_units,
        "total_certified_units": album_units + singles_units,
        "awards": parsed_awards,
    }


if __name__ == "__main__":
    test_cohort = ["Cam'ron", "Rich Homie Quan"]
    for artist in test_cohort:
        print(f"Querying RIAA database for '{artist}'...")
        html = query_artist_riaa(artist)
        parse_riaa_table(html, artist)