"""
scripts/download_verified_portraits_local.py
Downloads verified artist images locally to data/artist_avatars/.
Uses browser headers to prevent CDN 403 Forbidden errors.
"""

import os
import sys
import sqlite3
import requests
import re
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

BASE_DIR = Path(__file__).resolve().parent.parent
DB_PATH = BASE_DIR / "data" / "hiffi_catalog.db"
AVATAR_DIR = BASE_DIR / "data" / "artist_avatars"

# Ensure the local avatar storage folder exists
AVATAR_DIR.mkdir(parents=True, exist_ok=True)

# Headers that tell CDNs this is a standard web browser request
DEFAULT_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/svg+xml,image/*,*/*;q=0.8",
    "Accept-Language": "en-US,en;q=0.9",
    "Referer": "https://www.deezer.com/"
}

def clean_file_slug(name: str) -> str:
    """Turn artist names into safe file names (e.g. '$uicideboy$' -> '_uicideboy_')."""
    cleaned = re.sub(r'[^a-zA-Z0-9_-]', '_', name.strip().lower())
    return re.sub(r'_+', '_', cleaned).strip('_')

def download_single_artist(artist_id: str, artist_name: str, url: str):
    """Downloads one artist portrait and saves it locally."""
    if not url or url.startswith("data:image"):
        return artist_name, False, "Empty or invalid URL"

    slug = clean_file_slug(artist_name)
    ext = ".jpg"
    if ".png" in url.lower():
        ext = ".png"
    elif ".webp" in url.lower():
        ext = ".webp"

    local_filename = f"{slug}_{artist_id}{ext}"
    local_path = AVATAR_DIR / local_filename

    # If the file already exists and is not corrupt/empty (larger than 2 KB), reuse it
    if local_path.exists() and local_path.stat().st_size > 2048:
        return artist_name, True, "Already cached"

    # Set headers depending on the host
    headers = DEFAULT_HEADERS.copy()
    if "wikimedia.org" in url or "wikipedia.org" in url:
        headers["Referer"] = "https://en.wikipedia.org/"
        headers["User-Agent"] = "HIFFIArtistIntelligence/1.0 (contact@hiffi.ai)"

    try:
        response = requests.get(url, headers=headers, timeout=15)
        if response.status_code == 200:
            if len(response.content) < 1500:
                return artist_name, False, "Response too small to be an image"

            with open(local_path, "wb") as f:
                f.write(response.content)
            return artist_name, True, "Downloaded successfully"
        else:
            return artist_name, False, f"Server returned HTTP {response.status_code}"
    except Exception as exc:
        return artist_name, False, str(exc)

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

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    cursor.execute("SELECT artist_id, artist_name, image_url FROM artists WHERE image_url IS NOT NULL AND image_url != ''")
    records = cursor.fetchall()
    conn.close()

    print(f"[*] Starting local image download for {len(records)} artists...")

    successful = 0
    failures = []

    # Download multiple images at once using 8 parallel workers for speed
    with ThreadPoolExecutor(max_workers=8) as executor:
        future_to_artist = {
            executor.submit(download_single_artist, row[0], row[1], row[2]): (row[1], row[2])
            for row in records
        }

        for future in as_completed(future_to_artist):
            name, url = future_to_artist[future]
            artist_name, success, message = future.result()
            if success:
                successful += 1
            else:
                failures.append((artist_name, message, url))

    print(f"\n[+] Local Caching Complete: {successful}/{len(records)} saved to {AVATAR_DIR}")
    if failures:
        print(f"[-] Remaining failed images ({len(failures)}):")
        for fail_name, fail_msg, fail_url in failures[:10]:
            print(f"    - {fail_name}: {fail_msg} ({fail_url[:60]}...)")

if __name__ == "__main__":
    main()