import os
import re
from bs4 import BeautifulSoup
from curl_cffi import requests
import pandas as pd
from rapidfuzz import fuzz

DATA_DIR = "data"
CATALOG_PATH = os.path.join(DATA_DIR, "master_artists_intelligence_17k.csv")
OUTPUT_PATH = os.path.join(DATA_DIR, "master_artists_intelligence_17k.csv")

KWORB_STREAMS_URL = "https://kworb.net/spotify/artists.html"
KWORB_LISTENERS_URL = "https://kworb.net/spotify/listeners.html"


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_kworb_master_table() -> dict:
  print("Fetching Kworb master stream table...")
  session = requests.Session()
  resp = session.get(KWORB_STREAMS_URL, impersonate="chrome120", timeout=15)

  if resp.status_code != 200:
    raise ConnectionError(
        f"Failed to fetch Kworb streams table (Status: {resp.status_code})"
    )

  soup = BeautifulSoup(resp.text, "html.parser")
  table = soup.find("table")
  if not table:
    raise ValueError("No table found on Kworb page.")

  kworb_data = {}
  rows = table.find_all("tr")

  # Table format: Artist | Total Streams (M) | Daily (M) | As Lead (M) | Solo (M) | As Feature (M)
  for row in rows[1:]:
    cells = row.find_all("td")
    if len(cells) < 2:
      continue

    artist_anchor = cells[0].find("a")
    artist_name = (
        artist_anchor.get_text(strip=True)
        if artist_anchor
        else cells[0].get_text(strip=True)
    )

    try:
      # Kworb streams are formatted in millions (e.g., '143,818.8' means 143.8 billion)
      stream_str = cells[1].get_text(strip=True).replace(",", "")
      total_streams = int(float(stream_str) * 1_000_000)
    except (ValueError, IndexError):
      total_streams = 0

    norm_key = normalize_name(artist_name)
    if norm_key and total_streams > 0:
      kworb_data[norm_key] = {
          "kworb_artist_name": artist_name,
          "spotify_total_streams": total_streams,
      }

  print(f"Loaded {len(kworb_data):,} artists with lifetime stream metrics.")
  return kworb_data


def run_kworb_enrichment():
  df = pd.read_csv(CATALOG_PATH)
  kworb_dict = fetch_kworb_master_table()

  if "spotify_total_streams" not in df.columns:
    df["spotify_total_streams"] = None

  unverified_mask = df["catalog_intelligence_tier"].str.contains(
      "TIER_3_UNAUDITED", na=False
  )
  unverified_indices = df[unverified_mask].index

  print(
      f"Matching Kworb streaming metrics against {len(unverified_indices):,}"
      " unverified artists..."
  )

  matches = 0
  for idx in unverified_indices:
    a_name = str(df.at[idx, "artist_name"])
    c_name = normalize_name(a_name)

    # 1. Exact match
    if c_name in kworb_dict:
      entry = kworb_dict[c_name]
      df.at[idx, "spotify_total_streams"] = entry["spotify_total_streams"]
      df.at[idx, "catalog_intelligence_tier"] = (
          "TIER_3A_VERIFIED_DIGITAL_CATALOG"
      )
      matches += 1
      continue

    # 2. Strict token match for names longer than 5 chars
    if len(c_name) > 5:
      for k_key, entry in kworb_dict.items():
        if fuzz.token_sort_ratio(c_name, k_key) >= 92:
          df.at[idx, "spotify_total_streams"] = entry["spotify_total_streams"]
          df.at[idx, "catalog_intelligence_tier"] = (
              "TIER_3A_VERIFIED_DIGITAL_CATALOG"
          )
          matches += 1
          break

  print(
      f"\nEnrichment complete! Successfully matched {matches:,} previously"
      " unverified artists."
  )
  print("\nUpdated Catalog Intelligence Tiers:")
  print(df["catalog_intelligence_tier"].value_counts())

  df.to_csv(OUTPUT_PATH, index=False)
  print(f"Saved updated catalog to: {OUTPUT_PATH}")


if __name__ == "__main__":
  run_kworb_enrichment()