Fetch Youtube Script

"""
Viewer Minutes Calculation Method:
• Fetches video metadata (duration, view count) from YouTube API for given channels
• Calculates estimated viewer minutes = views × duration_minutes × average_watch_fraction
• Default average watch fraction is 50% (configurable via --avg-watch parameter)
• Sorts videos by estimated viewer minutes in descending order
• Outputs results to CSV with individual video stats and channel totals
"""

import argparse
import os
import time
from datetime import datetime, timezone
import pandas as pd
import isodate
from tqdm import tqdm
from dotenv import load_dotenv

# Google API
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError

# Channel mapping dictionary - add your channels here
CHANNEL_MAPPING = {
    "robert_miles": "UCLB7AzTwc6VFZrBsO2ucBMg",
    # "kurzgesagt": "UCsXVk37bltHxD1rDPwtNM8Q",
    # "veritasium": "UCHnyfMqiRRG1u-2MsSQLbXA",
    # "3blue1brown": "UCYO_jab_esuFRV4b17AJtAw",
    # "numberphile": "UCoxcjq-8xIDTYp3uz647V5A",
    "dwarkesh_podcast": "UC8aM4gqH7dFgHQWYX5_7k3Q",
    "doom_debates": "UCote8RH_wwSLza2Qb0GAQJw",
    "80k": "UCZYTClx2T1of7BRZ86-8fow",
    "rational_animations": "UCgqt1RE0k0MIr0LoyJRy2lg",
    "ai_species": "UCEENWVBdvDy-QWfuQoXC9HQ",
    "ai_in_context": "UCwicfou4Ewu-koikmga-bEg",
    "fli": "UC-rCCy3FQ-GItDimSR9lhzw",
    "cognitive_revolutions": "UCjNRVMBVI30Sak_p6HRWhIA",
    "the_ai_risk_network": "UCTEpIeib92TuKo2cSs4wg8A",
    # Add more channels as needed
    # "channel_name": "channel_id",
}

def parse_args():
    ap = argparse.ArgumentParser(description="Fetch YouTube videos and estimate viewer-minutes")
    ap.add_argument("--channels", type=str, nargs="+", help="Channel names to process (from CHANNEL_MAPPING)")
    ap.add_argument("--channel-id", type=str, default=None, help="YouTube channel ID starting with UC... (for single channel)")
    ap.add_argument("--list-channels", action="store_true", help="List available channel names and exit")
    ap.add_argument("--avg-watch", type=float, default=0.50, help="Assumed average watch fraction (0.50 = 50%)")
    ap.add_argument("--out", type=str, default="data/out.csv", help="Output CSV path")
    ap.add_argument("--max", type=int, default=None, help="Optional limit of videos to fetch per channel (for testing)")
    ap.add_argument("--sleep", type=float, default=0.0, help="Seconds to sleep between API pages (avoid quota spikes)")
    return ap.parse_args()

def get_env():
    load_dotenv()
    api_key = os.getenv("YOUTUBE_API_KEY")
    if not api_key:
        raise RuntimeError("Missing YOUTUBE_API_KEY in environment (.env)")
    return api_key

def list_available_channels():
    """Print available channel names and their IDs"""
    print("Available channels:")
    print("-" * 50)
    for name, channel_id in CHANNEL_MAPPING.items():
        print(f"{name:20} -> {channel_id}")
    print("-" * 50)
    print(f"Total: {len(CHANNEL_MAPPING)} channels")

def get_channel_ids(channel_names):
    """Convert channel names to channel IDs"""
    channel_ids = []
    invalid_channels = []
    
    for name in channel_names:
        if name in CHANNEL_MAPPING:
            channel_ids.append(CHANNEL_MAPPING[name])
        else:
            invalid_channels.append(name)
    
    if invalid_channels:
        print(f"Warning: Unknown channel names: {invalid_channels}")
        print("Available channels:")
        for name in CHANNEL_MAPPING.keys():
            print(f"  - {name}")
    
    return channel_ids

def get_uploads_playlist_id(youtube, channel_id: str) -> str:
    resp = youtube.channels().list(part="contentDetails", id=channel_id).execute()
    items = resp.get("items", [])
    if not items:
        raise RuntimeError(f"Channel not found or no contentDetails for {channel_id}")
    return items[0]["contentDetails"]["relatedPlaylists"]["uploads"]

def chunked(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i : i + n]

def fetch_videos(api_key: str, channel_id: str, limit: int | None = None, sleep_s: float = 0.0):
    youtube = build("youtube", "v3", developerKey=api_key)
    uploads_id = get_uploads_playlist_id(youtube, channel_id)

    video_ids = []
    next_token = None
    while True:
        req = youtube.playlistItems().list(
            part="contentDetails",
            playlistId=uploads_id,
            maxResults=50,
            pageToken=next_token,
        )
        resp = req.execute()
        for item in resp.get("items", []):
            video_ids.append(item["contentDetails"]["videoId"])
        next_token = resp.get("nextPageToken")
        if limit and len(video_ids) >= limit:
            video_ids = video_ids[:limit]
            break
        if not next_token:
            break
        if sleep_s > 0:
            time.sleep(sleep_s)

    results = []
    for batch in tqdm(list(chunked(video_ids, 50)), desc="Fetching video details"):
        vreq = build("youtube", "v3", developerKey=api_key).videos().list(
            part="snippet,contentDetails,statistics",
            id=",".join(batch),
        )
        vresp = vreq.execute()
        for it in vresp.get("items", []):
            vid = it["id"]
            sn = it["snippet"]
            cd = it["contentDetails"]
            st = it.get("statistics", {})

            title = sn.get("title", "")
            published_at = sn.get("publishedAt")
            try:
                duration_iso = cd.get("duration", "PT0S")
                duration_sec = isodate.parse_duration(duration_iso).total_seconds()
            except Exception:
                duration_sec = 0.0
            duration_min = duration_sec / 60.0 if duration_sec else 0.0

            views = int(st.get("viewCount", 0))

            results.append({
                "video_id": vid,
                "title": title,
                "published_at": published_at,
                "duration_min": round(duration_min, 3),
                "views": views,
            })

    return pd.DataFrame(results)

def compute_viewer_minutes(df: pd.DataFrame, avg_watch: float) -> pd.DataFrame:
    df = df.copy()
    df["avg_watch_fraction"] = avg_watch
    df["est_viewer_minutes"] = (df["views"] * df["duration_min"] * df["avg_watch_fraction"]).round(2)
    df.sort_values(["est_viewer_minutes"], ascending=False, inplace=True)
    return df

def main():
    args = parse_args()
    api_key = get_env()
    
    # Handle --list-channels option
    if args.list_channels:
        list_available_channels()
        return
    
    # Determine which channels to process
    if args.channels:
        # Process multiple channels by name
        channel_ids = get_channel_ids(args.channels)
        if not channel_ids:
            print("No valid channels specified. Use --list-channels to see available options.")
            return
    elif args.channel_id:
        # Process single channel by ID
        channel_ids = [args.channel_id]
    else:
        # Process all channels from CHANNEL_MAPPING
        channel_ids = list(CHANNEL_MAPPING.values())
        print(f"📋 No specific channels specified. Processing all {len(channel_ids)} channels from CHANNEL_MAPPING.")

    all_results = []
    
    for channel_id in channel_ids:
        # Find channel name for display
        channel_name = "unknown"
        for name, cid in CHANNEL_MAPPING.items():
            if cid == channel_id:
                channel_name = name
                break
        
        print(f"\n📺 Processing channel: {channel_name} ({channel_id})")
        
        try:
            raw = fetch_videos(api_key=api_key, channel_id=channel_id, limit=args.max, sleep_s=args.sleep)
            if not raw.empty:
                df = compute_viewer_minutes(raw, args.avg_watch)
                df['channel_name'] = channel_name
                df['channel_id'] = channel_id
                all_results.append(df)
                
                # Print channel summary
                total_views = int(df["views"].sum())
                total_vm = float(df["est_viewer_minutes"].sum())
                print(f"  ✅ {len(df)} videos | views: {total_views:,} | est. viewer-minutes: {total_vm:,.0f}")
            else:
                print(f"  ⚠️  No videos found for channel {channel_name}")
                
        except HttpError as e:
            print(f"  ❌ YouTube API error for {channel_name}: {e}")
        except Exception as e:
            print(f"  ❌ Error processing {channel_name}: {e}")

    if not all_results:
        print("No data to save.")
        return
    
    # Combine all results
    combined_df = pd.concat(all_results, ignore_index=True)
    
    # Sort by estimated viewer minutes
    combined_df.sort_values(["est_viewer_minutes"], ascending=False, inplace=True)

    # Save results
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    combined_df.to_csv(args.out, index=False, encoding="utf-8-sig")

    # Print overall summary
    total_views = int(combined_df["views"].sum())
    total_vm = float(combined_df["est_viewer_minutes"].sum())
    print(f"\n📊 Overall Results:")
    print(f"✅ Wrote {len(combined_df)} videos from {len(all_results)} channels to {args.out}")
    print(f"Totals — views: {total_views:,} | est. viewer-minutes (@{args.avg_watch:.0%}): {total_vm:,.0f}")
    
    # Print per-channel summary
    print(f"\n📈 Per-channel breakdown:")
    channel_summary = combined_df.groupby('channel_name').agg({
        'views': 'sum',
        'est_viewer_minutes': 'sum',
        'video_id': 'count'
    }).round(2)
    channel_summary.columns = ['Total Views', 'Est. Viewer Minutes', 'Video Count']
    channel_summary = channel_summary.sort_values('Est. Viewer Minutes', ascending=False)
    print(channel_summary.to_string())

if __name__ == "__main__":
    main()