import re
import urllib.parse
from bs4 import BeautifulSoup
from curl_cffi import requests
from rapidfuzz import fuzz

def normalize_text(text: str) -> str:
    if not text:
        return ""
    cleaned = re.sub(r"\(.*?\)|\[.*?\]", "", str(text))
    cleaned = re.sub(r"[^\w\s]", "", cleaned)
    return cleaned.strip().lower()

def extract_units_from_title(title_text: str) -> int:
    text = (title_text or "").upper().strip()
    mult_match = re.search(r"(\d+)\s*X", text)
    multiplier = int(mult_match.group(1)) if mult_match else 1
    
    if "DIAMOND" in text:
        return max(multiplier, 1) * 10_000_000
    elif "PLATINUM" in text or "MULTI-PLATINUM" in text:
        return multiplier * 1_000_000
    elif "GOLD" in text:
        return 500_000
    return 0

def fetch_riaa_certified_units(artist_name: str) -> dict:
    target_clean = normalize_text(artist_name)
    encoded = urllib.parse.quote_plus(target_clean)
    url = f"https://www.riaa.com/gold-platinum/?tab_active=default-award&se={encoded}#search_section"
    
    result = {
        "artist_name": artist_name,
        "riaa_total_units": 0,
        "riaa_awards_count": 0,
        "provenance": "UNVERIFIED_NO_RECORDS"
    }
    
    try:
        resp = requests.get(url, impersonate="chrome120", timeout=20)
        if resp.status_code != 200:
            return result
            
        soup = BeautifulSoup(resp.text, "html.parser")
        table = soup.find("table")
        if not table:
            return result
            
        rows = table.find_all("tr")
        total_units = 0
        award_count = 0
        
        for r in rows[1:]:
            cells = r.find_all("td")
            if len(cells) < 2:
                continue
            
            # Match artist name to avoid counting unrelated credits
            row_artist = cells[1].get_text(strip=True)
            if fuzz.token_sort_ratio(target_clean, normalize_text(row_artist)) < 70:
                # Also accept if target appears in featured song title (Cell 2)
                row_title = cells[2].get_text(strip=True) if len(cells) > 2 else ""
                if target_clean not in normalize_text(row_title):
                    continue
            
            # Extract award title from <img> tag in Cell 0
            img = cells[0].find("img")
            award_title = img.get("title", "") if img else ""
            units = extract_units_from_title(award_title)
            
            if units > 0:
                total_units += units
                award_count += 1
                
        if award_count > 0:
            result["riaa_total_units"] = total_units
            result["riaa_awards_count"] = award_count
            result["provenance"] = "VERIFIED_RIAA"
            
    except Exception as e:
        result["provenance"] = f"ERROR: {str(e)[:25]}"
        
    return result

print("--- EXECUTING VERIFIED RIAA BADGE PARSER ---")
benchmarks = ["Eminem", "Kendrick Lamar", "J. Cole", "$avage", "FakeArtist123XYZ"]
for name in benchmarks:
    res = fetch_riaa_certified_units(name)
    print(f"Artist: {res['artist_name']:<20} | Units: {res['riaa_total_units']:>12,d} | Awards: {res['riaa_awards_count']:>3} | Provenance: {res['provenance']}")