"""
Genesis Browserless Session Manager
====================================
Manages persistent browser sessions for Genesis agents using:
  - Browserless (Chromium) for remote browser automation
  - Elestio Redis for durable cookie/session storage across restarts

Accounts managed:
  kinan@agileadapt.com  → Gemini Ultra, GHL, Telnyx
  sunvision07@gmail.com → Jules Pro, YouTube

Usage:
    from core.browserless_session_manager import BrowserlessSessionManager

    manager = BrowserlessSessionManager()

    # Launch an authenticated browser session (cookies loaded from Redis)
    session_url = manager.launch_authenticated_session("gemini.google.com")

    # Check login state
    status = manager.check_login_state("kinan@agileadapt.com", "gemini.google.com")

    # Save cookies after manual login (used by setup script)
    manager.save_cookies("kinan@agileadapt.com", "gemini.google.com", cookies_list)
"""

import json
import logging
import os
import sys
import time
from typing import Any, Dict, List, Optional

import requests

# Elestio Redis — same pattern as elestio_config.py
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'data', 'genesis-memory'))
try:
    from elestio_config import RedisConfig
    import redis as redis_lib
    REDIS_AVAILABLE = True
except ImportError:
    REDIS_AVAILABLE = False

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Account → Service mapping
# ---------------------------------------------------------------------------

ACCOUNT_SERVICES: Dict[str, Dict[str, Any]] = {
    "kinan@agileadapt.com": {
        "display_name": "Kinan (AgileAdapt)",
        "services": {
            "gemini.google.com": {
                "login_url": "https://gemini.google.com",
                "auth_check_url": "https://gemini.google.com",
                # Cookie names that indicate a valid Google session
                "auth_cookie_keys": ["__Secure-1PSID", "__Secure-3PSID", "SID"],
                "auth_check_selector": "[data-test-id='user-menu']",
                "description": "Gemini Ultra — AI research and deep thinking",
            },
            "app.gohighlevel.com": {
                "login_url": "https://app.gohighlevel.com",
                "auth_check_url": "https://app.gohighlevel.com/dashboard",
                "auth_cookie_keys": ["gh_token", "gh_session", "access_token"],
                "auth_check_selector": ".sidebar-menu",
                "description": "GoHighLevel CRM — client management",
            },
            "portal.telnyx.com": {
                "login_url": "https://portal.telnyx.com",
                "auth_check_url": "https://portal.telnyx.com/#/app/dashboard",
                "auth_cookie_keys": ["telnyx_session", "_telnyx_session", "auth_token"],
                "auth_check_selector": "[data-testid='dashboard']",
                "description": "Telnyx — voice/telephony platform",
            },
        },
    },
    "sunvision07@gmail.com": {
        "display_name": "Sunvision (Gmail)",
        "services": {
            "youtube.com": {
                "login_url": "https://www.youtube.com",
                "auth_check_url": "https://www.youtube.com",
                "auth_cookie_keys": ["__Secure-1PSID", "__Secure-3PSID", "SID", "HSID"],
                "auth_check_selector": "#avatar-btn",
                "description": "YouTube — content management",
            },
            "jules.google.com": {
                "login_url": "https://jules.google.com",
                "auth_check_url": "https://jules.google.com",
                "auth_cookie_keys": ["__Secure-1PSID", "__Secure-3PSID", "SID"],
                "auth_check_selector": "[aria-label='Account']",
                "description": "Jules Pro — Google AI coding assistant",
            },
        },
    },
}

# Redis key prefix — namespaced to avoid collision with other Genesis data
REDIS_KEY_PREFIX = "genesis:browserless:sessions"

# Cookie TTL in Redis — 7 days (most auth cookies last this long)
COOKIE_TTL_SECONDS = 7 * 24 * 60 * 60


class BrowserlessSessionManager:
    """
    Manages persistent browser sessions for Genesis agents.

    Cookies are stored in Elestio Redis so they survive:
    - Container restarts
    - New agent sessions
    - System reboots

    After initial manual login (Collaborative Browser Handoff), subsequent
    sessions are fully automated — agents get authenticated browsers instantly.
    """

    def __init__(
        self,
        browserless_url: Optional[str] = None,
        browserless_token: Optional[str] = None,
    ):
        """
        Args:
            browserless_url:   URL of the Browserless instance.
                               Defaults to env var BROWSERLESS_URL or localhost:3000.
            browserless_token: Auth token for Browserless.
                               Defaults to env var BROWSERLESS_TOKEN.
        """
        self.browserless_url = (
            browserless_url
            or os.environ.get("BROWSERLESS_URL", "http://localhost:3000")
        ).rstrip("/")

        self.browserless_token = (
            browserless_token
            or os.environ.get("BROWSERLESS_TOKEN", "")
        )

        self._redis: Optional[Any] = None
        self._init_redis()

    # ------------------------------------------------------------------
    # Redis connection
    # ------------------------------------------------------------------

    def _init_redis(self) -> None:
        """Connect to Elestio Redis for persistent session storage."""
        if not REDIS_AVAILABLE:
            logger.warning(
                "redis-py or elestio_config not available — "
                "session persistence disabled, running in memory-only mode"
            )
            self._memory_store: Dict[str, str] = {}
            return

        try:
            params = RedisConfig.get_connection_params()
            self._redis = redis_lib.Redis(**params)
            self._redis.ping()
            logger.info(
                "BrowserlessSessionManager connected to Elestio Redis at %s:%s",
                params["host"],
                params["port"],
            )
        except Exception as exc:
            logger.error("Failed to connect to Elestio Redis: %s", exc)
            logger.warning("Falling back to in-memory session store (non-persistent)")
            self._redis = None
            self._memory_store = {}

    def _redis_get(self, key: str) -> Optional[str]:
        """Get a value from Redis (or in-memory fallback)."""
        if self._redis:
            try:
                val = self._redis.get(key)
                return val if isinstance(val, str) else (val.decode() if val else None)
            except Exception as exc:
                logger.error("Redis GET error for key %s: %s", key, exc)
                return None
        return self._memory_store.get(key)

    def _redis_set(self, key: str, value: str, ttl: int = COOKIE_TTL_SECONDS) -> bool:
        """Set a value in Redis (or in-memory fallback)."""
        if self._redis:
            try:
                self._redis.setex(key, ttl, value)
                return True
            except Exception as exc:
                logger.error("Redis SET error for key %s: %s", key, exc)
                return False
        self._memory_store[key] = value
        return True

    def _redis_delete(self, key: str) -> bool:
        """Delete a key from Redis (or in-memory fallback)."""
        if self._redis:
            try:
                self._redis.delete(key)
                return True
            except Exception as exc:
                logger.error("Redis DELETE error for key %s: %s", key, exc)
                return False
        self._memory_store.pop(key, None)
        return True

    # ------------------------------------------------------------------
    # Cookie persistence
    # ------------------------------------------------------------------

    def _make_redis_key(self, account: str, service: str) -> str:
        """Build the Redis key for a given account + service pair."""
        # Sanitize: replace dots/@ with underscores for clean keys
        safe_account = account.replace("@", "_at_").replace(".", "_")
        safe_service = service.replace(".", "_")
        return f"{REDIS_KEY_PREFIX}:{safe_account}:{safe_service}"

    def save_cookies(
        self,
        account: str,
        service: str,
        cookies: List[Dict[str, Any]],
        ttl: int = COOKIE_TTL_SECONDS,
    ) -> bool:
        """
        Save browser cookies for account+service to Elestio Redis.

        Args:
            account:  Email address (e.g., "kinan@agileadapt.com")
            service:  Service hostname (e.g., "gemini.google.com")
            cookies:  List of cookie dicts from Playwright/CDP
                      Each dict: {name, value, domain, path, expires, httpOnly, secure, sameSite}
            ttl:      Redis TTL in seconds (default 7 days)

        Returns:
            True if saved successfully
        """
        if account not in ACCOUNT_SERVICES:
            logger.error("Unknown account: %s", account)
            return False

        key = self._make_redis_key(account, service)
        payload = {
            "account": account,
            "service": service,
            "cookies": cookies,
            "saved_at": time.time(),
            "count": len(cookies),
        }

        success = self._redis_set(key, json.dumps(payload), ttl=ttl)
        if success:
            logger.info(
                "Saved %d cookies for %s @ %s (TTL: %ds)",
                len(cookies), account, service, ttl
            )
        return success

    def load_cookies(
        self, account: str, service: str
    ) -> Optional[List[Dict[str, Any]]]:
        """
        Load persisted cookies for account+service from Elestio Redis.

        Args:
            account:  Email address
            service:  Service hostname

        Returns:
            List of cookie dicts, or None if not found / expired
        """
        key = self._make_redis_key(account, service)
        raw = self._redis_get(key)

        if not raw:
            logger.info("No persisted cookies found for %s @ %s", account, service)
            return None

        try:
            payload = json.loads(raw)
            cookies = payload.get("cookies", [])
            saved_at = payload.get("saved_at", 0)
            age_hours = (time.time() - saved_at) / 3600
            logger.info(
                "Loaded %d cookies for %s @ %s (age: %.1fh)",
                len(cookies), account, service, age_hours
            )
            return cookies
        except json.JSONDecodeError as exc:
            logger.error("Failed to decode cookies for %s @ %s: %s", account, service, exc)
            return None

    def clear_cookies(self, account: str, service: str) -> bool:
        """
        Remove persisted cookies for account+service (forces re-login).

        Args:
            account:  Email address
            service:  Service hostname

        Returns:
            True if cleared
        """
        key = self._make_redis_key(account, service)
        success = self._redis_delete(key)
        if success:
            logger.info("Cleared cookies for %s @ %s", account, service)
        return success

    def list_sessions(self) -> List[Dict[str, Any]]:
        """
        List all persisted sessions currently in Redis.

        Returns:
            List of dicts with account, service, cookie_count, saved_at, age_hours
        """
        sessions = []

        for account, account_data in ACCOUNT_SERVICES.items():
            for service in account_data["services"]:
                key = self._make_redis_key(account, service)
                raw = self._redis_get(key)
                if raw:
                    try:
                        payload = json.loads(raw)
                        sessions.append({
                            "account": account,
                            "service": service,
                            "cookie_count": payload.get("count", 0),
                            "saved_at": payload.get("saved_at", 0),
                            "age_hours": (time.time() - payload.get("saved_at", 0)) / 3600,
                        })
                    except json.JSONDecodeError:
                        continue

        return sessions

    # ------------------------------------------------------------------
    # Login state check
    # ------------------------------------------------------------------

    def check_login_state(self, account: str, service: str) -> Dict[str, Any]:
        """
        Check whether stored cookies are present and appear valid.

        This is a quick Redis check — it does NOT launch a browser.
        To do a live browser check, use launch_authenticated_session()
        and inspect the resulting page.

        Args:
            account:  Email address
            service:  Service hostname

        Returns:
            {
                "logged_in": bool,
                "cookie_count": int,
                "age_hours": float,
                "needs_refresh": bool,  # True if cookies older than 6 days
                "account": str,
                "service": str,
            }
        """
        cookies = self.load_cookies(account, service)

        if not cookies:
            return {
                "logged_in": False,
                "cookie_count": 0,
                "age_hours": 0.0,
                "needs_refresh": True,
                "account": account,
                "service": service,
            }

        key = self._make_redis_key(account, service)
        raw = self._redis_get(key)
        payload = json.loads(raw) if raw else {}
        saved_at = payload.get("saved_at", 0)
        age_hours = (time.time() - saved_at) / 3600

        # Check auth cookie keys are present (quick sanity check)
        service_config = (
            ACCOUNT_SERVICES.get(account, {})
            .get("services", {})
            .get(service, {})
        )
        auth_keys = service_config.get("auth_cookie_keys", [])
        cookie_names = {c.get("name", "") for c in cookies}
        has_auth_cookie = any(k in cookie_names for k in auth_keys) if auth_keys else True

        return {
            "logged_in": has_auth_cookie and len(cookies) > 0,
            "cookie_count": len(cookies),
            "age_hours": round(age_hours, 1),
            "needs_refresh": age_hours > (6 * 24),  # Refresh if older than 6 days
            "account": account,
            "service": service,
            "auth_cookie_present": has_auth_cookie,
        }

    # ------------------------------------------------------------------
    # Authenticated session launch
    # ------------------------------------------------------------------

    def launch_authenticated_session(
        self,
        service: str,
        account: Optional[str] = None,
        navigate_to: Optional[str] = None,
        headless: bool = True,
        stealth: bool = True,
    ) -> Dict[str, Any]:
        """
        Launch a Browserless Chromium session pre-loaded with stored cookies.

        This uses the Browserless /function endpoint to:
        1. Open a new page
        2. Navigate to the service domain to set the cookie context
        3. Inject all stored cookies
        4. Navigate to the target URL

        Args:
            service:     Service hostname (e.g., "gemini.google.com")
            account:     Email address. If None, auto-selects based on service.
            navigate_to: URL to navigate to after injecting cookies.
                         Defaults to the service login URL.
            headless:    Run Chrome headless (True for agents, False for debugging)
            stealth:     Use stealth mode to avoid bot detection

        Returns:
            {
                "success": bool,
                "ws_endpoint": str,   # WebSocket URL for Playwright/CDP connection
                "session_id": str,
                "account": str,
                "service": str,
                "cookies_injected": int,
                "error": str or None,
            }
        """
        # Auto-select account if not specified
        if account is None:
            account = self._find_account_for_service(service)
            if account is None:
                return {
                    "success": False,
                    "error": f"No account configured for service: {service}",
                    "service": service,
                }

        # Load cookies from Redis
        cookies = self.load_cookies(account, service)

        service_config = (
            ACCOUNT_SERVICES.get(account, {})
            .get("services", {})
            .get(service, {})
        )
        target_url = navigate_to or service_config.get(
            "auth_check_url",
            service_config.get("login_url", f"https://{service}")
        )

        if not cookies:
            logger.warning(
                "No cookies found for %s @ %s — session will be unauthenticated. "
                "Run setup_browserless_profiles.py to perform initial login.",
                account, service
            )
            cookies = []

        # Build the Browserless /function payload
        # This JS code runs inside Browserless and returns the page state
        js_function = self._build_cookie_injection_script(cookies, target_url, stealth)

        try:
            response = requests.post(
                f"{self.browserless_url}/function",
                headers={
                    "Authorization": f"Bearer {self.browserless_token}",
                    "Content-Type": "application/json",
                },
                json={
                    "code": js_function,
                    "context": {
                        "cookies": cookies,
                        "targetUrl": target_url,
                    },
                    "launch": {
                        "headless": headless,
                        "args": [
                            "--no-sandbox",
                            "--disable-dev-shm-usage",
                            "--disable-gpu",
                            "--disable-blink-features=AutomationControlled",
                        ] + (["--disable-web-security"] if stealth else []),
                    },
                },
                timeout=60,
            )

            if response.status_code == 200:
                result = response.json()
                return {
                    "success": True,
                    "result": result,
                    "account": account,
                    "service": service,
                    "cookies_injected": len(cookies),
                    "target_url": target_url,
                    "error": None,
                }
            else:
                return {
                    "success": False,
                    "error": f"Browserless returned HTTP {response.status_code}: {response.text[:500]}",
                    "account": account,
                    "service": service,
                }

        except requests.RequestException as exc:
            return {
                "success": False,
                "error": f"Failed to connect to Browserless at {self.browserless_url}: {exc}",
                "account": account,
                "service": service,
            }

    def get_ws_endpoint(self, session_id: Optional[str] = None) -> str:
        """
        Get the WebSocket endpoint for connecting Playwright directly.

        For use when Genesis agents want to control the browser via
        Playwright's remote browser connection rather than /function API.

        Args:
            session_id: Optional specific session ID

        Returns:
            WebSocket URL string
        """
        if session_id:
            return f"ws://{self.browserless_url.replace('http://', '').replace('https://', '')}/chromium?token={self.browserless_token}&session={session_id}"
        return f"ws://{self.browserless_url.replace('http://', '').replace('https://', '')}/chromium?token={self.browserless_token}"

    # ------------------------------------------------------------------
    # Session health check
    # ------------------------------------------------------------------

    def health_check(self) -> Dict[str, Any]:
        """
        Check Browserless service health and session counts.

        Returns:
            {
                "browserless_online": bool,
                "active_sessions": int,
                "queued_sessions": int,
                "max_sessions": int,
                "redis_online": bool,
                "persisted_sessions": int,
            }
        """
        result: Dict[str, Any] = {
            "browserless_url": self.browserless_url,
            "browserless_online": False,
            "active_sessions": 0,
            "redis_online": False,
            "persisted_sessions": 0,
        }

        # Check Browserless
        try:
            resp = requests.get(
                f"{self.browserless_url}/health",
                headers={"Authorization": f"Bearer {self.browserless_token}"},
                timeout=10,
            )
            if resp.status_code == 200:
                result["browserless_online"] = True
                health_data = resp.json()
                result["active_sessions"] = health_data.get("running", 0)
                result["queued_sessions"] = health_data.get("queued", 0)
                result["max_sessions"] = health_data.get("maxConcurrent", 0)
        except Exception as exc:
            result["browserless_error"] = str(exc)

        # Check Redis
        try:
            if self._redis:
                self._redis.ping()
                result["redis_online"] = True
        except Exception:
            pass

        # Count persisted sessions
        result["persisted_sessions"] = len(self.list_sessions())

        return result

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _find_account_for_service(self, service: str) -> Optional[str]:
        """Find which account is configured for a given service hostname."""
        for account, data in ACCOUNT_SERVICES.items():
            if service in data["services"]:
                return account
        return None

    def _build_cookie_injection_script(
        self,
        cookies: List[Dict[str, Any]],
        target_url: str,
        stealth: bool = True,
    ) -> str:
        """
        Build the JavaScript function that Browserless executes to inject cookies.

        The script:
        1. Opens a page
        2. Navigates to the domain to establish context
        3. Injects all cookies
        4. Navigates to the target URL
        5. Returns page title and URL to confirm success
        """
        cookies_json = json.dumps(cookies)
        target_json = json.dumps(target_url)

        stealth_script = ""
        if stealth:
            stealth_script = """
    // Stealth: hide automation markers
    await page.evaluateOnNewDocument(() => {
        Object.defineProperty(navigator, 'webdriver', { get: () => false });
        Object.defineProperty(navigator, 'plugins', { get: () => [1, 2, 3, 4, 5] });
        Object.defineProperty(navigator, 'languages', { get: () => ['en-AU', 'en'] });
        window.chrome = { runtime: {} };
    });"""

        return f"""
module.exports = async ({{ page, context }}) => {{
    const cookies = {cookies_json};
    const targetUrl = {target_json};

    {stealth_script}

    // Set cookies if we have them
    if (cookies && cookies.length > 0) {{
        // Navigate to domain first to establish cookie context
        const urlObj = new URL(targetUrl);
        const domainUrl = urlObj.origin;

        try {{
            await page.goto(domainUrl, {{ waitUntil: 'domcontentloaded', timeout: 30000 }});
        }} catch (e) {{
            // Domain navigation may fail — that's OK, we just need the context
        }}

        // Inject cookies
        for (const cookie of cookies) {{
            try {{
                await page.setCookie(cookie);
            }} catch (e) {{
                // Skip invalid cookies silently
            }}
        }}
    }}

    // Navigate to the actual target
    await page.goto(targetUrl, {{ waitUntil: 'networkidle2', timeout: 45000 }});

    const title = await page.title();
    const url = page.url();
    const cookieCount = cookies ? cookies.length : 0;

    return {{
        title,
        url,
        cookiesInjected: cookieCount,
        success: true,
    }};
}};
"""


# ---------------------------------------------------------------------------
# Convenience functions for Genesis agents
# ---------------------------------------------------------------------------

def get_session_manager() -> BrowserlessSessionManager:
    """Get a configured BrowserlessSessionManager instance (singleton pattern)."""
    return BrowserlessSessionManager()


def get_authenticated_gemini_session() -> Dict[str, Any]:
    """
    Get an authenticated Gemini session for kinan@agileadapt.com.
    Returns Browserless session result dict.
    """
    manager = get_session_manager()
    return manager.launch_authenticated_session(
        service="gemini.google.com",
        account="kinan@agileadapt.com",
        navigate_to="https://gemini.google.com",
    )


def get_authenticated_ghl_session() -> Dict[str, Any]:
    """
    Get an authenticated GoHighLevel session for kinan@agileadapt.com.
    Returns Browserless session result dict.
    """
    manager = get_session_manager()
    return manager.launch_authenticated_session(
        service="app.gohighlevel.com",
        account="kinan@agileadapt.com",
        navigate_to="https://app.gohighlevel.com/dashboard",
    )


def get_authenticated_telnyx_session() -> Dict[str, Any]:
    """
    Get an authenticated Telnyx session for kinan@agileadapt.com.
    Returns Browserless session result dict.
    """
    manager = get_session_manager()
    return manager.launch_authenticated_session(
        service="portal.telnyx.com",
        account="kinan@agileadapt.com",
        navigate_to="https://portal.telnyx.com/#/app/dashboard",
    )


def get_authenticated_youtube_session() -> Dict[str, Any]:
    """
    Get an authenticated YouTube session for sunvision07@gmail.com.
    Returns Browserless session result dict.
    """
    manager = get_session_manager()
    return manager.launch_authenticated_session(
        service="youtube.com",
        account="sunvision07@gmail.com",
        navigate_to="https://www.youtube.com/feed/subscriptions",
    )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")

    manager = BrowserlessSessionManager()

    print("Genesis Browserless Session Manager")
    print("=" * 50)

    # Health check
    print("\nHealth Check:")
    health = manager.health_check()
    for k, v in health.items():
        print(f"  {k}: {v}")

    # List sessions
    print("\nPersisted Sessions:")
    sessions = manager.list_sessions()
    if sessions:
        for s in sessions:
            status = "STALE" if s["age_hours"] > 144 else "FRESH"
            print(f"  [{status}] {s['account']} @ {s['service']} — {s['cookie_count']} cookies, {s['age_hours']:.1f}h old")
    else:
        print("  No sessions persisted yet. Run setup_browserless_profiles.py")

    # Login state for each configured account+service
    print("\nLogin States:")
    for account, data in ACCOUNT_SERVICES.items():
        for service in data["services"]:
            state = manager.check_login_state(account, service)
            icon = "LOGGED IN" if state["logged_in"] else "NOT LOGGED IN"
            refresh = " [NEEDS REFRESH]" if state["needs_refresh"] else ""
            print(f"  {icon}{refresh}: {account} @ {service} ({state['cookie_count']} cookies)")