"""
test_riaa_probe.py
Targeted extraction script for Cam'ron and Rich Homie Quan.
Extracts RIAA search-award-table records and calculates certified Album and Single units.
"""

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 RIAA search table."""
    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):
    """Parses award image tags, formats, and releases for an artist."""
    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, "awards": []}

    album_units = 0
    singles_units = 0
    awards_list = []

    rows = table.find_all("tr")
    for r in rows:
        tds = r.find_all("td")
        if not tds or len(tds) < 6:
            continue

        # Cell 0: Award icon / image
        award_cell = tds[0]
        award_img = award_cell.find("img")
        award_str = ""
        if award_img:
            award_str = (award_img.get("alt") or award_img.get("src") or "").upper()
        else:
            award_str = award_cell.get_text(" ", strip=True).upper()

        # Cell 1: Artist
        artist_col = tds[1].get_text(" ", strip=True)

        # Cell 2: Title
        title_col = tds[2].get_text(" ", strip=True)

        # Cell 4: Format (ALBUM vs SINGLE)
        format_col = tds[4].get_text(" ", strip=True).upper()

        # Cell 5: Certification Date
        date_col = tds[5].get_text(" ", strip=True)

        # Multi-platinum regex check
        multiplier = 1
        mult_match = re.search(r"(\d+)\s*[xX]", award_str)
        if mult_match:
            multiplier = int(mult_match.group(1))

        # Award level assignment
        if "PLATINUM" in award_str:
            unit_val = multiplier * 1_000_000
            award_name = f"{multiplier}x Platinum" if multiplier > 1 else "Platinum"
        elif "GOLD" in award_str:
            unit_val = 500_000
            award_name = "Gold"
        elif "DIAMOND" in award_str:
            unit_val = 10_000_000
            award_name = "Diamond"
        else:
            # Fallback check across full row text if icon alt wasn't set
            row_full = r.get_text(" ", strip=True).upper()
            if "PLATINUM" in row_full:
                unit_val = 1_000_000
                award_name = "Platinum"
            elif "GOLD" in row_full:
                unit_val = 500_000
                award_name = "Gold"
            else:
                unit_val = 500_000
                award_name = "Gold (Default Tier)"

        # Classify by release format
        if "ALBUM" in format_col:
            album_units += unit_val
            fmt_type = "ALBUM"
        else:
            singles_units += unit_val
            fmt_type = "SINGLE"

        awards_list.append({
            "title": title_col,
            "artist": artist_col,
            "format": fmt_type,
            "award": award_name,
            "units": unit_val,
            "cert_date": date_col,
        })

    print(f"\n==========================================")
    print(f"RIAA Audit: {artist_name}")
    print(f"==========================================")
    print(f"Certified Album Units:  {album_units:,}")
    print(f"Certified Singles Units: {singles_units:,}")
    print(f"Total Awards Extracted: {len(awards_list)}")
    for a in awards_list:
        print(f"  • [{a['format']}] {a['title']} - {a['award']} ({a['units']:,} units) | {a['cert_date']}")

    return {
        "artist": artist_name,
        "album_units": album_units,
        "singles_units": singles_units,
        "awards": awards_list,
    }


if __name__ == "__main__":
    test_cohort = ["Cam'ron", "Rich Homie Quan"]
    for artist in test_cohort:
        html = query_artist_riaa(artist)
        parse_riaa_table(html, artist)