from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request
import os
import pickle
import requests
import logging
from pathlib import Path

# OAuth 2.0 scopes
SCOPES = ["https://www.googleapis.com/auth/youtube.readonly"]

# Configure logging
logger = logging.getLogger(__name__)

def get_token():
    """Perform OAuth authentication and save the token."""
    creds = None
    # The token file to save credentials
    token_file = "youtube_token.pickle"

    # Check if token file exists
    if os.path.exists(token_file):
        with open(token_file, "rb") as token:
            creds = pickle.load(token)

    # If no valid credentials, authenticate user
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            flow = InstalledAppFlow.from_client_secrets_file("client_secret.json", SCOPES)
            creds = flow.run_local_server(port=8080)
        # Save the credentials to a file
        with open(token_file, "wb") as token:
            pickle.dump(creds, token)
        print(f"Token saved to {token_file}")
    else:
        print("Token already exists and is valid.")
    
    return creds

def create_cookies_from_token(output_file="youtube_cookies.txt"):
    """Create a cookies file using the OAuth token for yt-dlp authentication."""
    try:
        # Get valid credentials
        creds = get_token()
        if not creds:
            logger.error("No valid credentials available")
            return False
        
        # Create a session with the OAuth token
        session = requests.Session()
        
        # Add authorization header
        headers = {
            'Authorization': f'Bearer {creds.token}',
            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        session.headers.update(headers)
        
        # Make a request to YouTube to establish session
        try:
            response = session.get('https://www.youtube.com/', timeout=30)
            if response.status_code == 200:
                logger.info("Successfully authenticated with YouTube using OAuth token")
                
                # Extract cookies from the session
                cookies_content = "# Netscape HTTP Cookie File\n# Generated from OAuth token\n\n"
                
                for cookie in session.cookies:
                    # Convert cookie to Netscape format
                    domain = cookie.domain if cookie.domain else '.youtube.com'
                    flag = 'TRUE' if domain.startswith('.') else 'FALSE'
                    path = cookie.path if cookie.path else '/'
                    secure = 'TRUE' if cookie.secure else 'FALSE'
                    expires = str(int(cookie.expires)) if cookie.expires else '0'
                    
                    cookies_content += f"{domain}\t{flag}\t{path}\t{secure}\t{expires}\t{cookie.name}\t{cookie.value}\n"
                
                # Add essential YouTube cookies if they're missing
                essential_cookies = [
                    ('.youtube.com', 'TRUE', '/', 'FALSE', '0', 'PREF', 'f1=50000000&hl=en'),
                    ('.youtube.com', 'TRUE', '/', 'FALSE', '0', 'CONSENT', 'YES+'),
                ]
                
                for domain, flag, path, secure, expires, name, value in essential_cookies:
                    if name not in [c.name for c in session.cookies]:
                        cookies_content += f"{domain}\t{flag}\t{path}\t{secure}\t{expires}\t{name}\t{value}\n"
                
                # Write cookies file
                with open(output_file, 'w') as f:
                    f.write(cookies_content)
                
                logger.info(f"Cookies file created: {output_file}")
                return True
            else:
                logger.error(f"Failed to authenticate with YouTube: {response.status_code}")
                return False
                
        except requests.RequestException as e:
            logger.error(f"Error making request to YouTube: {e}")
            return False
            
    except Exception as e:
        logger.error(f"Error creating cookies from token: {e}")
        return False

def get_oauth_headers():
    """Get OAuth authorization headers for API requests."""
    try:
        creds = get_token()
        if not creds or not creds.valid:
            return None
            
        return {
            'Authorization': f'Bearer {creds.token}',
            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36'
        }
    except Exception as e:
        logger.error(f"Error getting OAuth headers: {e}")
        return None

if __name__ == "__main__":
    import sys
    
    # Set up logging
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    
    if "--create-cookies" in sys.argv:
        # Create cookies file from OAuth token
        output_file = "youtube_cookies.txt"
        if "--output" in sys.argv:
            output_idx = sys.argv.index("--output") + 1
            if output_idx < len(sys.argv):
                output_file = sys.argv[output_idx]
        
        print("Creating cookies file from OAuth token...")
        success = create_cookies_from_token(output_file)
        if success:
            print(f"✅ Cookies file created successfully: {output_file}")
            print("This file can now be used with yt-dlp for authentication")
        else:
            print("❌ Failed to create cookies file")
            sys.exit(1)
    else:
        # Default behavior - just get/refresh token
        creds = get_token()
        if creds:
            print("✅ OAuth token is valid and ready")
            print("💡 To create cookies for yt-dlp, run: python3 auth.py --create-cookies")
        else:
            print("❌ Failed to get valid OAuth token")
            sys.exit(1)