import asyncio
import os
import sys
from playwright.async_api import async_playwright

async def wait_for_enter():
    # Print the prompt and flush immediately so it's visible.
    print("Press Enter to exit...", flush=True)
    loop = asyncio.get_event_loop()
    # Run sys.stdin.readline in the default executor so we don't block the event loop.
    await loop.run_in_executor(None, sys.stdin.readline)

async def save_storage_state():
    auth_state_path = "auth_state.json"
    print("\n=== Starting Authentication State Check ===")
    
    # Check if auth state exists and print the result
    has_auth_state = os.path.exists(auth_state_path)
    if has_auth_state:
        print(f"✓ Found existing auth_state.json at: {os.path.abspath(auth_state_path)}")
    else:
        print("✗ No existing auth_state.json found. Will create new authentication state.")
    
    print("\n=== Launching Browser ===")
    async with async_playwright() as playwright:
        print("• Launching Chromium...")
        browser = await playwright.chromium.launch(headless=False)
        
        if has_auth_state:
            print("• Loading stored authentication state...")
            context = await browser.new_context(storage_state=auth_state_path)
        else:
            print("• Starting fresh browser context...")
            context = await browser.new_context()
            
        print("• Opening new page...")
        page = await context.new_page()
        print("• Navigating to Hugging Face...")
        await page.goto("https://huggingface.co")
        await page.wait_for_load_state("networkidle")
        
        print("\n=== Instructions ===")
        print("1. Please check if you're logged in.")
        print("2. If not, log in manually in the browser window.")
        print("3. When ready to exit and save the current state, hit Enter below.\n")
        
        # Await the blocking call to wait for the user to press Enter.
        await wait_for_enter()
        
        print("\n=== Saving State ===")
        await context.storage_state(path=auth_state_path)
        print(f"✓ Authentication state saved to: {os.path.abspath(auth_state_path)}")
        
        print("\n=== Cleaning Up ===")
        print("• Closing browser...")
        await browser.close()
        print("✓ Done!")

if __name__ == "__main__":
    asyncio.run(save_storage_state())