#!/usr/bin/env python3
"""
spotify_tracklist_check.py

Reads non-priority playlists missing tracklist data from Snowflake, checks each against
the Spotify API, and writes results back to the same table immediately after
each call — so a crash loses at most one row of progress.

Run the setup SQL first:
  playlist_tracklist_check_setup.sql

Usage:
  export SPOTIFY_CLIENT_ID=46132ebb48a74ab1bc007f22b99680a7
  export SPOTIFY_CLIENT_SECRET=da5eca6bef81460689fa92037ddaab77
  python spotify_tracklist_check.py
  python spotify_tracklist_check.py --stale-hours 48   # re-check rows older than 48h
  python spotify_tracklist_check.py --limit 100        # cap rows per run (useful for testing)
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

try:
    from dotenv import load_dotenv
except ImportError:
    sys.exit("Error: python-dotenv is required.\n  pip install python-dotenv")

load_dotenv(Path(__file__).parent / ".env")

try:
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization
    import snowflake.connector
except ImportError:
    sys.exit(
        "Error: snowflake-connector-python and cryptography are required.\n"
        "  pip install snowflake-connector-python cryptography"
    )

# ── Spotify ──────────────────────────────────────────────────────────────────

TOKEN_URL        = "https://accounts.spotify.com/api/token"
API_BASE         = "https://api.spotify.com/v1"
REQUEST_TIMEOUT  = 10
MAX_RETRIES      = 4
RATE_LIMIT_PAUSE = 0.1

# ── Snowflake ─────────────────────────────────────────────────────────────────

SNOWFLAKE_TABLE = "DEV_ENGINEERING.MSTEAD.NON_PRIORITY_PLAYLIST_TRACKLIST_CHECK"

SELECT_PENDING = f"""
    WITH tracklist_has_data AS (
        SELECT DISTINCT t.store_playlist_id, 'non_priority' AS source
        FROM FACTS.PROD.PLAYLISTS_NON_PRIORITY_BY_PLAYLIST_CURRENT_TRACKLIST_DBT t
        INNER JOIN FACTS.PROD.PLAYLISTS_NON_PRIORITY_BY_PLAYLIST_DBT np
            ON np.store_playlist_id = t.store_playlist_id
        INNER JOIN FACTS.PROD.NON_PRIORITY_SPOTIFY_PLAYLIST_DBT nps on nps.playlist_id = t.store_playlist_id
        WHERE np.store_id = 286
          AND t.last_added_on_date IS NOT NULL
          AND (t.removed_on IS NULL OR t.removed_on < t.last_added_on_date)
          AND (
              t.playlist_type IN ('PERSONALIZED', 'ALGORITHMIC', 'STATION', 'RADIO')
              OR (t.current_position IS NOT NULL AND t.current_position <= t.playlist_track_count)
          )
          AND nps.followers_latest > 150
    ),
    universe AS (
        SELECT store_playlist_id, 'non_priority' AS source
        FROM FACTS.PROD.PLAYLISTS_NON_PRIORITY_BY_PLAYLIST_DBT
        INNER JOIN FACTS.PROD.NON_PRIORITY_SPOTIFY_PLAYLIST_DBT m on m.playlist_id = store_playlist_id
        WHERE store_id = 286
          AND m.followers_latest > 150
    )
    SELECT u.source, u.store_playlist_id
    FROM universe u
    LEFT JOIN tracklist_has_data t
        ON t.source = u.source AND t.store_playlist_id = u.store_playlist_id
    LEFT JOIN {SNOWFLAKE_TABLE} c
        ON c.source = u.source AND c.store_playlist_id = u.store_playlist_id
    WHERE t.store_playlist_id IS NULL
      AND (c.spotify_checked_at IS NULL OR c.spotify_checked_at < DATEADD(hour, -%s, CURRENT_TIMESTAMP()))
    ORDER BY c.spotify_checked_at ASC NULLS FIRST
    {{limit_clause}}
"""

UPSERT_ROW = f"""
    MERGE INTO {SNOWFLAKE_TABLE} AS tgt
    USING (SELECT %s AS source, %s AS store_playlist_id) AS src
        ON tgt.source = src.source AND tgt.store_playlist_id = src.store_playlist_id
    WHEN MATCHED THEN UPDATE SET
        spotify_track_count = %s,
        spotify_not_found   = %s,
        spotify_checked_at  = CURRENT_TIMESTAMP()
    WHEN NOT MATCHED THEN INSERT
        (source, store_playlist_id, has_current_tracklist, has_catalog_tracks, spotify_track_count, spotify_not_found, spotify_checked_at)
    VALUES
        (src.source, src.store_playlist_id, 'N', 'N', %s, %s, CURRENT_TIMESTAMP())
"""

REFRESH_CATALOG_TRACKS = f"""
    MERGE INTO {SNOWFLAKE_TABLE} AS tgt
    USING (
        SELECT
            tl.store_playlist_id,
            tl.source,
            CASE WHEN COUNT(dtc.isrc) > 0 THEN 'Y' ELSE 'N' END AS has_catalog_tracks
        FROM (
                        SELECT DISTINCT t.store_playlist_id, 'non_priority' AS source, t.isrc
                        FROM FACTS.PROD.PLAYLISTS_NON_PRIORITY_BY_PLAYLIST_CURRENT_TRACKLIST_DBT t
                        INNER JOIN FACTS.PROD.PLAYLISTS_NON_PRIORITY_BY_PLAYLIST_DBT np
                                ON np.store_playlist_id = t.store_playlist_id
                        WHERE np.store_id = 286
                            AND t.last_added_on_date IS NOT NULL
                            AND (t.removed_on IS NULL OR t.removed_on < t.last_added_on_date)
                            AND (
                                    t.playlist_type IN ('PERSONALIZED', 'ALGORITHMIC', 'STATION', 'RADIO')
                                    OR (t.current_position IS NOT NULL AND t.current_position <= t.playlist_track_count)
                            )
        ) tl
        LEFT JOIN FACTS.PROD.DIM_TRACK_CLEAN_MV dtc
            ON  dtc.isrc       = tl.isrc
            AND dtc.offer_type IN ('all', 'stream_only', 'album_download_stream', 'track_download_stream')
        LEFT JOIN FACTS.PROD.DIM_RELEASE dr
            ON  dr.releaseid            = dtc.upc
            AND dr.product_id           IS NOT NULL
            AND dr.not_for_distribution IN (
                    'N', 'EditoriallySuspectContent', 'SMEAnalyticsDummy',
                    'IncompleteAssets', 'SwitchboardDummy', 'AccountingDummy'
                )
            AND dr.product_type = 'digital'
            AND dr.deletions    = 'N'
        GROUP BY tl.store_playlist_id, tl.source
    ) AS src
    ON  tgt.source            = src.source
    AND tgt.store_playlist_id = src.store_playlist_id
    WHEN MATCHED AND COALESCE(tgt.has_catalog_tracks, 'N') != src.has_catalog_tracks THEN
        UPDATE SET has_catalog_tracks = src.has_catalog_tracks
"""


class TokenExpired(Exception):
    pass


# ── Spotify helpers ───────────────────────────────────────────────────────────

def get_token() -> str:
    client_id     = os.environ.get("SPOTIFY_CLIENT_ID")
    client_secret = os.environ.get("SPOTIFY_CLIENT_SECRET")
    if not client_id or not client_secret:
        sys.exit("Error: set SPOTIFY_CLIENT_ID and SPOTIFY_CLIENT_SECRET env vars")

    credentials = base64.b64encode(f"{client_id}:{client_secret}".encode()).decode()
    data = urllib.parse.urlencode({"grant_type": "client_credentials"}).encode()
    req  = urllib.request.Request(
        TOKEN_URL,
        data=data,
        headers={"Authorization": f"Basic {credentials}"},
    )
    with urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT) as resp:
        return json.loads(resp.read())["access_token"]


def get_track_count(token: str, playlist_id: str) -> tuple[int | None, bool]:
    """
    Returns (track_count, not_found).
      - (N, False)    — playlist found with N tracks
            - (None, True)  — 403/404: inaccessible, deleted, or private
      - Raises TokenExpired on 401
    """
    params = urllib.parse.urlencode({"fields": "id,name,tracks.total"})
    url    = f"{API_BASE}/playlists/{playlist_id}?{params}"

    for attempt in range(MAX_RETRIES):
        req = urllib.request.Request(url, headers={"Authorization": f"Bearer {token}"})
        try:
            with urllib.request.urlopen(req, timeout=REQUEST_TIMEOUT) as resp:
                data  = json.loads(resp.read())
                count = data.get("tracks", {}).get("total")
                return (count, False)

        except urllib.error.HTTPError as e:
            if e.code in (403, 404):
                return (None, True)
            if e.code == 401:
                raise TokenExpired()
            if e.code == 429:
                wait = int(e.headers.get("Retry-After", 5))
                print(f"\n  Rate limited — waiting {wait}s...", flush=True)
                time.sleep(wait)
                continue
            if e.code >= 500:
                time.sleep(2 ** attempt)
                continue
            raise

        except (urllib.error.URLError, TimeoutError, OSError):
            if attempt < MAX_RETRIES - 1:
                time.sleep(2 ** attempt)
                continue
            print(f"\n  Network error on {playlist_id} after {MAX_RETRIES} attempts, skipping.", flush=True)
            return (None, False)

    print(f"\n  Gave up on {playlist_id} after {MAX_RETRIES} retries.", flush=True)
    return (None, False)


# ── Snowflake helpers ─────────────────────────────────────────────────────────

def get_snowflake_connection():
    params = {
        "account":   os.getenv("SNOWFLAKE_ACCOUNT", "orchard"),
        "user":      os.getenv("SNOWFLAKE_USER"),
        "database":  "DEV_ENGINEERING",
        "schema":    "MSTEAD",
        "warehouse": os.getenv("SNOWFLAKE_WAREHOUSE", "DEV_PERFORMANCE_WAREHOUSE"),
    }

    key_path = os.getenv("SNOWFLAKE_PRIVATE_KEY_PATH")
    if not key_path:
        sys.exit("Error: SNOWFLAKE_PRIVATE_KEY_PATH env var is required")
    key_file = Path(key_path).expanduser()
    if not key_file.exists():
        sys.exit(f"Error: private key file not found: {key_file}")

    passphrase = os.getenv("SNOWFLAKE_KEY_PASSPHRASE")
    with open(key_file, "rb") as f:
        p_key = serialization.load_pem_private_key(
            f.read(),
            password=passphrase.encode() if passphrase else None,
            backend=default_backend(),
        )
    params["private_key"] = p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )

    return snowflake.connector.connect(**params)


def fmt_eta(seconds: float) -> str:
    m, s = divmod(int(seconds), 60)
    h, m = divmod(m, 60)
    if h:
        return f"{h}h{m:02d}m"
    if m:
        return f"{m}m{s:02d}s"
    return f"{s}s"


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="Check Spotify for playlists missing tracklist data")
    parser.add_argument("--stale-hours", type=int, default=24,
                        help="Re-check rows last checked more than N hours ago (default 24)")
    parser.add_argument("--limit", type=int, default=0,
                        help="Max rows to process per run, 0 = no limit (default 0)")
    args = parser.parse_args()

    limit_clause = f"LIMIT {args.limit}" if args.limit else ""
    query = SELECT_PENDING.format(limit_clause=limit_clause)

    print("Connecting to Snowflake...")
    conn   = get_snowflake_connection()
    cursor = conn.cursor()

    cursor.execute(query, (args.stale_hours,))
    rows = cursor.fetchall()  # [(source, store_playlist_id), ...]

    if not rows:
        print(f"No playlists need checking (stale threshold: {args.stale_hours}h).")
        cursor.close()
        conn.close()
        return

    print(f"Found {len(rows)} playlists to check. Fetching Spotify token...\n")
    token = get_token()

    empty      = 0
    has_tracks = 0
    not_found  = 0
    start      = time.time()

    for i, (source, pid) in enumerate(rows, 1):
        elapsed = time.time() - start
        rate    = i / elapsed if elapsed > 0 else 0
        eta     = (len(rows) - i) / rate if rate > 0 else 0
        print(f"  [{i:>4}/{len(rows)}] {pid}  ETA {fmt_eta(eta)}   ", end="\r", flush=True)

        try:
            count, not_found_flag = get_track_count(token, pid)
        except TokenExpired:
            print("\n  Token expired — refreshing...", flush=True)
            token = get_token()
            count, not_found_flag = get_track_count(token, pid)

        # Write result back immediately — crash-safe
        cursor.execute(UPSERT_ROW, (source, pid, count, not_found_flag, count, not_found_flag))

        if not_found_flag:
            not_found += 1
        elif count == 0:
            empty += 1
        else:
            has_tracks += 1

        time.sleep(RATE_LIMIT_PAUSE)

    print()  # clear progress line

    elapsed = time.time() - start
    print(f"\n{'='*55}")
    print(f"  Checked          : {len(rows)}")
    print(f"  Has tracks       : {has_tracks}")
    print(f"  Empty (0 tracks) : {empty}")
    print(f"  Not found (404)  : {not_found}")
    print(f"  Elapsed          : {fmt_eta(elapsed)}")
    print(f"{'='*55}")

    print("\nRefreshing has_catalog_tracks from Snowflake catalog...", flush=True)
    cursor.execute(REFRESH_CATALOG_TRACKS)
    print("Done.")

    cursor.close()
    conn.close()


if __name__ == "__main__":
    main()