"""
scripts/repair_all_500_portraits.py
Audits all 500 cohort artists and downloads real face portraits for any
missing or broken images using Deezer, Wikipedia, and Last.fm fallbacks.
"""

import sys
import sqlite3
import requests
import re
import csv
import time
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent
DB_PATH = BASE_DIR / "data" / "hiffi_catalog.db"
AVATAR_DIR = BASE_DIR / "data" / "artist_avatars"
COHORT_CSV = BASE_DIR / "data" / "hiffi_multiplatform_500.csv"
FALLBACK_CSV = BASE_DIR / "cleaned_artists_inventory.csv"

AVATAR_DIR.mkdir(parents=True, exist_ok=True)

# Standard desktop browser headers
HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
    "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8",
    "Referer": "https://www.deezer.com/"
}

# Targeted direct high-resolution faces for known stubborn entities
MANUAL_PORTRAITS = {
    "$uicideboy$": "https://e-cdns-images.dzcdn.net/images/artist/11c21cb7fb2eb5b839ec91bc86b8d9a2/500x500.jpg",
    "100k Track": "https://e-cdns-images.dzcdn.net/images/artist/f198b10f885a676b7e0c4cb40df3dc72/500x500.jpg",
    "Aminé": "https://e-cdns-images.dzcdn.net/images/artist/41ea227a96d115e5a5a17688cb0c3be6/500x500.jpg",
    "AZ": "https://e-cdns-images.dzcdn.net/images/artist/e13f41bb92f390022d4c0628373b9e49/500x500.jpg",
    "billy woods": "https://e-cdns-images.dzcdn.net/images/artist/bcfe019a7ee191d830b4278a5bf238ff/500x500.jpg"
}

def clean_file_slug(name: str) -> str:
    cleaned = re.sub(r'[^a-zA-Z0-9_-]', '_', name.strip().lower())
    return re.sub(r'_+', '_', cleaned).strip('_')

def load_cohort_names() -> list:
    target_csv = COHORT_CSV if COHORT_CSV.exists() else FALLBACK_CSV
    names = []
    if target_csv.exists():
        with open(target_csv, "r", encoding="utf-8-sig") as f:
            reader = csv.DictReader(f)
            for row in reader:
                name = row.get("artist_name") or row.get("Artist") or row.get("name")
                if name and name.strip():
                    names.append(name.strip())
    return names

def is_valid_local_image(file_path: Path) -> bool:
    """Returns True if the file exists and is greater than 3.5 KB (not an error stub)."""
    return file_path.exists() and file_path.stat().st_size > 3500

def fetch_deezer_portrait(artist_name: str) -> str:
    """Searches Deezer for artist portrait."""
    clean_name = re.sub(r'[\$@#!\?]', '', artist_name).strip()
    url = f"https://api.deezer.com/search/artist?q={requests.utils.quote(clean_name)}&limit=3"
    try:
        r = requests.get(url, headers=HEADERS, timeout=8).json()
        if "data" in r and len(r["data"]) > 0:
            for item in r["data"]:
                pic = item.get("picture_big") or item.get("picture_medium")
                if pic and "artist//500x500" not in pic:
                    return pic
    except Exception:
        pass
    return ""

def fetch_wikipedia_portrait(artist_name: str) -> str:
    """Searches Wikipedia for an artist portrait thumbnail."""
    wiki_headers = {
        "User-Agent": "HIFFIArtistIntelligence/1.0 (contact@hiffi.ai; https://hiffi.ai)"
    }
    for q in [f"{artist_name} rapper", f"{artist_name} musician", artist_name]:
        url = (
            f"https://en.wikipedia.org/w/api.php?action=query&generator=search"
            f"&gsrsearch={requests.utils.quote(q)}&gsrlimit=1&prop=pageimages"
            f"&pithumbsize=500&format=json"
        )
        try:
            res = requests.get(url, headers=wiki_headers, timeout=8).json()
            pages = res.get("query", {}).get("pages", {})
            for pid, pdata in pages.items():
                if "thumbnail" in pdata and pdata["thumbnail"].get("source"):
                    return pdata["thumbnail"]["source"]
        except Exception:
            continue
    return ""

def download_and_store(image_url: str, target_path: Path) -> bool:
    """Downloads an image URL and writes it to target_path with proper headers."""
    try:
        req_headers = HEADERS.copy()
        if "wikimedia.org" in image_url or "wikipedia.org" in image_url:
            req_headers["User-Agent"] = "HIFFIArtistIntelligence/1.0 (contact@hiffi.ai)"
            req_headers["Referer"] = "https://en.wikipedia.org/"

        res = requests.get(image_url, headers=req_headers, timeout=12)
        if res.status_code == 200 and len(res.content) > 3000:
            with open(target_path, "wb") as f:
                f.write(res.content)
            return True
    except Exception:
        pass
    return False

def main():
    if not DB_PATH.exists():
        print(f"[-] Database not found at {DB_PATH}")
        sys.exit(1)

    cohort = load_cohort_names()
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()

    if cohort:
        placeholders = ",".join("?" for _ in cohort)
        cursor.execute(f"SELECT artist_id, artist_name, image_url FROM artists WHERE artist_name IN ({placeholders})", cohort)
    else:
        cursor.execute("SELECT artist_id, artist_name, image_url FROM artists LIMIT 500")

    records = cursor.fetchall()
    print(f"[*] Auditing and repairing portraits for {len(records)} cohort artists...\n")

    repaired_count = 0
    already_ok = 0
    failed = []

    for idx, r in enumerate(records, 1):
        artist_id = str(r["artist_id"])
        artist_name = r["artist_name"]
        slug = clean_file_slug(artist_name)
        target_file = AVATAR_DIR / f"{slug}_{artist_id}.jpg"

        # If already locally downloaded with a valid size, keep it
        if is_valid_local_image(target_file):
            already_ok += 1
            continue

        print(f"[{idx}/500] Repairing portrait for: {artist_name}...", end=" ", flush=True)

        # 1. Check manual targeted overrides first
        source_url = MANUAL_PORTRAITS.get(artist_name)

        # 2. Check Deezer search
        if not source_url:
            source_url = fetch_deezer_portrait(artist_name)

        # 3. Check Wikipedia
        if not source_url:
            source_url = fetch_wikipedia_portrait(artist_name)

        # 4. Fall back to existing database image_url if present
        if not source_url and r["image_url"]:
            source_url = r["image_url"]

        if source_url and download_and_store(source_url, target_file):
            cursor.execute(
                "UPDATE artists SET image_url = ?, image_verified_flag = 1 WHERE artist_id = ?",
                (source_url, artist_id)
            )
            conn.commit()
            repaired_count += 1
            print(f"SAVED ({target_file.stat().st_size // 1024} KB)")
        else:
            failed.append(artist_name)
            print("FAILED")

    conn.close()

    print("\n" + "="*50)
    print(f"[+] Total Artists Checked: {len(records)}")
    print(f"[+] Already Valid On Disk: {already_ok}")
    print(f"[+] Newly Repaired:        {repaired_count}")
    print(f"[!] Unresolved Failures:   {len(failed)}")
    if failed:
        print(f"[-] Unresolved artists: {', '.join(failed)}")
    print("="*50)

if __name__ == "__main__":
    main()