#!/usr/bin/env python3
"""
Automated muse.ai sign-in with in-browser OTP approval.

Usage:
  signin.py --email <email> [--otp <code>]
  
If --otp is not provided, the script will:
1. Navigate through login flow up to OTP prompt
2. Print "APPROVAL_NEEDED: OTP required for <email>"
3. Exit with code 2

The operator then obtains the OTP (via chat with user) and re-runs:
  signin.py --email <email> --otp <code>

This keeps credentials out of the automation — the OTP is provided
transiently via the operator, never stored.
"""
import json, urllib.request, websocket, time, sys, argparse

CDP_URL = "http://127.0.0.1:9410/json/list"

def get_page():
    with urllib.request.urlopen(CDP_URL, timeout=5) as r:
        ts = json.load(r)
    pages = [t for t in ts if t.get('type') == 'page']
    if not pages:
        print("ERROR: No page found", file=sys.stderr)
        sys.exit(1)
    return pages[0]

def ev(ws, expr, await_p=False):
    ws.send(json.dumps({
        "id": 1, "method": "Runtime.evaluate",
        "params": {"expression": expr, "returnByValue": True, "awaitPromise": await_p}
    }))
    resp = json.loads(ws.recv())
    return resp.get('result', {}).get('result', {}).get('value')

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--email', required=True)
    p.add_argument('--otp', default=None)
    args = p.parse_args()

    page = get_page()
    ws = websocket.create_connection(page['webSocketDebuggerUrl'], timeout=15)

    # Step 1: Check if already logged in
    title = ev(ws, "document.title")
    body = ev(ws, "document.body.innerText.slice(0,200)")
    if "Connected" in body or "Chats" in body:
        print(f"Already logged in (title: {title})")
        ws.close()
        return 0

    # Step 2: Click Log in (if on landing page)
    if "Log in" in body:
        print("Clicking Log in...")
        ev(ws, """(async()=>{
          const b=[...document.querySelectorAll('button')].find(x=>x.innerText.includes('Log in'));
          if(b) b.click(); return !!b;
        })()""", True)
        time.sleep(3)

    # Step 3: Enter email
    print(f"Entering email: {args.email}")
    result = ev(ws, f"""(async()=>{{
      const inp=[...document.querySelectorAll('input')].find(i=>
        (i.placeholder&&i.placeholder.toLowerCase().includes('email'))||
        (i.getAttribute('aria-label')&&i.getAttribute('aria-label').toLowerCase().includes('email'))
      );
      if(!inp) return 'NOINPUT';
      inp.focus();
      document.execCommand('insertText',false,'{args.email}');
      await new Promise(r=>setTimeout(r,500));
      return 'entered:'+inp.value;
    }})()""", True)
    print(f"Email: {result}")
    if result == 'NOINPUT':
        print("ERROR: Email input not found", file=sys.stderr)
        ws.close()
        sys.exit(1)
    time.sleep(1)

    # Step 4: Click Continue
    print("Clicking Continue...")
    ev(ws, """(async()=>{
      const b=[...document.querySelectorAll('button')].find(x=>x.innerText.includes('Continue'));
      if(b) b.click(); return !!b;
    })()""", True)
    time.sleep(4)

    # Step 5: Check for OTP prompt
    body = ev(ws, "document.body.innerText.slice(0,300)")
    if "Enter your code" in body or "code we sent" in body:
        print(f"OTP prompt detected for {args.email}")
        if not args.otp:
            print(f"APPROVAL_NEEDED: OTP required for {args.email}")
            print("Re-run with: signin.py --email {} --otp <code>".format(args.email))
            ws.close()
            sys.exit(2)  # Approval needed
        
        # Enter OTP
        print("Entering OTP...")
        result = ev(ws, f"""(async()=>{{
          const inp=[...document.querySelectorAll('input')].find(i=>i.type==='text');
          if(!inp) return 'NOINPUT';
          inp.focus();
          document.execCommand('insertText',false,'{args.otp}');
          await new Promise(r=>setTimeout(r,500));
          return 'entered:'+inp.value;
        }})()""", True)
        print(f"OTP: {result}")
        time.sleep(1)
        
        # Click Next/Verify
        print("Submitting OTP...")
        ev(ws, """(async()=>{
          const b=[...document.querySelectorAll('button')].find(x=>
            x.innerText.includes('Next')||x.innerText.includes('Verify')
          );
          if(b) b.click(); return !!b;
        })()""", True)
        time.sleep(5)
        
        # Verify login
        body = ev(ws, "document.body.innerText.slice(0,200)")
        title = ev(ws, "document.title")
        if "Connected" in body or "Chats" in body:
            print(f"SUCCESS: Logged in as {args.email} (title: {title})")
            ws.close()
            return 0
        else:
            print(f"WARNING: Login may have failed. Title: {title}", file=sys.stderr)
            print(f"Body: {body[:100]}", file=sys.stderr)
            ws.close()
            sys.exit(1)
    else:
        print("No OTP prompt found - check page state", file=sys.stderr)
        print(f"Body: {body[:200]}", file=sys.stderr)
        ws.close()
        sys.exit(1)

if __name__ == '__main__':
    sys.exit(main())