"""
ingest_riaa_cohort.py
RIAA Complete Ingestion Engine for HIFFI.
Processes all remaining uncataloged artists across the 500 roster,
separates verified Album vs Single certifications, and saves progress.
"""

from pathlib import Path
import re
import time
import pandas as pd
import requests
from bs4 import BeautifulSoup

DATA_PATH = Path("data/hiffi_multiplatform_500.csv")

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) -> str:
    """Queries RIAA search table for an artist."""
    params = {
        "tab_active": "default-award",
        "se": artist_name,
    }
    try:
        resp = session.get("https://www.riaa.com/gold-platinum/", params=params, timeout=20)
        if resp.status_code == 200:
            return resp.text
    except Exception:
        pass
    return ""


def parse_awards(html_content: str, artist_name: str) -> dict:
    """Parses RIAA search rows with accurate Gold / Platinum / Multi-Platinum multipliers."""
    if not html_content:
        return {"album_units": 0, "singles_units": 0, "award_count": 0}

    soup = BeautifulSoup(html_content, "html.parser")
    table = soup.find("table", {"id": "search-award-table"})
    if not table:
        return {"album_units": 0, "singles_units": 0, "award_count": 0}

    album_units = 0
    singles_units = 0
    awards_found = 0

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

        award_cell_html = str(tds[0]).lower()
        full_row_text = r.get_text(" ", strip=True).upper()
        format_col = tds[4].get_text(" ", strip=True).upper()

        multiplier = 1
        mult_match = re.search(r"(\d+)\s*[xX]", award_cell_html + " " + full_row_text)
        if mult_match:
            multiplier = int(mult_match.group(1))

        if "diamond" in award_cell_html or "DIAMOND" in full_row_text:
            units = multiplier * 10_000_000
        elif "multi-platinum" in award_cell_html or "platinum" in award_cell_html or "PLATINUM" in full_row_text:
            units = multiplier * 1_000_000
        elif "gold" in award_cell_html or "GOLD" in full_row_text:
            units = 500_000
        else:
            units = 500_000

        if "ALBUM" in format_col:
            album_units += units
        else:
            singles_units += units

        awards_found += 1

    return {
        "album_units": album_units,
        "singles_units": singles_units,
        "award_count": awards_found,
    }


def run_batch_ingestion(batch_limit: int = 0, sleep_delay: float = 1.0):
    """
    Scrapes RIAA certifications for artists awaiting sales data.
    Setting batch_limit=0 processes all remaining pending artists.
    """
    if not DATA_PATH.exists():
        print(f"File not found: {DATA_PATH}")
        return

    df = pd.read_csv(DATA_PATH)

    for col in ["riaa_certified_album_units", "riaa_certified_singles_units"]:
        if col not in df.columns:
            df[col] = None

    # Find artists that have not yet been verified by CSPC or RIAA
    mask = (df["sales_data_status"].isin(["PENDING_INGESTION", "NOT_CATALOGED_YET"])) | (
        (df["sales_data_status"] != "VERIFIED_CSPC") & (df["sales_data_status"] != "VERIFIED_RIAA")
    )
    pending_indices = df[mask].index.tolist()

    total_pending = len(pending_indices)
    if total_pending == 0:
        print("All artists have already been processed! No pending artists remain.")
        return

    run_count = total_pending if batch_limit == 0 else min(batch_limit, total_pending)
    print(f"Total pending artists: {total_pending}")
    print(f"Running ingestion for next: {run_count} artists\n")

    processed = 0
    newly_verified = 0

    for idx in pending_indices[:run_count]:
        artist = df.at[idx, "name"]
        processed += 1
        print(f"[{processed}/{run_count}] Querying: {artist}...", end="", flush=True)

        try:
            html = query_artist_riaa(artist)
            res = parse_awards(html, artist)

            if res["award_count"] > 0:
                newly_verified += 1
                df.at[idx, "riaa_certified_album_units"] = res["album_units"]
                df.at[idx, "riaa_certified_singles_units"] = res["singles_units"]
                df.at[idx, "sales_data_status"] = "VERIFIED_RIAA"
                
                # Update total EAS baseline if empty
                if pd.isna(df.at[idx, "total_eas_cspc"]) or df.at[idx, "total_eas_cspc"] == 0:
                    df.at[idx, "total_eas_cspc"] = res["album_units"]
                    
                print(f" -> FOUND ({res['award_count']} awards | Albums: {res['album_units']:,} | Singles: {res['singles_units']:,})")
            else:
                # Keep marked as PENDING_INGESTION for Tier 3 streaming fallback
                df.at[idx, "sales_data_status"] = "PENDING_INGESTION"
                print(" -> None on file")

            # Auto-save to CSV every 10 artists so no work is ever lost
            if processed % 10 == 0:
                df.to_csv(DATA_PATH, index=False)

            time.sleep(sleep_delay)

        except Exception as e:
            print(f" -> ERROR: {e}")

    # Final save
    df.to_csv(DATA_PATH, index=False)
    print("\n==========================================")
    print(f"Ingestion Finished!")
    print(f"Total Processed: {processed}")
    print(f"Newly Verified:  {newly_verified}")
    print(f"Saved directly to: {DATA_PATH}")
    print("==========================================")


if __name__ == "__main__":
    # batch_limit=0 means: process all remaining artists until 100% complete
    run_batch_ingestion(batch_limit=0, sleep_delay=1.0)