"""
scripts/apply_13_curated_faces.py
Finishes the remaining 5 artists using a robust requests Session with
automatic retries and connection handling to eliminate ConnectionResetError 10054.
"""

import sqlite3
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
import re
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"
AVATAR_DIR.mkdir(parents=True, exist_ok=True)

# The 5 remaining artists with primary links and guaranteed unblocked fallback mirrors
REMAINING_5 = {
    "Mist": {
        "primary": "https://tse1.mm.bing.net/th/id/OIP.lKuOO2l0qZQmi1BrPL8VzAHaHa?r=0&pid=Api&h=220&P=0",
        "fallback": "https://upload.wikimedia.org/wikipedia/commons/thumb/6/60/Mist_rapper_2019.jpg/500px-Mist_rapper_2019.jpg"
    },
    "Post Malone": {
        "primary": "https://tse4.mm.bing.net/th/id/OIP.uxymC27M9hw-tBM-tKxgvgHaE7?r=0&pid=Api&h=220&P=0",
        "fallback": "https://upload.wikimedia.org/wikipedia/commons/thumb/1/12/Post_Malone_at_the_2019_American_Music_Awards.png/500px-Post_Malone_at_the_2019_American_Music_Awards.png"
    },
    "Robb Bank$": {
        "primary": "https://tse4.mm.bing.net/th/id/OIP.8_O6-BnfOysqf_1B4daq3QHaHa?r=0&pid=Api&h=220&P=0",
        "fallback": "https://upload.wikimedia.org/wikipedia/commons/thumb/1/14/Robb_Bank%24_in_2015.jpg/500px-Robb_Bank%24_in_2015.jpg"
    },
    "Sada Baby": {
        "primary": "https://tse4.mm.bing.net/th/id/OIP.DyQEYMq6ciz0wf1QCcV3iQHaE7?r=0&pid=Api&h=220&P=0",
        "fallback": "https://upload.wikimedia.org/wikipedia/commons/thumb/c/c5/Sada_Baby_in_2020.jpg/500px-Sada_Baby_in_2020.jpg"
    },
    "Zombie Juice": {
        "primary": "https://tse4.mm.bing.net/th/id/OIP.exYMEZrpaGcSUCb8CxCXRAHaE8?r=0&pid=Api&h=220&P=0",
        "fallback": "https://upload.wikimedia.org/wikipedia/commons/thumb/7/7f/Flatbush_Zombies_-_Splash_Festival_2017_03.jpg/500px-Flatbush_Zombies_-_Splash_Festival_2017_03.jpg"
    }
}

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36",
    "Accept": "image/avif,image/webp,image/apng,image/*,*/*;q=0.8"
}

def get_resilient_session():
    """Configures a requests session that auto-retries on connection drops."""
    session = requests.Session()
    retries = Retry(
        total=3,
        backoff_factor=1,
        status_forcelist=[500, 502, 503, 504]
    )
    adapter = HTTPAdapter(max_retries=retries)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session

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 main():
    if not DB_PATH.exists():
        print(f"[-] Database not found at {DB_PATH}")
        return

    session = get_resilient_session()
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    print("\n" + "="*60)
    print("      RESOLVING REMAINING 5 ARTIST FACE PORTRAITS")
    print("="*60 + "\n")

    saved_count = 0

    for artist_name, links in REMAINING_5.items():
        cursor.execute("SELECT artist_id FROM artists WHERE artist_name = ? COLLATE NOCASE", (artist_name,))
        row = cursor.fetchone()
        artist_id = str(row[0]) if row else "0"

        slug = clean_file_slug(artist_name)
        target_path = AVATAR_DIR / f"{slug}_{artist_id}.jpg"

        print(f"[*] Downloading face for: {artist_name}...", end=" ", flush=True)

        downloaded = False
        for url in [links["primary"], links["fallback"]]:
            try:
                time.sleep(0.5)  # Pause to avoid connection rate-limiting
                res = session.get(url, headers=HEADERS, timeout=12)
                if res.status_code == 200 and len(res.content) > 2000:
                    with open(target_path, "wb") as f:
                        f.write(res.content)
                    size_kb = len(res.content) / 1024
                    print(f"SAVED -> {target_path.name} ({size_kb:.1f} KB)")

                    cursor.execute(
                        "UPDATE artists SET image_url = ?, image_verified_flag = 1 WHERE artist_id = ?",
                        (url, artist_id)
                    )
                    conn.commit()
                    saved_count += 1
                    downloaded = True
                    break
            except Exception:
                continue

        if not downloaded:
            print("FAILED")

    conn.close()
    print("\n" + "="*60)
    print(f"[+] Complete! Successfully resolved {saved_count}/5 remaining portraits.")
    print("="*60 + "\n")

if __name__ == "__main__":
    main()