# backend/tools/browser_tools.py
import asyncio
import logging
import os
from contextlib import asynccontextmanager

try:
    from playwright.async_api import async_playwright
except ImportError:
    async_playwright = None

try:
    from playwright_stealth import Stealth
    stealth_async = Stealth().apply_stealth_async
except ImportError:
    stealth_async = None

_playwright = None
_browser = None
_context_pool_semaphore = asyncio.Semaphore(3)

PROFILE_DIR = os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS', 'Playwright_Profile')
STATE_FILE = os.path.join(PROFILE_DIR, 'state.json')

async def _init_browser():
    global _playwright, _browser
    if async_playwright is None:
        raise ImportError("playwright is not installed")
        
    if _browser is None:
        os.makedirs(PROFILE_DIR, exist_ok=True)
        _playwright = await async_playwright().start()
        _browser = await _playwright.chromium.launch(
            headless=True,
            args=["--disable-blink-features=AutomationControlled"]
        )

@asynccontextmanager
async def get_ephemeral_context():
    """
    Yields an isolated browser context while enforcing a max pool of 3 concurrent sessions.
    Ensures context is closed and state is saved in the finally block.
    """
    await _init_browser()
    async with _context_pool_semaphore:
        storage_state = STATE_FILE if os.path.exists(STATE_FILE) else None
        context = await _browser.new_context(storage_state=storage_state)
        try:
            yield context
        finally:
            try:
                # Save cookies before closing
                await context.storage_state(path=STATE_FILE)
            except Exception as e:
                logging.error(f"Failed to save storage state: {e}")
            await context.close()

async def navigate(url: str) -> dict:
    try:
        async with get_ephemeral_context() as context:
            page = await context.new_page()
            if stealth_async:
                await stealth_async(page)
                
            await page.goto(url)
            title = await page.title()
            await page.close()
            return {"status": "success", "title": title}
    except Exception as e:
        return {"error": str(e)}

async def close_browser():
    global _playwright, _browser
    if _browser:
        await _browser.close()
        _browser = None
    if _playwright:
        await _playwright.stop()
        _playwright = None