import os
import re
import sys
import time
import urllib.parse
from bs4 import BeautifulSoup
from curl_cffi import requests
import pandas as pd
from rapidfuzz import fuzz
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

DATA_DIR = "data"
CHECKPOINT_DIR = os.path.join(DATA_DIR, "checkpoints")
os.makedirs(CHECKPOINT_DIR, exist_ok=True)

MASTER_PATH = os.path.join(DATA_DIR, "master_artists_17k.csv")
CSPC_PATH = os.path.join(DATA_DIR, "chartmasters_audited.csv")
CHECKPOINT_FILE = os.path.join(CHECKPOINT_DIR, "sales_pipeline_progress.csv")
FINAL_OUTPUT = os.path.join(DATA_DIR, "master_sales_17k_verified.csv")

checkpoint_lock = threading.Lock()

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 is_strict_match(target_name: str, candidate_artist: str, candidate_title: str) -> bool:
    c_target = normalize_text(target_name)
    if not c_target or len(c_target) < 2:
        return False
        
    c_cand_artist = normalize_text(candidate_artist)
    c_cand_title = normalize_text(candidate_title)
    
    if c_target == c_cand_artist:
        return True
    if re.search(rf"\b{re.escape(c_target)}\b", c_cand_artist):
        return True
    if re.search(rf"\b(feat|ft|featuring|with)\s+.*?\b{re.escape(c_target)}\b", candidate_title.lower()):
        return True
    if len(c_target) > 5 and fuzz.token_sort_ratio(c_target, c_cand_artist) >= 88:
        return True
        
    return False

def fetch_riaa(artist_name: str) -> dict:
    target_clean = normalize_text(artist_name)
    if not target_clean or len(target_clean) < 2:
        return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
        
    encoded = urllib.parse.quote_plus(target_clean)
    url = f"https://www.riaa.com/gold-platinum/?tab_active=default-award&se={encoded}#search_section"
    
    try:
        resp = requests.get(url, impersonate="chrome120", timeout=12)
        if resp.status_code != 200:
            return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
            
        soup = BeautifulSoup(resp.text, "html.parser")
        table = soup.find("table")
        if not table:
            return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
            
        rows = table.find_all("tr")
        total_units = 0
        award_count = 0
        
        for r in rows[1:]:
            cells = r.find_all("td")
            if len(cells) < 3:
                continue
                
            row_artist = cells[1].get_text(strip=True)
            row_title = cells[2].get_text(strip=True)
            
            if not is_strict_match(artist_name, row_artist, row_title):
                continue
                    
            img = cells[0].find("img")
            units = extract_units_from_title(img.get("title", "") if img else "")
            if units > 0:
                total_units += units
                award_count += 1
                
        if award_count > 0:
            return {"units": total_units, "awards": award_count, "provenance": "VERIFIED_RIAA"}
            
    except Exception:
        pass
        
    return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}

def process_artist(row, cspc_dict):
    a_id = row["artist_id"]
    a_name = str(row["artist_name"])
    c_name = normalize_text(a_name)
    
    # Tier 1: CSPC
    if c_name in cspc_dict:
        c_data = cspc_dict[c_name]
        return {
            "artist_id": a_id,
            "artist_name": a_name,
            "audited_sales_units": float(c_data["total_eas_cspc"]),
            "awards_count": None,
            "provenance": "VERIFIED_CSPC"
        }
    
    # Tier 2: RIAA
    riaa_res = fetch_riaa(a_name)
    units = riaa_res["units"]
    return {
        "artist_id": a_id,
        "artist_name": a_name,
        "audited_sales_units": units if units > 0 else None,
        "awards_count": riaa_res["awards"] if riaa_res["awards"] > 0 else None,
        "provenance": riaa_res["provenance"]
    }

def run_pipeline(max_workers=5):
    print(f"Initializing Multi-Threaded Sales Pipeline ({max_workers} Workers)...")
    master_df = pd.read_csv(MASTER_PATH)
    cspc_df = pd.read_csv(CSPC_PATH)
    cspc_dict = {normalize_text(r["artist_name"]): r for _, r in cspc_df.iterrows()}
    
    if os.path.exists(CHECKPOINT_FILE):
        progress_df = pd.read_csv(CHECKPOINT_FILE)
        processed_ids = set(progress_df["artist_id"].dropna())
        records = progress_df.to_dict("records")
        print(f"Resuming from checkpoint: {len(processed_ids):,} artists already processed.")
    else:
        records = []
        processed_ids = set()

    queue_df = master_df[~master_df["artist_id"].isin(processed_ids)]
    total_remaining = len(queue_df)
    total_catalog = len(master_df)
    
    print(f"Running on remaining {total_remaining:,} artists out of {total_catalog:,}...\n")
    
    processed_count = len(processed_ids)
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_artist = {
            executor.submit(process_artist, row, cspc_dict): row 
            for _, row in queue_df.iterrows()
        }
        
        for future in as_completed(future_to_artist):
            res = future.result()
            records.append(res)
            processed_count += 1
            
            if res["audited_sales_units"] is not None and res["audited_sales_units"] > 0:
                print(f"[{processed_count:,}/{total_catalog:,}] {res['artist_name']:<22} -> {res['audited_sales_units']:>12,f} ({res['provenance']})")
            
            # Save checkpoint every 50 artists
            if processed_count % 50 == 0:
                with checkpoint_lock:
                    pd.DataFrame(records).to_csv(CHECKPOINT_FILE, index=False)

    # Final Rank & Export
    df_out = pd.DataFrame(records)
    has_sales = df_out["audited_sales_units"].notna()
    df_out["sales_rank"] = None
    df_out.loc[has_sales, "sales_rank"] = (
        df_out.loc[has_sales, "audited_sales_units"]
        .rank(ascending=False, method="min")
        .astype(int)
    )
    
    df_out.to_csv(FINAL_OUTPUT, index=False)
    print(f"\nPipeline successfully finished! Final catalog saved to: {FINAL_OUTPUT}")

if __name__ == "__main__":
    run_pipeline(max_workers=5)