"""
Genesis Browserless — Initial Profile Setup Script
====================================================
Implements the Collaborative Browser Handoff protocol:

1. Opens Browserless Chrome to a service login page
2. Waits for YOU (Kinan) to log in manually (bypasses 2FA, Cloudflare, etc.)
3. Captures all cookies after login
4. Saves cookies to Elestio Redis for permanent persistence
5. Marks the session as ready for agent use

Run this script ONCE after deploying Browserless.
After that, agents use BrowserlessSessionManager.launch_authenticated_session()
and get pre-authenticated sessions automatically.

Usage:
    python scripts/setup_browserless_profiles.py

    # Set up only specific accounts/services:
    python scripts/setup_browserless_profiles.py --account kinan@agileadapt.com --service gemini.google.com

    # Refresh a specific session (re-login):
    python scripts/setup_browserless_profiles.py --refresh --service portal.telnyx.com

    # Check current session status without doing any logins:
    python scripts/setup_browserless_profiles.py --status

Environment variables:
    BROWSERLESS_URL   - Browserless instance URL (default: http://localhost:3000)
    BROWSERLESS_TOKEN - Auth token for Browserless
"""

import argparse
import json
import logging
import os
import sys
import time

import requests

# Add genesis-system root to path
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, ROOT)
sys.path.insert(0, os.path.join(ROOT, "data", "genesis-memory"))

from core.browserless_session_manager import (
    ACCOUNT_SERVICES,
    BrowserlessSessionManager,
)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Setup helper
# ---------------------------------------------------------------------------

class ProfileSetup:
    """Orchestrates the Collaborative Browser Handoff login flow."""

    def __init__(self, manager: BrowserlessSessionManager):
        self.manager = manager
        self.browserless_url = manager.browserless_url
        self.token = manager.browserless_token

    def print_banner(self) -> None:
        print()
        print("=" * 60)
        print("  Genesis Browserless — Profile Setup")
        print("  Collaborative Browser Handoff Protocol")
        print("=" * 60)
        print()
        print("This script will open a Chrome browser for each service.")
        print("YOU log in manually. The script captures and saves cookies.")
        print("Agents will use those cookies for all future sessions.")
        print()

    def check_browserless_connection(self) -> bool:
        """Verify Browserless is reachable before starting."""
        print(f"Connecting to Browserless at: {self.browserless_url}")
        try:
            resp = requests.get(
                f"{self.browserless_url}/health",
                headers={"Authorization": f"Bearer {self.token}"},
                timeout=10,
            )
            if resp.status_code == 200:
                data = resp.json()
                print(f"  Browserless online — max sessions: {data.get('maxConcurrent', 'N/A')}")
                return True
            else:
                print(f"  ERROR: Browserless returned HTTP {resp.status_code}")
                return False
        except requests.RequestException as exc:
            print(f"  ERROR: Cannot reach Browserless: {exc}")
            print()
            print("  Troubleshooting:")
            print("  1. Is the Docker container running?")
            print("     cd infra/browserless && docker compose up -d")
            print("  2. Is BROWSERLESS_URL set correctly?")
            print(f"     Current: {self.browserless_url}")
            print("  3. Is the token correct?")
            return False

    def print_status_table(self) -> None:
        """Print a table of all session states."""
        print()
        print(f"{'ACCOUNT':<30} {'SERVICE':<25} {'STATUS':<15} {'COOKIES':<10} {'AGE':<10}")
        print("-" * 90)

        for account, data in ACCOUNT_SERVICES.items():
            for service, svc_config in data["services"].items():
                state = self.manager.check_login_state(account, service)
                status = "LOGGED IN" if state["logged_in"] else "NOT SET"
                if state["logged_in"] and state["needs_refresh"]:
                    status = "STALE"
                cookies = str(state["cookie_count"]) if state["cookie_count"] > 0 else "-"
                age = f"{state['age_hours']:.1f}h" if state["age_hours"] > 0 else "-"
                print(f"  {account:<28} {service:<25} {status:<15} {cookies:<10} {age:<10}")

        print()

    def setup_service(self, account: str, service: str, service_config: dict) -> bool:
        """
        Run the Collaborative Browser Handoff for a single account+service.

        Opens a visible Chrome window, waits for manual login, captures cookies.
        """
        print()
        print(f"  Setting up: {account} @ {service}")
        print(f"  Service: {service_config.get('description', service)}")
        print(f"  Login URL: {service_config['login_url']}")
        print()

        # Check if already logged in
        state = self.manager.check_login_state(account, service)
        if state["logged_in"] and not state["needs_refresh"]:
            print(f"  Already logged in ({state['cookie_count']} cookies, {state['age_hours']:.1f}h old)")
            answer = input("  Re-login anyway? [y/N]: ").strip().lower()
            if answer != "y":
                print("  Skipping.")
                return True

        print("  Opening Chrome browser...")
        print("  PLEASE LOG IN when the browser opens.")
        print("  The script will wait up to 5 minutes for you to complete login.")
        print()

        # Launch a VISIBLE browser (headless=False) and navigate to login URL
        session_id = f"setup_{account.replace('@', '_').replace('.', '_')}_{int(time.time())}"

        # Use Browserless /session endpoint to get a live browser
        launch_payload = {
            "launch": {
                "headless": False,  # Visible window — Kinan can log in
                "args": [
                    "--no-sandbox",
                    "--disable-dev-shm-usage",
                    "--disable-blink-features=AutomationControlled",
                    "--start-maximized",
                ],
            }
        }

        try:
            # Open the browser and navigate to login page
            open_resp = requests.post(
                f"{self.browserless_url}/function",
                headers={
                    "Authorization": f"Bearer {self.token}",
                    "Content-Type": "application/json",
                },
                json={
                    "code": self._build_open_and_wait_script(
                        service_config["login_url"],
                        service_config.get("auth_cookie_keys", []),
                        wait_seconds=300,  # 5 min window
                    ),
                    "launch": launch_payload["launch"],
                },
                timeout=360,  # Allow up to 6 minutes total
            )

            if open_resp.status_code != 200:
                print(f"  ERROR: Browserless /function returned {open_resp.status_code}")
                print(f"  Response: {open_resp.text[:200]}")
                return self._fallback_cookie_capture(account, service, service_config)

            result = open_resp.json()
            cookies = result.get("cookies", [])

            if cookies:
                # Save to Redis
                saved = self.manager.save_cookies(account, service, cookies)
                if saved:
                    print(f"  SUCCESS: Saved {len(cookies)} cookies for {account} @ {service}")
                    return True
                else:
                    print("  ERROR: Failed to save cookies to Redis")
                    return False
            else:
                print("  No cookies captured — login may not have completed")
                return self._fallback_cookie_capture(account, service, service_config)

        except requests.RequestException as exc:
            print(f"  ERROR contacting Browserless: {exc}")
            return self._fallback_cookie_capture(account, service, service_config)

    def _build_open_and_wait_script(
        self,
        login_url: str,
        auth_cookie_keys: list,
        wait_seconds: int = 300,
    ) -> str:
        """
        JS script that opens login page, waits for auth cookies to appear,
        then returns all cookies.
        """
        login_url_json = json.dumps(login_url)
        auth_keys_json = json.dumps(auth_cookie_keys)
        wait_ms = wait_seconds * 1000

        return f"""
module.exports = async ({{ page }}) => {{
    const loginUrl = {login_url_json};
    const authCookieKeys = {auth_keys_json};
    const waitMs = {wait_ms};
    const pollIntervalMs = 3000;

    // Navigate to login page
    await page.goto(loginUrl, {{ waitUntil: 'domcontentloaded', timeout: 30000 }});

    console.log('Browser opened at: ' + loginUrl);
    console.log('Waiting for login... (max ' + (waitMs / 60000).toFixed(0) + ' minutes)');

    // Poll for auth cookies
    const startTime = Date.now();
    let isLoggedIn = false;
    let lastCookieCount = 0;

    while (Date.now() - startTime < waitMs) {{
        await new Promise(r => setTimeout(r, pollIntervalMs));

        const cookies = await page.cookies();
        const cookieNames = cookies.map(c => c.name);

        // Check if auth cookies are present
        if (authCookieKeys.length > 0) {{
            isLoggedIn = authCookieKeys.some(key => cookieNames.includes(key));
        }} else {{
            // No auth keys defined — assume logged in if cookie count grew
            isLoggedIn = cookies.length > lastCookieCount && cookies.length > 5;
        }}

        if (cookies.length !== lastCookieCount) {{
            console.log('Cookies changed: ' + cookies.length + ' total');
            lastCookieCount = cookies.length;
        }}

        if (isLoggedIn) {{
            console.log('Auth cookies detected — login successful!');
            // Wait a moment for any post-login redirects to settle
            await new Promise(r => setTimeout(r, 2000));

            const allCookies = await page.cookies();
            const title = await page.title();
            const url = page.url();

            return {{
                success: true,
                cookies: allCookies,
                cookie_count: allCookies.length,
                page_title: title,
                page_url: url,
                elapsed_seconds: Math.round((Date.now() - startTime) / 1000),
            }};
        }}
    }}

    // Timeout — return whatever cookies we have
    const allCookies = await page.cookies();
    return {{
        success: false,
        cookies: allCookies,
        cookie_count: allCookies.length,
        message: 'Timeout waiting for login — returning current cookies',
    }};
}};
"""

    def _fallback_cookie_capture(
        self, account: str, service: str, service_config: dict
    ) -> bool:
        """
        Manual fallback: ask user to paste cookies as JSON.
        Used when Browserless cannot open a visible window.
        """
        print()
        print("  FALLBACK: Manual cookie capture")
        print("  --------------------------------")
        print("  1. Open a browser manually and log into:")
        print(f"     {service_config['login_url']}")
        print("  2. Open DevTools (F12) → Application → Cookies")
        print(f"     Select the {service} domain")
        print("  3. Copy all cookies as JSON")
        print("     (In Chrome: DevTools → Console → paste this and copy the output)")
        print()
        print("     JSON.stringify([...document.cookie.split(';').map(c => {")
        print("       const [name, ...v] = c.trim().split('=');")
        print("       return {name: name.trim(), value: v.join('='), domain: location.hostname};")
        print("     })])")
        print()
        print("  Then paste the JSON here (or press Enter to skip):")
        print("  > ", end="", flush=True)

        try:
            raw = sys.stdin.readline().strip()
        except (KeyboardInterrupt, EOFError):
            print("\n  Skipped.")
            return False

        if not raw:
            print("  Skipped.")
            return False

        try:
            cookies = json.loads(raw)
            if not isinstance(cookies, list):
                print("  ERROR: Expected a JSON array of cookies")
                return False

            saved = self.manager.save_cookies(account, service, cookies)
            if saved:
                print(f"  Saved {len(cookies)} cookies for {account} @ {service}")
                return True
            else:
                print("  ERROR: Failed to save cookies to Redis")
                return False

        except json.JSONDecodeError as exc:
            print(f"  ERROR: Invalid JSON: {exc}")
            return False


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Genesis Browserless — Initial Profile Setup",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Set up all accounts and services
  python scripts/setup_browserless_profiles.py

  # Set up only Gemini
  python scripts/setup_browserless_profiles.py --service gemini.google.com

  # Set up all services for one account
  python scripts/setup_browserless_profiles.py --account kinan@agileadapt.com

  # Check status only
  python scripts/setup_browserless_profiles.py --status

  # Force refresh a stale session
  python scripts/setup_browserless_profiles.py --refresh --service portal.telnyx.com
        """,
    )
    parser.add_argument(
        "--account",
        help="Only set up this account (e.g., kinan@agileadapt.com)",
    )
    parser.add_argument(
        "--service",
        help="Only set up this service (e.g., gemini.google.com)",
    )
    parser.add_argument(
        "--status",
        action="store_true",
        help="Show session status and exit (no logins)",
    )
    parser.add_argument(
        "--refresh",
        action="store_true",
        help="Force re-login even for already-valid sessions",
    )
    parser.add_argument(
        "--browserless-url",
        default=os.environ.get("BROWSERLESS_URL", "http://localhost:3000"),
        help="Browserless URL (default: http://localhost:3000)",
    )
    parser.add_argument(
        "--token",
        default=os.environ.get("BROWSERLESS_TOKEN", ""),
        help="Browserless auth token",
    )

    args = parser.parse_args()

    manager = BrowserlessSessionManager(
        browserless_url=args.browserless_url,
        browserless_token=args.token,
    )
    setup = ProfileSetup(manager)

    setup.print_banner()

    # Status-only mode
    if args.status:
        print("Current Session Status:")
        setup.print_status_table()
        health = manager.health_check()
        print(f"Browserless: {'ONLINE' if health['browserless_online'] else 'OFFLINE'}")
        print(f"Redis: {'ONLINE' if health['redis_online'] else 'OFFLINE'}")
        print(f"Persisted sessions: {health['persisted_sessions']}")
        return

    # Check Browserless connection
    if not setup.check_browserless_connection():
        print()
        print("ERROR: Cannot proceed without a Browserless connection.")
        print("Start the container first:")
        print("  cd E:/genesis-system/infra/browserless")
        print("  docker compose up -d")
        sys.exit(1)

    print()
    print("Current Status Before Setup:")
    setup.print_status_table()

    # Determine what to set up
    accounts_to_process = {}

    for account, data in ACCOUNT_SERVICES.items():
        # Filter by account if specified
        if args.account and account != args.account:
            continue

        services_to_process = {}
        for service, svc_config in data["services"].items():
            # Filter by service if specified
            if args.service and service != args.service:
                continue

            # Skip already-valid sessions unless --refresh
            if not args.refresh:
                state = manager.check_login_state(account, service)
                if state["logged_in"] and not state["needs_refresh"]:
                    logger.info("Skipping %s @ %s — already logged in", account, service)
                    continue

            services_to_process[service] = svc_config

        if services_to_process:
            accounts_to_process[account] = {
                "display_name": data["display_name"],
                "services": services_to_process,
            }

    if not accounts_to_process:
        print("All sessions are already valid. Nothing to do.")
        print("Use --refresh to force re-login, or --status to check.")
        return

    # Count total setups needed
    total = sum(len(d["services"]) for d in accounts_to_process.values())
    done = 0
    failed = 0

    print(f"Setting up {total} session(s)...")
    print()
    print("-" * 60)

    for account, data in accounts_to_process.items():
        print(f"\nAccount: {data['display_name']} ({account})")

        for service, svc_config in data["services"].items():
            done += 1
            print(f"\n  [{done}/{total}] {service}")

            success = setup.setup_service(account, service, svc_config)
            if not success:
                failed += 1
                print(f"  FAILED: {account} @ {service}")

            # Brief pause between services to avoid overwhelming Browserless
            if done < total:
                time.sleep(2)

    # Final status
    print()
    print("=" * 60)
    print(f"Setup complete: {done - failed}/{total} successful")
    print()
    print("Final Session Status:")
    setup.print_status_table()

    if failed > 0:
        print(f"WARNING: {failed} session(s) failed to set up.")
        print("Run with --refresh to retry failed sessions.")
        sys.exit(1)
    else:
        print("All sessions ready. Genesis agents can now use authenticated browsers.")
        print()
        print("Test with:")
        print("  python core/browserless_session_manager.py")


if __name__ == "__main__":
    main()