"""
scripts/enrich_cohort_images.py
Fetches high-resolution artist avatar images via YouTube Data API v3 in batches of 50
strictly for the 500 cohort, updating image_url and image_verified_flag in SQLite.
"""
import os
import sys
from pathlib import Path
import pandas as pd
import requests
from dotenv import load_dotenv

ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

load_dotenv(ROOT_DIR / ".env")
from core.database import get_connection

YOUTUBE_API_KEY = os.getenv("YOUTUBE_API_KEY")

def chunk_list(lst, size=50):
    for i in range(0, len(lst), size):
        yield lst[i:i + size]

def main():
    print("=" * 65)
    print("HIFFI INTELLIGENCE: COHORT PROFILE IMAGE ENRICHMENT")
    print("=" * 65)

    if not YOUTUBE_API_KEY:
        print("ERROR: YOUTUBE_API_KEY is missing from .env")
        sys.exit(1)

    csv_path = ROOT_DIR / "data" / "hiffi_multiplatform_500.csv"
    if not csv_path.exists():
        csv_path = ROOT_DIR / "cleaned_artists_inventory.csv"

    df_csv = pd.read_csv(csv_path)
    name_col = "name" if "name" in df_csv.columns else "artist_name"
    cohort_names = [str(n).strip() for n in df_csv[name_col].dropna()]

    conn = get_connection()
    cur = conn.cursor()

    placeholders = ",".join(["?"] * len(cohort_names))
    cur.execute(f"""
        SELECT a.artist_id, a.artist_name, a.youtube_url
        FROM artists a
        WHERE a.artist_name IN ({placeholders})
          AND a.youtube_url LIKE 'https://www.youtube.com/channel/%'
    """, cohort_names)
    rows = cur.fetchall()

    channel_to_artist = {}
    for artist_id, name, yt_url in rows:
        ch_id = yt_url.split("/channel/")[-1].strip("/")
        if ch_id:
            channel_to_artist[ch_id] = (artist_id, name)

    all_channels = list(channel_to_artist.keys())
    print(f"Targeting {len(all_channels)} cohort channel IDs (10 batches of 50)...\n")

    updated_count = 0

    for batch_num, batch in enumerate(chunk_list(all_channels, 50), start=1):
        ids_str = ",".join(batch)
        url = "https://www.googleapis.com/youtube/v3/channels"
        params = {
            "part": "snippet",
            "id": ids_str,
            "key": YOUTUBE_API_KEY,
            "maxResults": 50,
        }

        try:
            resp = requests.get(url, params=params, timeout=15)
            if resp.status_code == 200:
                data = resp.json()
                items = data.get("items", [])
                for item in items:
                    ch_id = item["id"]
                    thumbs = item.get("snippet", {}).get("thumbnails", {})
                    img_url = (
                        thumbs.get("high", {}).get("url")
                        or thumbs.get("medium", {}).get("url")
                        or thumbs.get("default", {}).get("url")
                    )

                    if img_url and ch_id in channel_to_artist:
                        artist_id, artist_name = channel_to_artist[ch_id]
                        cur.execute("""
                            UPDATE artists
                            SET image_url = ?,
                                image_verified_flag = 1
                            WHERE artist_id = ?
                        """, (img_url, artist_id))
                        updated_count += 1
                conn.commit()
                print(f"Batch {batch_num:>2}: Saved {len(items)} official avatars.")
            else:
                print(f"Batch {batch_num:>2} error {resp.status_code}: {resp.text[:100]}")
        except Exception as e:
            print(f"Batch {batch_num:>2} failed: {e}")

    conn.close()

    print("\n" + "=" * 65)
    print(f"Completed! Saved verified profile images for {updated_count} artists.")
    print("=" * 65)

if __name__ == "__main__":
    main()