import os
import re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import pandas as pd
from curl_cffi import requests
from dotenv import load_dotenv

load_dotenv()

DATA_DIR = "data"
INPUT_PATH = os.path.join(DATA_DIR, "master_artists_intelligence_17k_final.csv")
OUTPUT_PATH = os.path.join(DATA_DIR, "master_artists_intelligence_17k_final.csv")
CHECKPOINT_PATH = os.path.join(DATA_DIR, "checkpoints", "lastfm_verification_checkpoint.csv")

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

LASTFM_API_KEY = os.getenv("LASTFM_API_KEY")
NUM_WORKERS = 2  # Safe concurrency limit


def normalize_name(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 fetch_artist_lastfm(session: requests.Session, a_name: str) -> dict:
    target_norm = normalize_name(a_name)
    url = "https://ws.audioscrobbler.com/2.0/"
    params = {
        "method": "artist.getinfo",
        "artist": a_name,
        "api_key": LASTFM_API_KEY,
        "format": "json",
        "autocorrect": "0"
    }
    headers = {"User-Agent": "HIFFIArtistAudit/1.0 (annotation783@gmail.com)"}

    for attempt in range(4):
        try:
            resp = session.get(url, params=params, headers=headers, impersonate="chrome120", timeout=12)
            if resp.status_code == 200:
                data = resp.json()

                # Handle Last.fm Error 29: Rate Limit Exceeded
                if data.get("error") == 29:
                    time.sleep(15)
                    continue

                artist_data = data.get("artist")
                if artist_data and "name" in artist_data:
                    ret_norm = normalize_name(artist_data.get("name", ""))
                    if ret_norm == target_norm:
                        stats = artist_data.get("stats", {})
                        listeners = int(stats.get("listeners", 0))
                        playcount = int(stats.get("playcount", 0))
                        url_link = artist_data.get("url", "")

                        return {
                            "artist_name": a_name,
                            "lastfm_listeners": listeners,
                            "lastfm_playcount": playcount,
                            "lastfm_url": url_link,
                            "verified": listeners > 0
                        }
                break

            elif resp.status_code in [429, 503]:
                time.sleep(12)
                continue

        except Exception:
            time.sleep(3)

    return {
        "artist_name": a_name,
        "lastfm_listeners": None,
        "lastfm_playcount": None,
        "lastfm_url": None,
        "verified": False
    }


def run_lastfm_verification():
    if not LASTFM_API_KEY:
        print("ERROR: LASTFM_API_KEY not found in .env.")
        return

    df = pd.read_csv(INPUT_PATH)

    for col in ["lastfm_listeners", "lastfm_playcount", "lastfm_url"]:
        if col not in df.columns:
            df[col] = None

    # Load verified artists from checkpoint
    verified_map = {}
    if os.path.exists(CHECKPOINT_PATH):
        chk = pd.read_csv(CHECKPOINT_PATH)
        for _, row in chk.iterrows():
            norm_k = normalize_name(row["artist_name"])
            verified_map[norm_k] = row.to_dict()
        print(f"Resuming with {len(verified_map):,} verified artists already preserved.")

    # Populate preserved checkpoint rows into the dataframe
    for idx in df.index:
        norm_k = normalize_name(df.at[idx, "artist_name"])
        if norm_k in verified_map:
            record = verified_map[norm_k]
            df.at[idx, "lastfm_listeners"] = record.get("lastfm_listeners")
            df.at[idx, "lastfm_playcount"] = record.get("lastfm_playcount")
            df.at[idx, "lastfm_url"] = record.get("lastfm_url")
            df.at[idx, "catalog_intelligence_tier"] = "TIER_3_LASTFM_VERIFIED_ENTITY"

    # Only process rows that have not yet been verified
    unverified_mask = df["catalog_intelligence_tier"] == "TIER_3_UNAUDITED_DIGITAL_CATALOG"
    target_rows = [
        (idx, str(df.at[idx, "artist_name"]))
        for idx in df[unverified_mask].index
    ]

    total_to_process = len(target_rows)
    print(f"Verifying {total_to_process:,} unverified artists with rate-safe spacing...\n")

    if total_to_process == 0:
        print("All artists already verified.")
        return

    session = requests.Session()
    checkpoint_batch = []
    completed_count = 0
    new_verified = 0

    with ThreadPoolExecutor(max_workers=NUM_WORKERS) as executor:
        future_to_idx = {
            executor.submit(fetch_artist_lastfm, session, name): (idx, name)
            for idx, name in target_rows
        }

        for future in as_completed(future_to_idx):
            idx, name = future_to_idx[future]
            completed_count += 1
            time.sleep(0.35)

            try:
                result = future.result()
            except Exception:
                result = {"artist_name": name, "verified": False}

            if result.get("verified"):
                new_verified += 1
                df.at[idx, "lastfm_listeners"] = result["lastfm_listeners"]
                df.at[idx, "lastfm_playcount"] = result["lastfm_playcount"]
                df.at[idx, "lastfm_url"] = result["lastfm_url"]
                df.at[idx, "catalog_intelligence_tier"] = "TIER_3_LASTFM_VERIFIED_ENTITY"

                checkpoint_batch.append({
                    "artist_name": name,
                    "lastfm_listeners": result["lastfm_listeners"],
                    "lastfm_playcount": result["lastfm_playcount"],
                    "lastfm_url": result["lastfm_url"],
                    "catalog_intelligence_tier": "TIER_3_LASTFM_VERIFIED_ENTITY"
                })

            # Checkpoint every 50 newly verified artists
            if len(checkpoint_batch) >= 50 or completed_count == total_to_process:
                if checkpoint_batch:
                    chk_df = pd.DataFrame(checkpoint_batch)
                    chk_df.to_csv(CHECKPOINT_PATH, mode="a", header=not os.path.exists(CHECKPOINT_PATH), index=False)
                    checkpoint_batch = []
                print(f"[{completed_count:,}/{total_to_process:,}] Checked | New Verified: {new_verified:,} (Total Verified: {len(verified_map) + new_verified:,})")

    df.to_csv(OUTPUT_PATH, index=False)
    print(f"\nVerification complete! Saved to {OUTPUT_PATH}")


if __name__ == "__main__":
    run_lastfm_verification()