"""
scripts/fast_sales_ingestion_17k.py
High-throughput multithreaded commercial sales ingestion.
Uses ThreadPoolExecutor (8-10 workers) to ingest 2,500-3,500 artists per hour.
Safe SQLite transactions via threading Lock.
"""

import sqlite3
import json
import time
import re
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

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

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=5) 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=6) 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):
    clean_text = clean_wiki_markup(wikitext)
    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

    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)

    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)

    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

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

    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 process_single_artist(artist_row):
    aid = artist_row["artist_id"]
    name = artist_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"
    return (aid, sales_units, awards, prov, 99999, sales_units, 0.0, 0.0, 0.0)

def run_fast_ingestion(limit=3000):
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cur = conn.cursor()

    rows = cur.execute("""
        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 ?
    """, (limit,)).fetchall()
    
    total = len(rows)
    if total == 0:
        print("✓ All catalog artists already ingested.")
        conn.close()
        return

    print(f"🚀 Starting high-speed ingestion for {total} artists with {MAX_WORKERS} concurrent threads...")
    start_time = time.time()

    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 (?, ?, ?, ?, ?, ?, ?, ?, ?)
    """

    completed = 0
    buffer = []

    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
        futures = {executor.submit(process_single_artist, dict(r)): r for r in rows}
        for future in as_completed(futures):
            try:
                record = future.result()
                buffer.append(record)
                completed += 1

                if len(buffer) >= 50:
                    with db_lock:
                        cur.executemany(insert_sql, buffer)
                        conn.commit()
                    elapsed = time.time() - start_time
                    rate = round(completed / elapsed, 1)
                    print(f"  ⚡ Ingested {completed}/{total} | {rate} artists/sec | Elapsed: {int(elapsed)}s")
                    buffer.clear()

            except Exception as e:
                pass

        if buffer:
            with db_lock:
                cur.executemany(insert_sql, buffer)
                conn.commit()

    conn.close()
    total_time = round(time.time() - start_time, 1)
    print(f"✓ Fast batch complete: {completed} artists processed in {total_time}s.")

if __name__ == "__main__":
    import sys
    count = int(sys.argv[1]) if len(sys.argv) > 1 else 3000
    run_fast_ingestion(count)