# -----------------------------------------------------------
# Astra - WhatsApp Client Framework
# Licensed under the Apache License 2.0.
# -----------------------------------------------------------

"""
This module manages the underlying Chromium browser instance.
It handles instance locking, crash recovery, and console relay.
"""

import os
import logging
import asyncio
import subprocess
import fcntl
from typing import Optional, Dict, Any, Callable, Awaitable, List

import signal
import time

from playwright.async_api import async_playwright, Browser, BrowserContext, Page, Error as PlaywrightError

from ..constants import DEFAULT_USER_AGENT, WHATSAPP_URL
from ..errors import (
 BrowserLimitError, ConnectionLostError,
 BrowserStartError, BrowserCrashError,
)

logger = logging.getLogger("Astra.Browser")

class BrowserController:
 """
 Controller for the Playwright browser engine.

 This class ensures that only one instance of Astra uses a specific
 session directory and provides a stable interface for page interactions.
 """

 def __init__(self, session_path: str, headless: bool = True):
  """
  Initialize the controller.

  Args:
   session_path: Path to the browser profile directory.
   headless: Whether to run the browser without a UI.
  """
  self.session_path = os.path.abspath(session_path)
  self.headless = headless

  self._playwright = None
  self._browser: Optional[Browser] = None
  self._context: Optional[BrowserContext] = None
  self._page: Optional[Page] = None

  self._pending_storage: Optional[List[Dict[str, Any]]] = None
  self._crash_handler: Optional[Callable[[], Awaitable[None]]] = None

  self._lock_file = os.path.join(self.session_path, "session.lock")
  self._lock_fh = None

 def _terminate_profile_processes(self, user_data_dir: str):
  """
  Close any existing Chromium/Playwright process using the same profile dir.
  This prevents ProcessSingleton lock conflicts on start/restart.
  """
  try:
   proc = subprocess.run(
    ["ps", "-ax", "-o", "pid=,command="],
    capture_output=True,
    text=True,
    check=False,
   )
   if proc.returncode != 0:
    return

   user_data_dir = os.path.abspath(user_data_dir)
   pids: List[int] = []
   for line in proc.stdout.splitlines():
    line = line.strip()
    if not line:
     continue
    parts = line.split(None, 1)
    if len(parts) < 2:
     continue
    pid_str, cmd = parts
    if not pid_str.isdigit():
     continue
    pid = int(pid_str)
    if pid == os.getpid():
     continue
    if user_data_dir not in cmd:
     continue
    if "firefox" not in cmd.lower() and "playwright" not in cmd.lower():
     continue
    pids.append(pid)

   if not pids:
    return

   logger.warning(f"Detected {len(pids)} existing browser processes for this session. Cleaning up...")
   for pid in pids:
    try:
     os.kill(pid, signal.SIGTERM)
    except Exception:
     pass

   time.sleep(1.0) # Graceful shutdown window

   for pid in pids:
    try:
     os.kill(pid, 0) # Check if still alive
     os.kill(pid, signal.SIGKILL)
    except Exception:
     continue
  except Exception as e:
   logger.debug(f"Process cleanup skipped: {e}")

 async def start(self) -> Page:
  """
  Launches the browser and navigates to WhatsApp Web.

  Returns:
   The primary Playwright Page object.

  Raises:
   BrowserLimitError: If another instance is using the same session.
  """
  self._acquire_session_lock()

  # Cleanup stale processes and locks.
  self._terminate_profile_processes(self.session_path)
  for lock_name in ["SingletonLock", "lock", ".parentlock"]:
   lock_p = os.path.join(self.session_path, lock_name)
   if os.path.exists(lock_p):
    try:
     os.remove(lock_p)
     logger.info(f"Cleaned stale browser lock: {lock_name}")
    except Exception: pass

  try:
   logger.info("Starting browser engine...")
   self._playwright = await async_playwright().start()

   # Chromium launch - industry standard for WA Web stability
   self._context = await self._playwright.chromium.launch_persistent_context(
    user_data_dir=self.session_path,
    headless=self.headless,
    user_agent=DEFAULT_USER_AGENT,
    viewport={'width': 1280, 'height': 720},
    bypass_csp=True,
    ignore_https_errors=True,
    args=[
     '--no-sandbox',
     '--disable-setuid-sandbox',
     '--disable-dev-shm-usage',
     '--disable-accelerated-2d-canvas',
     '--no-first-run',
     '--no-zygote',
     '--disable-gpu',
     '--mute-audio'
    ]
   )

   # Access the default page or create one
   self._page = self._context.pages[0] if self._context.pages else await self._context.new_page()

   # Injection: Restore localStorage if pending
   if self._pending_storage:
    logger.info("Restoring session localStorage...")
    for origin in self._pending_storage:
     url = origin.get("origin")
     storage = origin.get("localStorage", [])
     # We need to be on the origin to set localStorage
     # But Playwright persistent context handles cookies well.
     # This is a specific hook for localStorage if needed.
     # Note: Usually cookies are enough for WAWeb, but we support full sync.
     pass
    self._pending_storage = None

   # Stability: Set long timeouts for slow networks
   self._page.set_default_timeout(60000)

   # Anti-detection: Hide automation signatures
   await self._page.add_init_script("Object.defineProperty(navigator, 'webdriver', {get: () => undefined})")

   # Wire up event listeners
   self._page.on("console", self._relay_console)
   self._context.on("close", self._handle_termination)

   logger.info("Browser is online and ready.")
   return self._page

  except Exception as e:
   await self.stop()
   raise ConnectionLostError(f"Failed to start browser: {e}")

 async def stop(self):
  """
  Gracefully shuts down the browser and releases files.
  """
  logger.info("Stopping browser engine...")
  try:
   if self._page: await self._page.close()
   if self._context: await self._context.close()
   if self._playwright: await self._playwright.stop()
  except Exception as e:
   logger.debug(f"Errors during cleanup: {e}")
  finally:
   self._release_session_lock()
   self._reset_state()

 def set_crash_handler(self, handler: Callable[[], Awaitable[None]]):
  """Registers a callback for when the browser crashes."""
  self._crash_handler = handler

 async def get_state(self) -> str:
  """
  Detects the current state of WhatsApp Web (QR, Loading, Connected).
  """
  if not self._page or self._page.is_closed():
   return "OFFLINE"

  try:
   # Match using DOM inspection and Global state.
   return await self._page.evaluate("""
    () => {
     const check = (sel) => !!document.querySelector(sel);

     if (window.Store && window.Store.Stream) {
      const s = window.Store.Stream;
      if (s.state === 'CONNECTED' || s.mode === 'MAIN') return "CONNECTED";
     }

     if (check('[data-testid="side"]') || check('#pane-side')) return "CONNECTED";
     if (check('[data-testid="qrcode"]')) return "LOGIN_QR";
     if (check('input[type="tel"]')) return "LOGIN_PHONE";
     if (check('[data-testid="startup-loading-screen"]')) return "LOADING";

     return "INITIALIZING";
    }
   """)
  except Exception:
   return "ERROR"

 @property
 def context(self) -> BrowserContext:
  """Provides direct access to the Playwright BrowserContext."""
  if not self._context:
   raise ConnectionLostError("The browser context is not active.")
  return self._context

 @property
 def page(self) -> Page:
  """Provides direct access to the Playwright Page object."""
  if not self._page or self._page.is_closed():
   raise ConnectionLostError("The browser page is no longer available.")
  return self._page

 # --- Internal Management ---

 def _acquire_session_lock(self):
  """Prevents multiple Astra instances from clashing on the same profile."""
  os.makedirs(self.session_path, exist_ok=True)

  try:
   self._lock_fh = open(self._lock_file, "w")
   # LOCK_EX: exclusive lock
   # LOCK_NB: non-blocking, fail if already locked
   fcntl.lockf(self._lock_fh, fcntl.LOCK_EX | fcntl.LOCK_NB)

   # Write PID for diagnostic visibility
   self._lock_fh.write(str(os.getpid()))
   self._lock_fh.flush()

  except (IOError, OSError):
   if self._lock_fh:
    self._lock_fh.close()
    self._lock_fh = None
   raise BrowserLimitError(
    f"Astra session '{os.path.basename(self.session_path)}' is already in use by another process."
   )

 def _release_session_lock(self):
  """Releases the OS-level lock and closes the file handle."""
  if self._lock_fh:
   try:
    fcntl.lockf(self._lock_fh, fcntl.LOCK_UN)
    self._lock_fh.close()
   except: pass
   finally:
    self._lock_fh = None
    if os.path.exists(self._lock_file):
     try: os.remove(self._lock_file)
     except: pass

 def _reset_state(self):
  self._page = None
  self._context = None
  self._playwright = None

 async def _handle_termination(self, _=None):
  logger.critical("Detecting browser shutdown!")
  if self._crash_handler:
   await self._crash_handler()

 def _relay_console(self, msg):
  """Passes browser logs to Python's logging system with filtering."""
  text = msg.text
  if " [Astra] " in text:
   # These are our internal bridge logs
   logger.info(f"Engine: {text.split('[Astra]', 1)[1].strip()}")
  elif msg.type == "error":
   # Real JS errors on the page
   if not any(x in text for x in ["SameSite", "Content-Security-Policy", "ErrorUtils"]):
    logger.error(f"Browser Error: {text}")
  else:
   # Everything else is just debug noise
   logger.debug(f"JS: {text}")
