"""
scripts/fix_cohort_500.py
Realign cohort strictly to the 500 artists present in artist_metrics_snapshots
and rebalance Tiers (50 Icon, 150 Mainstream, 300 Rising).
"""

import sqlite3

DB_PATH = "data/hiffi_catalog.db"

conn = sqlite3.connect(DB_PATH)
cur = conn.cursor()

print("1. Identifying true 500 snapshot cohort artists...")
cur.execute("SELECT DISTINCT artist_id FROM artist_metrics_snapshots;")
snapshot_artist_ids = set(row[0] for row in cur.fetchall())
print(f"Total snapshot artists: {len(snapshot_artist_ids)}")

# Reset image_verified_flag = 0 for any artist NOT in the 500 snapshot cohort
print("\n2. Resetting image_verified_flag = 0 for non-cohort artists...")
cur.execute("""
    UPDATE artists 
    SET image_verified_flag = 0 
    WHERE image_verified_flag = 1 
      AND artist_id NOT IN (SELECT DISTINCT artist_id FROM artist_metrics_snapshots);
""")
updated_resets = cur.rowcount
print(f"Artists demoted from verified cohort flag: {updated_resets}")

# Ensure all 500 snapshot artists have image_verified_flag = 1
print("\n3. Ensuring all 500 snapshot cohort artists have image_verified_flag = 1...")
cur.execute("""
    UPDATE artists 
    SET image_verified_flag = 1 
    WHERE artist_id IN (SELECT DISTINCT artist_id FROM artist_metrics_snapshots);
""")
updated_cohort = cur.rowcount
print(f"Cohort verified count: {updated_cohort}")

# Rebalance Tiers: The top 500 verified cohort artists strictly map to:
# Ranks 1-50: TIER_1_ICON (50)
# Ranks 51-200: TIER_2_MAINSTREAM (150)
# Ranks 201-500: TIER_3_RISING (300)
# All others: TIER_4_DEVELOPING
print("\n4. Rebalancing tier classifications...")
# First assign tier ranks within verified cohort
cur.execute("""
    SELECT a.artist_id, r.catalog_rank
    FROM artists a
    JOIN rankings r ON a.artist_id = r.artist_id
    WHERE a.image_verified_flag = 1
    ORDER BY r.catalog_rank ASC;
""")
cohort_ranked = cur.fetchall()

for idx, (aid, cat_rank) in enumerate(cohort_ranked, start=1):
    if idx <= 50:
        tier = "TIER_1_ICON"
    elif idx <= 200:
        tier = "TIER_2_MAINSTREAM"
    else:
        tier = "TIER_3_RISING"

    cur.execute("""
        UPDATE rankings 
        SET catalog_intelligence_tier = ?, tier_rank = ? 
        WHERE artist_id = ?;
    """, (tier, idx, aid))

# Ensure any non-verified artist is strictly TIER_4_DEVELOPING
cur.execute("""
    UPDATE rankings 
    SET catalog_intelligence_tier = 'TIER_4_DEVELOPING' 
    WHERE artist_id NOT IN (SELECT artist_id FROM artists WHERE image_verified_flag = 1);
""")

conn.commit()

# Final Verification
cur.execute("SELECT COUNT(*) FROM artists WHERE image_verified_flag = 1;")
final_cohort_count = cur.fetchone()[0]

cur.execute("""
    SELECT catalog_intelligence_tier, COUNT(*) 
    FROM rankings r 
    JOIN artists a ON r.artist_id = a.artist_id 
    WHERE a.image_verified_flag = 1 
    GROUP BY catalog_intelligence_tier;
""")
final_tiers = dict(cur.fetchall())

conn.close()

print("\n" + "=" * 55)
print("REBALANCE COMPLETED")
print("=" * 55)
print(f"Verified Cohort Count : {final_cohort_count} (Strict 500)")
print(f"Tier Breakdown        : {final_tiers}")