import os
import re
import time
import pandas as pd
from curl_cffi import requests
from dotenv import load_dotenv

load_dotenv()

DATA_PATH = "data/master_artists_intelligence_17k_final.csv"
OUTPUT_PATH = "data/master_artists_intelligence_17k_final.csv"
LASTFM_API_KEY = os.getenv("LASTFM_API_KEY")


def clean_primary_artist(raw_name: str) -> str:
    """Extracts primary artist if string is a collaboration or feature."""
    if not raw_name or pd.isna(raw_name):
        return ""
    text = str(raw_name)
    # Strip (feat. ...), [prod. ...], etc.
    text = re.sub(r"\(.*?\)|\[.*?\]", "", text)
    # Split on explicit collaboration delimiters and take primary artist
    for delim in [" feat. ", " ft. ", " featuring ", " Feat. ", " Ft. ", " & ", " x ", " X ", " vs. ", " vs "]:
        if delim in text:
            text = text.split(delim)[0]
    return text.strip()


def query_lastfm_search(session: requests.Session, artist_name: str) -> dict:
    url = "https://ws.audioscrobbler.com/2.0/"
    params = {
        "method": "artist.search",
        "artist": artist_name,
        "api_key": LASTFM_API_KEY,
        "format": "json",
        "limit": 1
    }
    headers = {"User-Agent": "HIFFIRecovery/1.0 (annotation783@gmail.com)"}

    try:
        resp = session.get(url, params=params, headers=headers, impersonate="chrome120", timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            matches = data.get("results", {}).get("artistmatches", {}).get("artist", [])
            if matches:
                top_match = matches[0] if isinstance(matches, list) else matches
                listeners = int(top_match.get("listeners", 0))
                if listeners > 0:
                    return {
                        "verified_name": top_match.get("name"),
                        "lastfm_listeners": listeners,
                        "lastfm_url": top_match.get("url"),
                        "source": "LASTFM_SEARCH"
                    }
    except Exception:
        pass
    return None


def query_musicbrainz(session: requests.Session, artist_name: str) -> dict:
    """Open authoritative music registry fallback."""
    url = "https://musicbrainz.org/ws/2/artist/"
    params = {"query": f'artist:"{artist_name}"', "fmt": "json", "limit": 1}
    headers = {"User-Agent": "HIFFIArtistIntelligence/1.0 ( srilekhar86@gmail.com )"}

    try:
        resp = session.get(url, params=params, headers=headers, timeout=10)
        if resp.status_code == 200:
            data = resp.json()
            artists = data.get("artists", [])
            if artists:
                top_artist = artists[0]
                score = int(top_artist.get("score", 0))
                if score >= 90:  # High confidence registry match
                    mbid = top_artist.get("id")
                    return {
                        "verified_name": top_artist.get("name"),
                        "lastfm_listeners": None,
                        "lastfm_url": f"https://musicbrainz.org/artist/{mbid}",
                        "source": "MUSICBRAINZ_REGISTRY"
                    }
    except Exception:
        pass
    return None


def run_recovery():
    df = pd.read_csv(DATA_PATH)
    unverified_mask = df["catalog_rank"].isna() | (df["catalog_intelligence_tier"] == "TIER_3_UNAUDITED_DIGITAL_CATALOG")
    indices_to_recover = df[unverified_mask].index.tolist()

    print(f"Attempting secondary recovery on {len(indices_to_recover):,} unverified artists...\n")

    session = requests.Session()
    recovered_count = 0

    for idx in indices_to_recover:
        raw_name = str(df.at[idx, "artist_name"])
        cleaned = clean_primary_artist(raw_name)

        # 1. Try Last.fm search endpoint with cleaned name
        res = query_lastfm_search(session, cleaned if cleaned else raw_name)
        time.sleep(0.3)

        # 2. If no result, try MusicBrainz open registry
        if not res:
            res = query_musicbrainz(session, cleaned if cleaned else raw_name)
            time.sleep(1.0)  # MusicBrainz rate limit compliance

        if res:
            recovered_count += 1
            if res["lastfm_listeners"]:
                df.at[idx, "lastfm_listeners"] = res["lastfm_listeners"]
            if res["lastfm_url"]:
                df.at[idx, "lastfm_url"] = res["lastfm_url"]

            tier_label = "TIER_3_LASTFM_VERIFIED_ENTITY" if res["source"] == "LASTFM_SEARCH" else "TIER_3_MUSICBRAINZ_VERIFIED"
            df.at[idx, "catalog_intelligence_tier"] = tier_label
            print(f"✅ Recovered: '{raw_name}' -> '{res['verified_name']}' ({tier_label})")

    print(f"\nRecovery complete: Successfully recovered {recovered_count:,} additional entities.")
    df.to_csv(OUTPUT_PATH, index=False)
    print(f"Updated catalog saved to {OUTPUT_PATH}")


if __name__ == "__main__":
    run_recovery()