"""
Playwright-based Microsoft Account Authenticator for Minecraft.

Automates the Microsoft login flow via headless Chromium to obtain:
1. Microsoft OAuth access_token
2. Xbox Live XSTS token
3. Minecraft Xbox access token (bearer token for Mojang API)

Flow:
  1. Navigate to login.live.com/oauth20_authorize.srf (Xbox Live OAuth)
  2. Enter email → click Next
  3. Handle verification: click "Other ways to sign in" → "Use my password"
  4. Enter password → click Sign in
  5. Handle 2FA if prompted (approval number or phone)
  6. Handle "Stay signed in?" prompt
  7. Extract OAuth token from redirect URL fragment
  8. Exchange for Xbox Live tokens (XBL → XSTS)
  9. Exchange for Minecraft Xbox access token

Token caching:
  - Tokens saved to ~/.ms_auth_cache/ as JSON
  - Auto-refresh before expiry using refresh_token
  - Bulk auth: load accounts.txt → auth all → save tokens

Usage:
    from playwright_microsoft_auth import MicrosoftAuthenticator

    auth = MicrosoftAuthenticator()
    
    # Single account
    result = auth.authenticate("email@example.com", "password")
    print(result.access_token)  # Minecraft bearer token
    
    # Bulk auth
    results = auth.authenticate_bulk("accounts.txt")  # email:password per line
    
    # Get cached token
    token = auth.get_cached_token("email@example.com")
"""

import json
import os
import re
import time
import logging
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional, List, Tuple
from urllib.parse import urlparse, parse_qs

from playwright.sync_api import sync_playwright, Page, Browser, BrowserContext

logger = logging.getLogger(__name__)

# Xbox Live OAuth constants (Minecraft Launcher)
XBOX_CLIENT_ID = "00000000402b5328"
XBOX_SCOPE = "service::user.auth.xboxlive.com::MBI_SSL"
MC_XBOX_SCOPE = "XboxLive.signin"
MC_XBOX_CLIENT_ID = "00000000402b5328"

# Microsoft OAuth authorize URL
# Using login.microsoftonline.com instead of login.live.com to bypass FIDO/biometric
# prompt that appears for accounts with Windows Hello configured. The microsoftonline
# endpoint presents a standard email→password flow without the FIDO intermediate step.
MS_AUTH_URL = (
    f"https://login.microsoftonline.com/common/oauth2/v2.0/authorize"
    f"?client_id={XBOX_CLIENT_ID}"
    f"&response_type=token"
    f"&redirect_uri=https://login.live.com/oauth20_desktop.srf"
    f"&scope={MC_XBOX_SCOPE}"
)

# Xbox Live API endpoints
XBL_AUTH_URL = "https://user.auth.xboxlive.com/user/authenticate"
XSTS_AUTH_URL = "https://xsts.auth.xboxlive.com/xsts/authorize"
MC_XBOX_TOKEN_URL = "https://api.minecraftservices.com/authentication/login_with_xbox"


@dataclass
class AuthResult:
    """Result of a successful authentication."""
    email: str
    username: str  # Minecraft username
    msa_access_token: str       # Microsoft OAuth token
    msa_refresh_token: str      # Microsoft refresh token
    xbl_token: str              # Xbox Live token
    xbl_token_issue_instant: str
    xbl_token_expiry: str
    xsts_token: str             # XSTS token
    xsts_token_issue_instant: str
    xsts_token_expiry: str
    mc_access_token: str        # Minecraft bearer token (what we actually use)
    mc_token_expiry: str
    xuid: str                   # Xbox User ID
    gamertag: str
    success: bool = True
    error: Optional[str] = None
    
    def to_dict(self):
        return asdict(self)
    
    @classmethod
    def failure(cls, email: str, error: str) -> 'AuthResult':
        return cls(
            email=email, username="", msa_access_token="", msa_refresh_token="",
            xbl_token="", xbl_token_issue_instant="", xbl_token_expiry="",
            xsts_token="", xsts_token_issue_instant="", xsts_token_expiry="",
            mc_access_token="", mc_token_expiry="", xuid="", gamertag="",
            success=False, error=error
        )


class MicrosoftAuthenticator:
    """
    Authenticates Microsoft accounts via Playwright headless browser.
    
    Handles the full login flow including password entry, 2FA prompts,
    and token extraction. Supports bulk authentication and token caching.
    """
    
    def __init__(self, cache_dir: Optional[str] = None, headless: bool = True):
        self.cache_dir = Path(cache_dir or os.path.expanduser("~/.ms_auth_cache"))
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.headless = headless
        self._playwright = None
        self._browser = None
        self._context = None
    
    def _ensure_browser(self):
        """Lazy-initialize Playwright browser."""
        if self._browser is None:
            self._playwright = sync_playwright().start()
            self._browser = self._playwright.chromium.launch(
                headless=self.headless,
                args=[
                    '--disable-blink-features=AutomationControlled',
                    '--no-sandbox',
                ]
            )
            self._context = self._browser.new_context(
                user_agent=(
                    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) '
                    'AppleWebKit/537.36 (KHTML, like Gecko) '
                    'Chrome/120.0.0.0 Safari/537.36'
                ),
                viewport={'width': 1280, 'height': 720},
            )
            # Stealth: remove navigator.webdriver
            self._context.add_init_script("""
                Object.defineProperty(navigator, 'webdriver', {get: () => undefined});
            """)
    
    def close(self):
        """Close the browser."""
        if self._browser:
            self._browser.close()
            self._browser = None
        if self._playwright:
            self._playwright.stop()
            self._playwright = None
    
    def __enter__(self):
        self._ensure_browser()
        return self
    
    def __exit__(self, *args):
        self.close()
    
    def _new_page(self) -> Page:
        """Create a new page in the browser context."""
        self._ensure_browser()
        return self._context.new_page()
    
    # -------------------------------------------------------------------------
    # Main authentication flow
    # -------------------------------------------------------------------------
    
    def authenticate(self, email: str, password: str, 
                     mfa_handler=None) -> AuthResult:
        """
        Authenticate a single Microsoft account.
        
        Args:
            email: Microsoft account email
            password: Account password
            mfa_handler: Optional callable(page) -> bool for handling 2FA.
                        Return True to continue, False to abort.
        
        Returns:
            AuthResult with tokens or error
        """
        page = self._new_page()
        try:
            logger.info(f"Authenticating {email}...")
            
            # Step 1: Navigate to Microsoft OAuth
            page.goto(MS_AUTH_URL, wait_until="domcontentloaded", timeout=30000)
            self._wait_stable(page, 2000)
            
            # Step 2: Enter email
            self._enter_email(page, email)
            self._wait_stable(page, 3000)
            
            # Step 3: Handle what comes after email (password, FIDO, verification)
            # Loop to handle cases where password submission triggers FIDO → Back → password again
            max_password_attempts = 5
            for pw_attempt in range(max_password_attempts):
                self._handle_post_email(page, email, password, mfa_handler)
                # Check if we ended up stuck on FIDO/intermediate page
                self._wait_stable(page, 2000)
                current_url = page.url.lower()
                current_body = page.inner_text('body').lower()
                if 'fido' in current_url or ('face' in current_body and 'fingerprint' in current_body):
                    logger.info(f"  Still on FIDO page after password attempt {pw_attempt + 1}, trying again")
                    self._switch_from_biometric_to_password(page)
                    continue
                # Check if we got the token redirect (success)
                if 'access_token' in current_url or 'oauth20_desktop.srf' in current_url:
                    logger.info(f"  Got token redirect after password attempt {pw_attempt + 1}")
                    break
                # Otherwise we're on a normal page (stay signed in, etc.) — proceed
                break
            else:
                logger.error(f"  Failed to escape FIDO loop after {max_password_attempts} attempts")

            
            # Step 4: Handle "Stay signed in?" prompt
            self._handle_stay_signed_in(page)
            
            # Step 5: Wait for redirect and extract token
            token_data = self._extract_oauth_token(page)
            if not token_data:
                return AuthResult.failure(email, "Failed to extract OAuth token from redirect")
            
            msa_access_token = token_data.get('access_token', '')
            msa_refresh_token = token_data.get('refresh_token', '')
            
            if not msa_access_token:
                return AuthResult.failure(email, "No access_token in OAuth response")
            
            logger.info(f"Got Microsoft OAuth token for {email} (len={len(msa_access_token)}, prefix={msa_access_token[:20]}...)")
            
            # Step 6: Exchange for Xbox Live tokens
            xbl_result = self._authenticate_xbl(msa_access_token)
            if not xbl_result:
                return AuthResult.failure(email, "XBL authentication failed — check logs for XBL response details")
            
            # Step 7: Exchange for XSTS token
            xsts_result = self._authenticate_xsts(xbl_result)
            if not xsts_result:
                return AuthResult.failure(email, "XSTS authentication failed")
            
            # Step 8: Exchange for Minecraft token
            mc_result = self._authenticate_minecraft(xsts_result)
            if not mc_result:
                return AuthResult.failure(email, "Minecraft authentication failed")
            
            # Step 9: Get Minecraft username
            username = self._get_minecraft_username(mc_result['access_token'])
            
            # Step 10: Save to cache
            result = AuthResult(
                email=email,
                username=username,
                msa_access_token=msa_access_token,
                msa_refresh_token=msa_refresh_token,
                xbl_token=xbl_result['Token'],
                xbl_token_issue_instant=xbl_result['TokenIssueInstant'],
                xbl_token_expiry=xbl_result['NotAfter'],
                xsts_token=xsts_result['Token'],
                xsts_token_issue_instant=xsts_result['TokenIssueInstant'],
                xsts_token_expiry=xsts_result['NotAfter'],
                mc_access_token=mc_result['access_token'],
                mc_token_expiry=mc_result.get('expires_in', '3600'),
                xuid=xsts_result.get('XUID', ''),
                gamertag=xsts_result.get('DisplayClaim', ''),
            )
            
            self._cache_result(result)
            logger.info(f"✓ Authenticated {email} as {username}")
            return result
            
        except Exception as e:
            logger.error(f"Auth failed for {email}: {e}")
            return AuthResult.failure(email, str(e))
        finally:
            page.close()
    
    # -------------------------------------------------------------------------
    # Login flow steps
    # -------------------------------------------------------------------------
    
    def _enter_email(self, page: Page, email: str):
        """Enter email and click Next."""
        # Wait for email input (try multiple selectors)
        email_selectors = [
            'input[type="email"]',
            'input[id="usernameEntry"]',
            'input[name="loginfmt"]',  # Microsoft's actual email input name
            'input[type="text"]',  # Fallback
        ]
        email_input = None
        for selector in email_selectors:
            try:
                email_input = page.wait_for_selector(selector, timeout=5000)
                if email_input and email_input.is_visible():
                    break
            except Exception:
                continue
        
        if not email_input:
            logger.error(f"  Could not find email input field. URL: {page.url}")
            # Take diagnostic screenshot
            try:
                page.screenshot(path=str(self.cache_dir / f"email_input_debug_{time.time()}.png"))
            except:
                pass
            raise Exception("Could not find email input field")
        
        email_input.fill(email)
        self._wait_stable(page, 500)
        
        # Click Next (submit button)
        next_btn = page.query_selector('button[type="submit"]')
        if not next_btn:
            next_btn = page.query_selector('button:has-text("Next")')
        if next_btn:
            next_btn.click()
        else:
            logger.warning(f"  Could not find Next button, trying to submit form")
            # Try pressing Enter in the email field as fallback
            email_input.press('Enter')
    
    def _handle_post_email(self, page: Page, email: str, password: str, mfa_handler):
        """
        Handle what comes after entering email.
        
        Uses a retry loop to handle intermediate pages (FIDO, verification, CAPTCHA)
        that may appear before the password input.
        """
        if 'access_token' in page.url:
            logger.info(f"  Already redirected with token")
            return
        
        # If we're already on FIDO, don't try to enter email/password — bail out
        # so the caller's retry loop can handle it
        url_lower = page.url.lower()
        body_text_check = page.inner_text('body').lower()
        if 'fido' in url_lower or ('face' in body_text_check and 'fingerprint' in body_text_check):
            logger.info(f"  Already on FIDO page, caller should handle recovery")
            return
        
        max_retries = 5
        for attempt in range(max_retries):
            self._wait_stable(page, 2000)
            body_text = page.inner_text('body')
            pw_input = page.query_selector('input[type="password"]')
            
            if pw_input:
                logger.info(f"  Password page found (attempt {attempt + 1})")
                pw_input.fill(password)
                self._wait_stable(page, 500)
                sign_in_btn = page.query_selector('button[type="submit"]')
                if not sign_in_btn:
                    sign_in_btn = page.query_selector('button:has-text("Sign in")')
                if sign_in_btn:
                    sign_in_btn.click()
                self._handle_post_password(page, mfa_handler)
                return
            
            url_lower = page.url.lower()
            body_lower = body_text.lower()
            
            # FIDO / biometric page
            if 'fido' in url_lower or 'face' in body_lower or 'fingerprint' in body_lower or 'security key' in body_lower or 'windows hello' in body_lower:
                logger.info(f"  [attempt {attempt + 1}] FIDO/biometric page detected, switching to password")
                self._switch_from_biometric_to_password(page)
                continue
            
            # Verification page
            if 'Send code' in body_text or 'Verify' in body_text or 'proofConfirmation' in url_lower:
                logger.info(f"  [attempt {attempt + 1}] Verification page detected, switching to password")
                self._switch_to_password_flow(page)
                continue
            
            # CAPTCHA / unusual activity
            if 'unusual' in body_lower or 'captcha' in body_lower or 'are you sure' in body_lower:
                logger.warning(f"  [attempt {attempt + 1}] CAPTCHA/unusual activity — trying to continue")
                continue_btn = page.query_selector('button:has-text("Continue"), button:has-text("Yes"), button[type="submit"]')
                if continue_btn:
                    continue_btn.click()
                continue
            
            # If we're on the email page (back navigation from FIDO), re-enter email
            email_input = page.query_selector('input[type="email"], input[name="loginfmt"]')
            if email_input and 'email' in body_lower:
                logger.info(f"  [attempt {attempt + 1}] Back on email page, re-entering email")
                email_input.fill(email)
                self._wait_stable(page, 500)
                next_btn = page.query_selector('button[type="submit"]')
                if next_btn:
                    next_btn.click()
                self._wait_stable(page, 3000)
                continue
            
            # Unknown page — log and break
            logger.warning(f"  [attempt {attempt + 1}] Unknown page. URL: {page.url}, Body: {body_text[:200]}")
            form_info = page.evaluate("""() => {
                const inputs = Array.from(document.querySelectorAll('input')).map(i => ({type: i.type, id: i.id, name: i.name}));
                const buttons = Array.from(document.querySelectorAll('button, [role="button"]')).map(b => b.textContent.trim().substring(0, 50));
                return {inputs, buttons, url: window.location.href};
            }""")
            logger.warning(f"  Form elements: {json.dumps(form_info, indent=2)[:500]}")
            break
    
    def _switch_to_password_flow(self, page: Page):
        """
        On the verification page, click 'Other ways to sign in' → 'Use my password'.
        """
        # Try clicking 'Other ways to sign in' (idA_PWD_SwitchToCredPicker)
        other_ways_selectors = [
            '#idA_PWD_SwitchToCredPicker',
            'span[role="button"]:has-text("Other ways")',
            '[data-testid="credentialPickerLink"]',
        ]
        
        clicked = False
        for selector in other_ways_selectors:
            try:
                el = page.query_selector(selector)
                if el and el.is_visible():
                    el.click()
                    clicked = True
                    self._wait_stable(page, 2000)
                    break
            except Exception:
                continue
        
        if clicked:
            # Now click 'Use my password'
            pw_selectors = [
                'button:has-text("Use my password")',
                'button:has-text("password")',
            ]
            for selector in pw_selectors:
                try:
                    btn = page.query_selector(selector)
                    if btn and btn.is_visible():
                        btn.click()
                        self._wait_stable(page, 2000)
                        return
                except Exception:
                    continue
        
        # Fallback: use JS to find and click
        page.evaluate("""
            const spans = Array.from(document.querySelectorAll('span[role="button"]'));
            for (const span of spans) {
                if (span.textContent.includes('Other') && span.textContent.includes('sign')) {
                    span.click();
                    break;
                }
            }
        """)
        self._wait_stable(page, 2000)
        
        page.evaluate("""
            const buttons = Array.from(document.querySelectorAll('button'));
            for (const btn of buttons) {
                if (btn.textContent.includes('password')) {
                    btn.click();
                    break;
                }
            }
        """)
        self._wait_stable(page, 2000)

    def _switch_from_biometric_to_password(self, page: Page):
        """
        Handle Windows Hello / FIDO / biometric page.
        
        When MS shows 'Face, fingerprint, PIN or security key' or similar,
        we need to find the 'Try another way' or 'Use my password instead' option.
        """
        logger.info(f"    Handling biometric/FIDO page...")
        
        # Strategy 0: Use JS to find and click "Sign-in options" (most reliable)
        # The element is a div[role="button"] with className "table" inside promoted-fed-cred-box
        try:
            clicked = page.evaluate("""() => {
                // Find the Sign-in options container
                const fedBox = document.querySelector('.promoted-fed-cred-box, .ext-promoted-fed-cred-box');
                if (fedBox) {
                    // Find the clickable button inside it
                    const btn = fedBox.querySelector('[role="button"], button, a');
                    if (btn) {
                        btn.click();
                        return {clicked: true, text: btn.textContent.trim().substring(0, 50)};
                    }
                    // Try clicking the container itself
                    fedBox.click();
                    return {clicked: true, text: 'fed-cred-box'};
                }
                
                // Fallback: search all elements
                const all = Array.from(document.querySelectorAll('[role="button"], button, div[role="button"]'));
                for (const el of all) {
                    const text = el.textContent.trim().toLowerCase();
                    if (text.includes('sign-in options') || text.includes('sign in options')) {
                        el.click();
                        return {clicked: true, text: el.textContent.trim().substring(0, 50)};
                    }
                }
                return {clicked: false, text: 'not found'};
            }""")
            if clicked.get('clicked'):
                logger.info(f"    Clicked Sign-in options via JS: {clicked.get('text')}")
                self._wait_stable(page, 4000)
                # After clicking, look for password option in the expanded panel
                pw_option = page.query_selector('button:has-text("Password"), [role="button"]:has-text("Password"), a:has-text("Password"), li:has-text("Password")')
                if pw_option:
                    logger.info(f"    Clicking 'Password' option")
                    pw_option.click()
                    self._wait_stable(page, 5000)
                    # Check if we got to password page
                    if page.query_selector('input[type="password"]'):
                        logger.info(f"    Got password page after Sign-in options!")
                        return True
                    return True
                # Check if we're off the FIDO page
                if 'fido' not in page.url.lower():
                    logger.info(f"    Sign-in options navigated away from FIDO")
                    return True
        except Exception as e:
            logger.info(f"    Sign-in options JS click failed: {e}")
        
        # Strategy 1: Click "Back" button (present on some FIDO pages)
        try:
            back_btn = page.query_selector('button:has-text("Back"), button:has-text("back")')
            if not back_btn:
                # Try Microsoft's standard Back button ID (input type=button)
                back_btn = page.query_selector('#idBtn_Back')
            if back_btn:
                logger.info(f"    Clicking 'Back' to escape FIDO page")
                back_btn.click()
                self._wait_stable(page, 5000)
                new_body = page.inner_text('body')
                if 'password' in new_body.lower() and 'fido' not in page.url.lower():
                    logger.info(f"    Successfully backed out to password page")
                    return True
                if 'email' in new_body.lower() or 'account' in new_body.lower():
                    logger.info(f"    Backed out to account selection page")
                    return True
        except Exception as e:
            logger.info(f"    'Back' button click failed: {e}")
        
        # Strategy 2: Click "Try another way" link
        try:
            another_way = page.query_selector('a:has-text("Try another way"), a:has-text("try another way")')
            if another_way:
                logger.info(f"    Clicking 'Try another way'")
                another_way.click()
                self._wait_stable(page, 3000)
                return True
        except Exception:
            pass
        
        # Strategy 3: JS fallback — find and click anything useful
        result = page.evaluate("""() => {
            const all = Array.from(document.querySelectorAll('a, button, [role="button"], div[role="button"]'));
            for (const el of all) {
                const text = el.textContent.toLowerCase().trim();
                if (text.includes('sign-in options') || text.includes('sign in options') || 
                    text.includes('another way') || text.includes('password instead') || 
                    text === 'back' || (text.includes('password') && !text.includes('enter'))) {
                    el.click();
                    return true;
                }
            }
            return false;
        }""")
        if result:
            logger.info(f"    JS fallback clicked alternative option")
            self._wait_stable(page, 3000)
            return True
        
        logger.warning(f"    Could not escape FIDO page with any strategy")
        return False

    def _handle_post_password(self, page: Page, mfa_handler):
        """Handle what comes after entering password (2FA, stay signed in, etc.)."""
        self._wait_stable(page, 5000)
        
        body_text = page.inner_text('body')
        
        # Check for FIDO/biometric page appearing after password (MS sometimes redirects here)
        if 'fido' in page.url.lower() or 'face' in body_text.lower() or 'fingerprint' in body_text.lower() or 'security key' in body_text.lower():
            logger.info(f"  FIDO/biometric page after password, switching away")
            self._switch_from_biometric_to_password(page)
            self._wait_stable(page, 3000)
            body_text = page.inner_text('body')
        
        # Check for 2FA / approval number
        if 'approval' in body_text.lower() or 'verify' in body_text.lower() or 'two-step' in body_text.lower():
            logger.info(f"  2FA prompt detected")
            if mfa_handler:
                if not mfa_handler(page):
                    raise Exception("MFA handler aborted")
            else:
                logger.warning(f"  2FA required but no handler provided")
                # Try to wait and see if it auto-resolves
                self._wait_stable(page, 5000)
        
        # Check for phone/SMS verification
        if 'phone' in body_text.lower() and ('verify' in body_text.lower() or 'code' in body_text.lower()):
            logger.info(f"  Phone verification detected")
            if mfa_handler:
                mfa_handler(page)
    
    def _handle_stay_signed_in(self, page: Page):
        """Handle the 'Stay signed in?' prompt if it appears."""
        try:
            self._wait_stable(page, 2000)
            body_text = page.inner_text('body')
            
            if 'stay signed in' in body_text.lower() or 'keep me signed' in body_text.lower():
                logger.info(f"  Stay signed in prompt detected")
                # Click "Yes" or the primary button (try multiple selectors)
                yes_btn = page.query_selector('button:has-text("Yes")')
                if not yes_btn:
                    yes_btn = page.query_selector('button:has-text("Keep me")')
                if not yes_btn:
                    yes_btn = page.query_selector('button[type="submit"]')
                if yes_btn:
                    yes_btn.click()
                    self._wait_stable(page, 3000)
        except Exception:
            pass  # Not critical
    
    def _extract_oauth_token(self, page: Page) -> Optional[dict]:
        """
        Extract the OAuth token from the redirect URL.
        
        After successful auth, Microsoft redirects to:
        https://login.live.com/oauth20_desktop.srf?access_token=***&token_type=Bearer&expires_in=xxx
        or
        https://login.live.com/oauth20_authorize.srf#access_token=xxx...
        
        Uses navigation event listener to capture tokens that may flash briefly
        in intermediate redirects before the final page loads.
        """
        # --- Navigation listener: capture ALL URLs visited during this phase ---
        # Tokens can appear in intermediate redirects that get overridden by the
        # final page. We need to capture them as they pass through.
        captured_urls = []
        token_found = {'result': None}
        
        def on_navigate(frame):
            url = frame.url
            captured_urls.append(url)
            # Check if this URL contains a token
            if 'access_token' in url:
                parsed = urlparse(url)
                params = parse_qs(parsed.query)
                if 'access_token' in params:
                    token_found['result'] = {
                        'access_token': params['access_token'][0],
                        'token_type': params.get('token_type', ['Bearer'])[0],
                        'expires_in': params.get('expires_in', ['3600'])[0],
                        'refresh_token': params.get('refresh_token', [''])[0],
                    }
                    logger.info("  Token captured from navigation event (query)")
                frag_params = parse_qs(parsed.fragment)
                if 'access_token' in frag_params:
                    token_found['result'] = {
                        'access_token': frag_params['access_token'][0],
                        'token_type': frag_params.get('token_type', ['Bearer'])[0],
                        'expires_in': frag_params.get('expires_in', ['3600'])[0],
                        'refresh_token': frag_params.get('refresh_token', [''])[0],
                    }
                    logger.info("  Token captured from navigation event (fragment)")
        
        page.on("frame navigated", on_navigate)
        
        try:
            # Wait for token-bearing redirect with generous timeout
            page.wait_for_function("""() => {
                const url = window.location.href;
                return url.includes('access_token') || 
                       url.includes('oauth20_desktop.srf');
            }""", timeout=45000)
        except Exception:
            # Timeout is OK — we may have already captured the token via navigation events
            pass
        finally:
            # Remove the listener before further processing
            page.remove_listener("frame navigated", on_navigate)
        
        # --- Check if navigation listener already found the token ---
        if token_found['result']:
            logger.info("  Using token captured from navigation event")
            return token_found['result']
        
        # --- Take a snapshot of current state ---
        url = page.url
        content = page.content()
        body_text = page.inner_text('body')
        logger.info(f"  Final URL: {url[:300]}")
        logger.info(f"  Body preview: {body_text[:200]}")
        logger.info(f"  Captured {len(captured_urls)} navigation events")
        
        # Check for known failure states early
        failure_indicators = [
            'ppsecure', 'proofConfirmation', 'consent',
            "Can't sign in", 'Try again later', 'security challenge',
            'Unusual sign-in', 'We need more information',
        ]
        for indicator in failure_indicators:
            if indicator in url.lower() or indicator in body_text.lower():
                logger.warning(f"  Detected failure state: '{indicator}' in URL/body")
                logger.warning(f"  Navigations seen: {captured_urls[-5:]}")
                break
        
        # Recovery: if stuck on FIDO/biometric page, try to switch away
        if 'fido' in url.lower() or 'face' in body_text.lower() or 'fingerprint' in body_text.lower() or 'security key' in body_text.lower():
            logger.warning(f"  Stuck on FIDO/biometric page during token extraction, attempting recovery")
            self._switch_from_biometric_to_password(page)
            self._wait_stable(page, 5000)
            # Re-capture url/content after recovery
            url = page.url
            content = page.content()
            body_text = page.inner_text('body')
            logger.info(f"  After recovery URL: {url[:300]}")
        
        # --- Method 1: Parse access_token from query params ---
        parsed = urlparse(url)
        params = parse_qs(parsed.query)
        
        if 'access_token' in params:
            logger.info("  Found token in URL query params")
            return {
                'access_token': params['access_token'][0],
                'token_type': params.get('token_type', ['Bearer'])[0],
                'expires_in': params.get('expires_in', ['3600'])[0],
                'refresh_token': params.get('refresh_token', [''])[0],
            }
        
        # --- Method 2: Try fragment (hash) params ---
        fragment = parsed.fragment
        if fragment:
            frag_params = parse_qs(fragment)
            if 'access_token' in frag_params:
                logger.info("  Found token in URL fragment")
                return {
                    'access_token': frag_params['access_token'][0],
                    'token_type': frag_params.get('token_type', ['Bearer'])[0],
                    'expires_in': frag_params.get('expires_in', ['3600'])[0],
                    'refresh_token': frag_params.get('refresh_token', [''])[0],
                }
        
        # --- Method 3: Extract from page HTML content (token embedded in JS) ---
        match = re.search(r'access_token=([^&"\s]+)', content)
        if match:
            logger.info("  Found token in page HTML")
            return {
                'access_token': match.group(1),
                'refresh_token': '',
            }
        
        # --- Method 4: Look for token in localStorage or sessionStorage ---
        try:
            stored_token = page.evaluate("""() => {
                for (let i = 0; i < localStorage.length; i++) {
                    const key = localStorage.key(i);
                    const val = localStorage.getItem(key);
                    if (val && val.includes('access_token')) return val;
                }
                for (let i = 0; i < sessionStorage.length; i++) {
                    const key = sessionStorage.key(i);
                    const val = sessionStorage.getItem(key);
                    if (val && val.includes('access_token')) return val;
                }
                return null;
            }""")
            if stored_token:
                logger.info("  Found token in browser storage")
                if isinstance(stored_token, str) and '=' in stored_token:
                    token_params = parse_qs(stored_token.split('#', 1)[-1] if '#' in stored_token else stored_token)
                    if 'access_token' in token_params:
                        return {
                            'access_token': token_params['access_token'][0],
                            'refresh_token': token_params.get('refresh_token', [''])[0],
                        }
                try:
                    stored_json = json.loads(stored_token)
                    if 'access_token' in stored_json:
                        return {
                            'access_token': stored_json['access_token'],
                            'refresh_token': stored_json.get('refresh_token', ''),
                        }
                except:
                    pass
        except Exception as e:
            logger.info(f"  Could not check browser storage: {e}")
        
        # --- Method 5: Wait more and retry (sometimes redirect is delayed) ---
        logger.info("  Token not found yet, waiting longer...")
        self._wait_stable(page, 10000)
        url = page.url
        content = page.content()
        logger.info(f"  Retry URL: {url[:300]}")
        
        parsed = urlparse(url)
        params = parse_qs(parsed.query)
        if 'access_token' in params:
            logger.info("  Found token on retry (query params)")
            return {
                'access_token': params['access_token'][0],
                'refresh_token': params.get('refresh_token', [''])[0],
            }
        
        fragment = parsed.fragment
        if fragment:
            frag_params = parse_qs(fragment)
            if 'access_token' in frag_params:
                logger.info("  Found token on retry (fragment)")
                return {
                    'access_token': frag_params['access_token'][0],
                    'refresh_token': frag_params.get('refresh_token', [''])[0],
                }
        
        match = re.search(r'access_token=([^&"\s]+)', content)
        if match:
            logger.info("  Found token on retry (HTML)")
            return {'access_token': match.group(1), 'refresh_token': ''}
        
        # --- Method 6: Check all captured navigation URLs for a token ---
        for nav_url in captured_urls:
            if 'access_token' in nav_url:
                parsed = urlparse(nav_url)
                params = parse_qs(parsed.query)
                if 'access_token' in params:
                    logger.info("  Found token in captured navigation history")
                    return {
                        'access_token': params['access_token'][0],
                        'token_type': params.get('token_type', ['Bearer'])[0],
                        'expires_in': params.get('expires_in', ['3600'])[0],
                        'refresh_token': params.get('refresh_token', [''])[0],
                    }
                frag_params = parse_qs(parsed.fragment)
                if 'access_token' in frag_params:
                    logger.info("  Found token in captured navigation history (fragment)")
                    return {
                        'access_token': frag_params['access_token'][0],
                        'refresh_token': frag_params.get('refresh_token', [''])[0],
                    }
        
        # --- Method 7: Search page content for JWT tokens (LAST RESORT only) ---
        # Microsoft OAuth access tokens are JWTs (eyJ...). BUT the login pages
        # also contain many other JWTs (CSRF, FIDO, internal state) that are NOT
        # valid for XBL auth. Only use this if we're reasonably confident we're
        # on the token redirect page, not the login page.
        jwt_matches = re.findall(r'(eyJ[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,}\.[A-Za-z0-9_-]{10,})', content)
        if jwt_matches:
            logger.info(f"  Found {len(jwt_matches)} JWT(s) in page content")
            # Only consider JWT fallback if we're NOT on a login/FIDO page
            is_login_page = any(indicator in url.lower() for indicator in [
                'login.microsoft.com', 'login.live.com', 'fido', 'ppsecure',
                'proofConfirmation', 'consent', 'xboxlive', 'account.microsoft.com',
            ])
            if is_login_page:
                logger.warning(f"  Found JWTs but still on a login/intermediate page ({url[:200]}), skipping JWT fallback")
                logger.warning(f"  JWTs on login pages are almost certainly NOT valid OAuth tokens")
                # Don't return these — let it fail so we get diagnostic info
            else:
                # We're on some other page — try to find a valid access token JWT
                for jwt in jwt_matches:
                    try:
                        header_b64 = jwt.split('.')[0]
                        header_b64 += '=' * (4 - len(header_b64) % 4)
                        header = json.loads(__import__('base64').b64decode(header_b64))
                        # Microsoft access tokens have specific audience claims
                        payload_b64 = jwt.split('.')[1]
                        payload_b64 += '=' * (4 - len(payload_b64) % 4)
                        payload = json.loads(__import__('base64').b64decode(payload_b64))
                        aud = payload.get('aud', '')
                        typ = header.get('typ', '')
                        logger.info(f"    JWT header.typ={typ}, payload.aud={aud}")
                        # Valid Microsoft OAuth tokens target Xbox Live or Microsoft services
                        if any(k in str(aud) for k in ['xbox', 'xbl', 'login.live.com', 'consumers', '00000000402b5328']):
                            logger.info("  Using JWT from page content as access token (valid audience)")
                            return {'access_token': jwt, 'refresh_token': ''}
                    except:
                        pass
                # If we found JWTs but none had valid audience, use the first one as last resort
                logger.warning("  Using first JWT as fallback access token (no valid audience found)")
                logger.info(f"  Fallback JWT length: {len(jwt_matches[0])}")
                return {'access_token': jwt_matches[0], 'refresh_token': ''}
        
        # --- Final diagnostic dump ---
        logger.error(f"  Could not find token after all methods")
        logger.error(f"  Current URL: {url}")
        logger.error(f"  Page title: {page.title()}")
        logger.error(f"  Body text (500 chars): {body_text[:500]}")
        logger.error(f"  Navigation history (last 10): {captured_urls[-10:]}")
        
        # Take a screenshot for debugging
        try:
            screenshot_path = str(self.cache_dir / f"auth_debug_{time.time()}.png")
            page.screenshot(path=screenshot_path)
            logger.error(f"  Debug screenshot saved: {screenshot_path}")
        except Exception as e:
            logger.error(f"  Could not save screenshot: {e}")
        
        return None
    
    # -------------------------------------------------------------------------
    # Xbox Live token exchange (pure HTTP, no browser needed)
    # -------------------------------------------------------------------------
    
    def _authenticate_xbl(self, msa_token: str) -> Optional[dict]:
        """Exchange Microsoft token for Xbox Live token."""
        import requests
        
        headers = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
        }
        
        payload = {
            "Properties": {
                "AuthMethod": "RPS",
                "SiteName": "user.auth.xboxlive.com",
                "RpsTicket": f"d={msa_token}"
            },
            "RelyingParty": "http://auth.xboxlive.com",
            "TokenType": "JWT"
        }
        
        try:
            resp = requests.post(XBL_AUTH_URL, json=payload, headers=headers, timeout=30)
            logger.info(f"  XBL response status: {resp.status_code}")
            logger.info(f"  XBL response body: {resp.text[:500]}")
            resp.raise_for_status()
            result = resp.json()
            # XBL success response should have 'Token' field
            if 'Token' not in result:
                logger.error(f"  XBL response missing 'Token' field: {result}")
                return None
            return result
        except Exception as e:
            logger.error(f"XBL auth failed: {e}")
            return None
    
    def _authenticate_xsts(self, xbl_result: dict) -> Optional[dict]:
        """Exchange Xbox Live token for XSTS token."""
        import requests
        
        headers = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
        }
        
        payload = {
            "Properties": {
                "UserTokens": [xbl_result['Token']]
            },
            "RelyingParty": "rp://api.minecraftservices.com/",
            "TokenType": "JWT"
        }
        
        try:
            resp = requests.post(XSTS_AUTH_URL, json=payload, headers=headers, timeout=30)
            logger.info(f"  XSTS response status: {resp.status_code}")
            if resp.status_code != 200:
                logger.error(f"  XSTS response body: {resp.text[:500]}")
            resp.raise_for_status()
            result = resp.json()
            if 'Token' not in result:
                logger.error(f"  XSTS response missing 'Token' field: {result}")
                return None
            return result
        except Exception as e:
            logger.error(f"XSTS auth failed: {e}")
            return None
    
    def _authenticate_minecraft(self, xsts_result: dict) -> Optional[dict]:
        """Exchange XSTS token for Minecraft access token."""
        import requests
        
        headers = {
            'Accept': 'application/json',
            'Content-Type': 'application/json',
        }
        
        # XSTS DisplayClaims may contain the uhs (unique auth string)
        uhs = xsts_result.get('uhs', '')
        if not uhs and 'DisplayClaims' in xsts_result:
            # Decode the base64 DisplayClaims to find uhs
            try:
                import base64
                claims_b64 = xsts_result['DisplayClaims']
                # Add padding if needed
                claims_b64 += '=' * (4 - len(claims_b64) % 4)
                claims = json.loads(base64.b64decode(claims_b64))
                uhs = claims.get('xui', [{}])[0].get('uhs', '')
            except Exception as e:
                logger.warning(f"  Could not decode XSTS DisplayClaims for uhs: {e}")
        
        payload = {
            "identityToken": f"XBL3.0 x={uhs},{xsts_result['Token']}"
        }
        
        try:
            resp = requests.post(MC_XBOX_TOKEN_URL, json=payload, headers=headers, timeout=30)
            logger.info(f"  MC auth response status: {resp.status_code}")
            if resp.status_code != 200:
                logger.error(f"  MC auth response body: {resp.text[:500]}")
            resp.raise_for_status()
            result = resp.json()
            if 'access_token' not in result:
                logger.error(f"  MC auth response missing 'access_token': {result}")
                return None
            return result
        except Exception as e:
            logger.error(f"Minecraft auth failed: {e}")
            return None
    
    def _get_minecraft_username(self, mc_token: str) -> str:
        """Get Minecraft username from access token."""
        import requests
        
        try:
            resp = requests.get(
                'https://api.minecraftservices.com/minecraft/profile',
                headers={'Authorization': f'Bearer {mc_token}'}
            )
            if resp.status_code == 200:
                data = resp.json()
                return data.get('name', 'Unknown')
        except Exception as e:
            logger.error(f"Failed to get username: {e}")
        return 'Unknown'
    
    # -------------------------------------------------------------------------
    # Token caching
    # -------------------------------------------------------------------------
    
    def _cache_filename(self, email: str) -> str:
        """Get cache file path for an email."""
        safe_name = re.sub(r'[^a-zA-Z0-9._-]', '_', email)
        return str(self.cache_dir / f"{safe_name}.json")
    
    def _cache_result(self, result: AuthResult):
        """Save auth result to cache file."""
        path = self._cache_filename(result.email)
        data = result.to_dict()
        data['cached_at'] = time.time()
        with open(path, 'w') as f:
            json.dump(data, f, indent=2)
    
    def get_cached_token(self, email: str) -> Optional[str]:
        """
        Get cached Minecraft access token for an email.
        Returns None if not cached or expired.
        """
        path = self._cache_filename(email)
        if not os.path.exists(path):
            return None
        
        with open(path) as f:
            data = json.load(f)
        
        # Check if token is still valid (with 5-minute buffer)
        cached_at = data.get('cached_at', 0)
        expires_in = int(data.get('mc_token_expiry', 3600))
        if time.time() - cached_at < (expires_in - 300):
            return data.get('mc_access_token')
        
        # Try refresh
        refresh_token = data.get('msa_refresh_token', '')
        if refresh_token:
            new_token = self._refresh_token(refresh_token)
            if new_token:
                return new_token
        
        return None
    
    def _refresh_token(self, refresh_token: str) -> Optional[str]:
        """Refresh a Microsoft OAuth token."""
        import requests
        
        try:
            resp = requests.post(
                'https://login.live.com/oauth20_token.srf',
                data={
                    'grant_type': 'refresh_token',
                    'client_id': XBOX_CLIENT_ID,
                    'refresh_token': refresh_token,
                    'scope': XBOX_SCOPE,
                }
            )
            if resp.status_code == 200:
                data = resp.json()
                return data.get('access_token')
        except Exception as e:
            logger.error(f"Token refresh failed: {e}")
        return None
    
    def list_cached_accounts(self) -> List[dict]:
        """List all cached accounts."""
        accounts = []
        for f in self.cache_dir.glob('*.json'):
            with open(f) as fp:
                data = json.load(fp)
                accounts.append({
                    'email': data.get('email', f.stem),
                    'username': data.get('username', ''),
                    'cached_at': data.get('cached_at', 0),
                    'path': str(f),
                })
        return accounts
    
    # -------------------------------------------------------------------------
    # Bulk authentication
    # -------------------------------------------------------------------------
    
    def authenticate_bulk(self, accounts_file: str, 
                          mfa_handler=None) -> List[AuthResult]:
        """
        Authenticate multiple accounts from a file.
        
        File format: email:password per line, or email,password per line.
        Lines starting with # are comments.
        """
        accounts = self._parse_accounts_file(accounts_file)
        logger.info(f"Authenticating {len(accounts)} accounts from {accounts_file}")
        
        results = []
        for i, (email, password) in enumerate(accounts):
            logger.info(f"  [{i+1}/{len(accounts)}] {email}")
            result = self.authenticate(email, password, mfa_handler)
            results.append(result)
            
            # Small delay between accounts to avoid rate limiting
            if i < len(accounts) - 1:
                time.sleep(1)
        
        # Print summary
        success = sum(1 for r in results if r.success)
        failed = len(results) - success
        logger.info(f"  Bulk auth complete: {success} success, {failed} failed")
        
        return results
    
    def _parse_accounts_file(self, path: str) -> List[Tuple[str, str]]:
        """Parse accounts file into list of (email, password) tuples."""
        accounts = []
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                # Try email:password first, then email,password
                if ':' in line:
                    parts = line.split(':', 1)
                elif ',' in line:
                    parts = line.split(',', 1)
                else:
                    continue
                if len(parts) == 2:
                    accounts.append((parts[0].strip(), parts[1].strip()))
        return accounts
    
    # -------------------------------------------------------------------------
    # Utility
    # -------------------------------------------------------------------------
    
    @staticmethod
    def _wait_stable(page: Page, ms: int = 1000):
        """Wait for page to stabilize."""
        page.wait_for_timeout(ms)


# ===========================================================================
# CLI interface
# ===========================================================================

def main():
    """CLI for bulk authentication."""
    import argparse
    
    parser = argparse.ArgumentParser(description='Microsoft Account Authenticator for Minecraft')
    parser.add_argument('action', choices=['auth', 'bulk', 'list', 'get'],
                       help='Action to perform')
    parser.add_argument('--email', help='Email address')
    parser.add_argument('--password', help='Password')
    parser.add_argument('--file', help='Accounts file (email:password per line)')
    parser.add_argument('--cache-dir', default='~/.ms_auth_cache', help='Cache directory')
    parser.add_argument('--visible', action='store_true', help='Run browser in visible mode')
    
    args = parser.parse_args()
    
    auth = MicrosoftAuthenticator(
        cache_dir=args.cache_dir,
        headless=not args.visible
    )
    
    try:
        if args.action == 'auth':
            if not args.email or not args.password:
                print("Error: --email and --password required for auth")
                return
            result = auth.authenticate(args.email, args.password)
            if result.success:
                print(json.dumps(result.to_dict(), indent=2))
            else:
                print(f"Failed: {result.error}")
        
        elif args.action == 'bulk':
            if not args.file:
                print("Error: --file required for bulk auth")
                return
            results = auth.authenticate_bulk(args.file)
            for r in results:
                status = "✓" if r.success else "✗"
                print(f"  {status} {r.email} -> {r.username} {'(' + r.error + ')' if not r.success else ''}")
        
        elif args.action == 'list':
            accounts = auth.list_cached_accounts()
            if accounts:
                for acc in accounts:
                    print(f"  {acc['email']} -> {acc['username']} (cached: {time.ctime(acc['cached_at'])})")
            else:
                print("No cached accounts")
        
        elif args.action == 'get':
            if not args.email:
                print("Error: --email required for get")
                return
            token = auth.get_cached_token(args.email)
            if token:
                print(token)
            else:
                print("No valid cached token")
    
    finally:
        auth.close()


if __name__ == '__main__':
    main()
