hiffi-artist-intelligence/scripts/run_sales_ingestion_17k.py

"""
scripts/run_sales_ingestion_17k.py
Automated ingestion of pure commercial sales and certifications across the 17K catalog.
Directly updates `commercial_sales` in data/hiffi_catalog.db.
"""

import sqlite3
import json
import time
import re
import urllib.parse
import urllib.request
from typing import Dict, Any, Tuple

DB_PATH = "data/hiffi_catalog.db"
WIKI_API_URL = "https://en.wikipedia.org/w/api.php"
HEADERS = {"User-Agent": "HIFFI-Artist-Intelligence/1.0 (sri.lekha@superlabs.co)"}


def clean_artist_query(name: str) -> str:
    clean = re.sub(r'[,$\'\"]', ' ', name)
    return re.sub(r'\s+', ' ', clean).strip()


def search_best_discography_title(artist_name: str) -> str:
    clean_name = clean_artist_query(artist_name)
    queries = [
        f"{artist_name} albums discography",
        f"{clean_name} discography",
        f"{artist_name} discography",
        artist_name
    ]
    for q in queries:
        params = {
            "action": "opensearch",
            "search": q,
            "limit": "4",
            "namespace": "0",
            "format": "json"
        }
        url = f"{WIKI_API_URL}?{urllib.parse.urlencode(params)}"
        try:
            req = urllib.request.Request(url, headers=HEADERS)
            with urllib.request.urlopen(req, timeout=6) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                titles = data[1] if len(data) > 1 else []
                for t in titles:
                    if "albums discography" in t.lower():
                        return t
                for t in titles:
                    if "discography" in t.lower():
                        return t
                if titles and (q == artist_name or q == clean_name):
                    return titles[0]
        except Exception:
            continue
    return f"{artist_name} discography"


def fetch_wiki_content(exact_title: str) -> str:
    params = {
        "action": "query",
        "prop": "revisions",
        "rvprop": "content",
        "rvslots": "main",
        "titles": exact_title,
        "format": "json",
        "redirects": "1"
    }
    url = f"{WIKI_API_URL}?{urllib.parse.urlencode(params)}"
    req = urllib.request.Request(url, headers=HEADERS)
    try:
        with urllib.request.urlopen(req, timeout=8) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            pages = data.get("query", {}).get("pages", {})
            for page_id, page_data in pages.items():
                if page_id != "-1":
                    slots = page_data.get("revisions", [{}])[0].get("slots", {})
                    return slots.get("main", {}).get("*", "")
    except Exception:
        pass
    return ""


def clean_wiki_markup(text: str) -> str:
    text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL | re.IGNORECASE)
    text = re.sub(r'<ref[^/>]*/>', '', text, flags=re.IGNORECASE)
    text = re.sub(r'\[\[(?:[^|\]]*\|)?([^\]]+)\]\]', r'\1', text)
    return text


def parse_studio_sales(wikitext: str) -> Tuple[float, int]:
    clean_text = clean_wiki_markup(wikitext)
    
    # Matches == Studio albums ==, === Studio albums ===, or Commercial mixtapes
    match = re.search(
        r'={2,4}\s*(?:Studio albums|Commercial mixtapes|Albums)\s*={2,4}.*?(?===={2,4}|\Z)',
        clean_text,
        re.DOTALL | re.IGNORECASE
    )
    album_section = match.group(0) if match else clean_text

    pure_album_sales = 0.0
    awards_count = 0

    # 1. Reported US Sales Figures
    sales_matches = re.findall(
        r'(?:US:\s*|sales:\s*|US sales:\s*)([0-9]{1,3}(?:,[0-9]{3})+)',
        album_section,
        re.IGNORECASE
    )
    seen_sales = set()
    for s in sales_matches:
        num = float(s.replace(",", ""))
        if 40_000 <= num <= 25_000_000 and num not in seen_sales:
            pure_album_sales += num
            seen_sales.add(num)

    # 2. Multi-Platinum RIAA
    multi_plat = re.findall(r'RIAA[:\s|=]*(\d+)[\xd7x\s]*Platinum', album_section, re.IGNORECASE)
    for m in multi_plat:
        mult = float(m)
        pure_album_sales += mult * 1_000_000
        awards_count += int(mult)

    # 3. Single Platinum RIAA
    cleaned_multi = re.sub(r'\d+[\xd7x\s]*Platinum', '', album_section, flags=re.IGNORECASE)
    single_plat = len(re.findall(r'RIAA[:\s|=]*Platinum', cleaned_multi, re.IGNORECASE))
    pure_album_sales += single_plat * 1_000_000
    awards_count += single_plat

    # 4. Gold RIAA
    gold = len(re.findall(r'RIAA[:\s|=]*Gold', album_section, re.IGNORECASE))
    pure_album_sales += gold * 500_000
    awards_count += gold

    # 5. General Platinum fallback if RIAA tag was omitted
    if pure_album_sales == 0.0 and awards_count == 0:
        general_multi = re.findall(r'(\d+)[\xd7x]\s*Platinum', album_section, re.IGNORECASE)
        for m in general_multi[:10]:
            mult = float(m)
            pure_album_sales += mult * 1_000_000
            awards_count += int(mult)

    return pure_album_sales, awards_count


def ingest_batch(limit: int = 500):
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cur = conn.cursor()

    # Query only artists not yet in commercial_sales
    unprocessed_query = """
        SELECT a.artist_id, a.artist_name 
        FROM artists a
        LEFT JOIN commercial_sales cs ON a.artist_id = cs.artist_id
        WHERE cs.sales_id IS NULL
        LIMIT ?
    """
    rows = cur.execute(unprocessed_query, (limit,)).fetchall()
    total_to_process = len(rows)

    if total_to_process == 0:
        print("✓ All 17,135 artists already exist in commercial_sales table.")
        conn.close()
        return

    print(f"[*] Commencing sales ingestion for next batch of {total_to_process} artists...")

    insert_sql = """
        INSERT INTO commercial_sales (
            artist_id, audited_sales_units, awards_count, provenance_flag,
            sales_rank, pure_album_sales, physical_singles_sales,
            digital_singles_sales, streaming_eas
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    """

    processed_count = 0
    for row in rows:
        aid = row["artist_id"]
        name = row["artist_name"]

        title = search_best_discography_title(name)
        wikitext = fetch_wiki_content(title) if title else ""
        
        if wikitext:
            sales_units, awards = parse_studio_sales(wikitext)
            prov = "WIKI_DISCOGRAPHY_AUDIT" if sales_units > 0 else "UNAUDITED"
        else:
            sales_units, awards = 0.0, 0
            prov = "UNAUDITED"

        # Temporary sales_rank placeholder (recalculated globally afterwards)
        cur.execute(insert_sql, (
            aid, sales_units, awards, prov, 99999,
            sales_units, 0.0, 0.0, 0.0
        ))

        processed_count += 1
        if processed_count % 25 == 0:
            conn.commit()
            print(f"  -> Ingested {processed_count}/{total_to_process} records...")

        time.sleep(0.15)

    conn.commit()
    conn.close()
    print(f"✓ Batch completed. Processed {processed_count} artists.")


if __name__ == "__main__":
    import sys
    batch_size = int(sys.argv[1]) if len(sys.argv) > 1 else 100
    ingest_batch(batch_size)