#!/usr/bin/env python3
"""
Fetch all playlists from YouTube Music brand account and save each playlist's
information to a separate JSON file.
"""

import os
import sys
import json
import time
from ytmusicapi import YTMusic
from ytmusicapi.auth.browser import setup_browser

BROWSER_AUTH_FILE = "browser.json"
BRAND_ACCOUNT_ID = "107494778873257953135"  # cuidas brand account
OUTPUT_DIR = "playlists"


def main():
    # Check if browser.json exists
    if not os.path.exists(BROWSER_AUTH_FILE):
        print(f"Error: {BROWSER_AUTH_FILE} not found!")
        print(f"Please follow the instructions in SETUP.md to create {BROWSER_AUTH_FILE}")
        sys.exit(1)
    
    # Initialize YTMusic with browser authentication
    account_info = f" (brand account: {BRAND_ACCOUNT_ID})"
    print(f"Initializing YouTube Music API with {BROWSER_AUTH_FILE}{account_info}...")
    try:
        # Load browser.json as dict
        with open(BROWSER_AUTH_FILE, 'r') as f:
            browser_data = json.load(f)
        
        # Ensure we have required keys for browser auth
        required_keys = ['Cookie', 'User-Agent']
        missing_keys = [k for k in required_keys if k not in browser_data]
        if missing_keys:
            print(f"Error: Missing required keys in {BROWSER_AUTH_FILE}: {missing_keys}")
            sys.exit(1)
        
        # Convert dict to headers string format
        headers_string = "\n".join([f"{k}: {v}" for k, v in browser_data.items()])
        
        # Use setup_browser to properly format and validate browser headers
        auth_string = setup_browser(headers_raw=headers_string, filepath=None)
        
        # Parse the result to check auth type detection
        from ytmusicapi.auth.auth_parse import determine_auth_type, parse_auth_str
        from ytmusicapi.auth.types import AuthType
        
        parsed_headers, _ = parse_auth_str(auth_string)
        detected_type = determine_auth_type(parsed_headers)
        
        # Workaround: determine_auth_type() only detects browser auth if there's an
        # "authorization" header with "SAPISIDHASH". Since browser.json doesn't have it,
        # we need to add a dummy one to pass detection. ytmusicapi will regenerate it properly.
        if detected_type == AuthType.OAUTH_CUSTOM_CLIENT:
            parsed_headers["authorization"] = "SAPISIDHASH dummy_for_detection"
            detected_type = determine_auth_type(parsed_headers)
            auth_dict = dict(parsed_headers)
            auth_string = json.dumps(auth_dict)
        
        # Initialize with the properly formatted auth string
        yt = YTMusic(auth=auth_string, user=BRAND_ACCOUNT_ID if BRAND_ACCOUNT_ID else None)
    except json.JSONDecodeError as e:
        print(f"Error parsing {BROWSER_AUTH_FILE}: {e}")
        print("Please check that your browser.json file is valid JSON.")
        sys.exit(1)
    except Exception as e:
        print(f"Error initializing YTMusic: {e}")
        print("Please check that your browser.json file is correctly formatted.")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    
    # Create output directory
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)
        print(f"Created output directory: {OUTPUT_DIR}")
    
    # Fetch all playlists
    print(f"\nFetching all playlists from brand account...")
    print("-" * 60)
    
    try:
        # Get all playlists (limit=None should get all)
        playlists = yt.get_library_playlists(limit=None)
        print(f"Found {len(playlists)} playlist(s)")
        
        if not playlists:
            print("No playlists found!")
            return
        
        # Process each playlist
        for i, playlist in enumerate(playlists, 1):
            playlist_id = playlist.get('playlistId')
            playlist_title = playlist.get('title', 'Unknown')
            
            print(f"\n[{i}/{len(playlists)}] Processing: {playlist_title} (ID: {playlist_id})")
            
            try:
                # Get full playlist details
                playlist_details = yt.get_playlist(playlist_id, limit=None)
                
                if not playlist_details:
                    print(f"  ⚠ Warning: get_playlist returned None for {playlist_id}")
                    # Save what we have from library_info
                    playlist_info = {
                        "library_info": playlist,
                        "full_details": None,
                        "error": "get_playlist returned None"
                    }
                else:
                    # Combine library playlist info with full details
                    playlist_info = {
                        "library_info": playlist,  # Info from get_library_playlists
                        "full_details": playlist_details  # Full details from get_playlist
                    }
                
                # Create safe filename from playlist title
                safe_filename = "".join(c for c in playlist_title if c.isalnum() or c in (' ', '-', '_')).rstrip()
                safe_filename = safe_filename.replace(' ', '_')
                if not safe_filename:
                    safe_filename = f"playlist_{playlist_id}"
                
                # Add playlist ID to filename to ensure uniqueness
                filename = f"{safe_filename}_{playlist_id}.json"
                filepath = os.path.join(OUTPUT_DIR, filename)
                
                # Save to file
                with open(filepath, 'w', encoding='utf-8') as f:
                    json.dump(playlist_info, f, indent=2, ensure_ascii=False)
                
                # Print summary
                if playlist_details:
                    tracks_count = playlist_details.get('trackCount', 0)
                    print(f"  ✓ Saved to: {filepath}")
                    print(f"    Tracks: {tracks_count}")
                    print(f"    Title: {playlist_details.get('title', 'Unknown')}")
                    description = playlist_details.get('description', 'N/A')
                    if description:
                        print(f"    Description: {description[:50]}...")
                    else:
                        print(f"    Description: N/A")
                else:
                    print(f"  ⚠ Saved (library info only) to: {filepath}")
                    print(f"    Title: {playlist_title}")
                
            except Exception as e:
                print(f"  ✗ Error processing playlist {playlist_title}: {e}")
                import traceback
                traceback.print_exc()
                continue
        
        # Summary
        print("\n" + "=" * 60)
        print("SUMMARY")
        print("=" * 60)
        print(f"Total playlists processed: {len(playlists)}")
        print(f"Output directory: {OUTPUT_DIR}")
        print(f"✓ Done! All playlist information saved to individual files.")
        
    except Exception as e:
        print(f"Error fetching playlists: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()