# -*- coding: utf-8 -*-
from playwright.async_api import async_playwright, Browser, BrowserContext
import json
import shutil
import errno
from pathlib import Path
import asyncio
import uuid
from typing import List, Optional

import logging

logging.getLogger(__name__)


# Base directories for storing profile data and storage states (cross-platform)
BASE_PROFILE_DIR = Path("data/chromium_profiles")  # Directory for user data/profile directories
BASE_STATE_DIR = Path("data/browser_states")      # Directory for storage state JSON files

def copytree_ignore_errors(src: Path, dst: Path, ignore_patterns: list = ["SingletonLock", "SingletonCookie", "SingletonSocket"]):
    """Copy a directory tree, ignoring specified files and handling errors gracefully."""
    # Define a function to ignore files matching specified patterns (e.g., Singleton* files)
    def ignore_files(directory, files):
        return [f for f in files if any(pattern in f for pattern in ignore_patterns)]
    
    try:
        # Copy the directory tree, ignoring specified files
        shutil.copytree(src, dst, ignore=ignore_files, dirs_exist_ok=False)
    except shutil.Error as e:
        # Log non-critical copy errors (e.g., permission issues)
        logging.info(f"Non-critical errors during copy: {e}")
    except OSError as e:
        # Handle specific OS errors (e.g., file not found, device busy) gracefully
        if e.errno not in (errno.ENOENT, errno.EAGAIN, errno.EBUSY):
            raise
        logging.info(f"Skipped problematic file during copy: {e}")

class PersistentBrowser:
    """Simulates a Playwright Browser object using either launch_persistent_context or browser.launch."""
    
    def __init__(self, playwright, base_profile_dir: Path, use_persistent_context: bool, extension_paths: List[Path] = None, **kwargs):
        """Initialize the simulated browser."""
        # Store Playwright instance for browser operations
        self._playwright = playwright
        # Base directory for profile data
        self._base_profile_dir = base_profile_dir
        # Flag to choose between launch_persistent_context and browser.launch
        self._use_persistent_context = use_persistent_context
        # List of extension paths to copy to each context (converted to Path objects)
        self._extension_paths = [Path(p) for p in (extension_paths or [])]
        # Dictionary to store context_id -> BrowserContext
        self._contexts = {}
        # Browser instance for browser.launch mode
        self._browser = None
        # Flag to track if the browser is closed
        self._is_closed = False

        self._kwargs = kwargs

    def contexts(self):
       return self._browser.contexts

    async def launch(self):
        """Simulate browser.launch() by preparing the environment."""
        # Create base profile directory if it doesn't exist
        self._base_profile_dir.mkdir(parents=True, exist_ok=True)
        # If using browser.launch mode, launch a single Chromium instance
        if not self._use_persistent_context:
            self._browser = await self._playwright.chromium.launch(headless=False, **self._kwargs)
        logging.info("Persistent browser launched.")
        return self

    async def new_context(self, context_id: str = None, **kwargs) -> BrowserContext:
        """Create a new context with its own persistent state."""
        # Check if the browser is closed
        if self._is_closed:
            raise RuntimeError("Browser is closed.")

        # Generate a unique context ID if not provided
        context_id = context_id or str(uuid.uuid4())
        
        # Create or restore profile directory for this context
        profile_dir = self._get_profile_dir(context_id)
        if not profile_dir.exists():
            restore_profile_dir(context_id, self._base_profile_dir)
            create_profile_dir(context_id, self._base_profile_dir)

        # Create context based on mode
        if self._use_persistent_context:
            # Use launch_persistent_context for native persistence via user_data_dir
            context = await self._playwright.chromium.launch_persistent_context(
                user_data_dir=profile_dir,
                headless=False,
                viewport=kwargs.get("viewport", {"width": 1280, "height": 720}),
                **{k: v for k, v in kwargs.items() if k != "extension_paths"}  # Filter out extension_paths
            )
        else:
            # Use browser.launch with manual persistence
            if not self._browser:
                raise RuntimeError("Browser not launched.")
            if kwargs.get("viewport"):
               vp = kwargs.get("viewport", {"width": 1280, "height": 720})
            else:
               vp = None
            context = await self._browser.new_context(
                viewport=vp,
                **{k: v for k, v in kwargs.items() if k != "extension_paths" and k!="viewport"}
            )

        # Load storage state manually if it exists
        state_file = BASE_STATE_DIR / f"browser_state_{context_id}.json"
        if state_file.exists():
            with open(state_file, 'r', encoding='utf-8') as f:
                storage_state = json.load(f)
                # Load cookies from storage state
                if "cookies" in storage_state:
                    await context.add_cookies(storage_state["cookies"])
                # Load local storage from storage state
                if "origins" in storage_state:
                    for origin in storage_state["origins"]:
                        for item in origin.get("localStorage", []):
                            await context.evaluate(
                                """
                                ({name, value}) => {
                                    window.localStorage.setItem(name, value);
                                }
                                """,
                                {"name": item["name"], "value": item["value"]}
                            )
            logging.info(f"Loaded browser state for context {context_id} from {state_file}")

        # Store the context
        self._contexts[context_id] = context
        logging.info(f"Created new context with ID {context_id} using profile {profile_dir}")

        # Install pre-set extensions for this context
        for ext_path in self._extension_paths:
            await install_extension(context, ext_path, context_id)

        return context

    async def close(self):
        """Close all contexts and mark the browser as closed."""
        if not self._is_closed:
            for context_id, context in self._contexts.items():
                # Save storage state before closing
                await save_browser_state(context, context_id)
                # Close context to release Singleton* files
                await context.close()
                # Backup profile directory after closing (only for persistent context mode)
                if self._use_persistent_context:
                    backup_profile_dir(context_id, self._base_profile_dir)
            # Close the browser if in browser.launch mode
            if self._browser:
                await self._browser.close()
            self._contexts.clear()
            self._is_closed = True
            logging.info("Persistent browser closed.")

    def _get_profile_dir(self, context_id: str) -> Path:
        """Get the profile directory for a context."""
        return self._base_profile_dir / f"chromium_profile_{context_id}"

async def save_browser_state(context: BrowserContext, context_id: str):
    """Save the browser's storage state for a specific context."""
    state_file = BASE_STATE_DIR / f"browser_state_{context_id}.json"
    await context.storage_state(path=state_file)
    logging.info(f"Saved browser state for context {context_id} to {state_file}")

def create_profile_dir(context_id: str, base_profile_dir: Path):
    """Create a profile directory for a specific context if it doesn't exist."""
    profile_dir = base_profile_dir / f"chromium_profile_{context_id}"
    profile_dir.mkdir(parents=True, exist_ok=True)
    logging.info(f"Created profile directory for context {context_id} at {profile_dir}")

def backup_profile_dir(context_id: str, base_profile_dir: Path):
    """Copy the profile directory for a specific context to a backup location."""
    profile_dir = base_profile_dir / f"chromium_profile_{context_id}"
    backup_dir = base_profile_dir / f"chromium_profile_backup_{context_id}"
    if profile_dir.exists():
        if backup_dir.exists():
            shutil.rmtree(backup_dir)  # Remove old backup
        copytree_ignore_errors(profile_dir, backup_dir)
        logging.info(f"Backed up profile directory for context {context_id} to {backup_dir}")

def restore_profile_dir(context_id: str, base_profile_dir: Path):
    """Restore the profile directory for a specific context from the backup."""
    profile_dir = base_profile_dir / f"chromium_profile_{context_id}"
    backup_dir = base_profile_dir / f"chromium_profile_backup_{context_id}"
    if backup_dir.exists():
        if profile_dir.exists():
            shutil.rmtree(profile_dir)  # Remove current profile dir
        copytree_ignore_errors(backup_dir, profile_dir)
        logging.info(f"Restored profile directory for context {context_id} from {backup_dir}")
    else:
        logging.info(f"No backup profile directory found for context {context_id}.")

async def install_extension(context: BrowserContext, extension_path: Path, context_id: str):
    """Load an extension for a specific context."""
    extension_path = Path(extension_path)
    profile_dir = BASE_PROFILE_DIR / f"chromium_profile_{context_id}"
    if extension_path.exists():
        # Copy extension to the user data directory's Extensions folder
        extension_dest = profile_dir / "Extensions" / extension_path.name
        if not extension_dest.exists():
            shutil.copytree(extension_path, extension_dest)
            logging.info(f"Copied extension to {extension_dest} for context {context_id}")
        # Add a script to simulate extension interaction (optional)
        await context.add_init_script(
            script="""chrome.runtime.sendMessage({ message: 'Extension loaded' });"""
        )
        logging.info(f"Extension loaded from {extension_path} for context {context_id}")
    else:
        logging.info(f"Extension path {extension_path} does not exist for context {context_id}.")

async def run_context(persistent_browser: PersistentBrowser, context_id: str, extension_paths: List[Path]):
    """Run a single browser context with its own persistent state."""
    # Create a new context with the pre-set extensions
    context = await persistent_browser.new_context(
        context_id=context_id,
        extension_paths=extension_paths
    )

    # Open a new page in the context
    page = await context.new_page()

    # Navigate to a sample website
    await page.goto("https://example.com")

    # Print the page title as an example action
    logging.info(f"Page title for context {context_id}: {await page.title()}")

async def run(playwright, num_contexts=3, use_persistent_context=True):
    """Run multiple browser contexts with persistent states."""
    # Ensure base directories exist
    BASE_PROFILE_DIR.mkdir(parents=True, exist_ok=True)
    BASE_STATE_DIR.mkdir(parents=True, exist_ok=True)

    # Define pre-set extensions to be copied to all contexts
    EXTENSION_PATHS = [
        Path("./my_extension1"),  # Example: use forward slashes for cross-platform compatibility
        Path("./my_extension2")   # Replace with actual paths to unpacked extension directories
    ]

    # Create simulated browser with the chosen mode and extensions
    persistent_browser = PersistentBrowser(playwright, BASE_PROFILE_DIR, use_persistent_context, EXTENSION_PATHS)
    await persistent_browser.launch()

    # Create tasks for multiple contexts
    tasks = []
    for _ in range(num_contexts):
        context_id = str(uuid.uuid4())  # Unique ID for each context
        tasks.append(run_context(persistent_browser, context_id, EXTENSION_PATHS))

    # Run all contexts concurrently
    await asyncio.gather(*tasks)

    while True:
       asyncio.sleep(0.1)
    # Close the simulated browser, saving states and backing up profiles
    await persistent_browser.close()

async def main():
    """Main entry point to test both modes."""
    """
    HOW TO USE
    ----------
    This script provides a unified solution for managing multiple Playwright browser contexts with persistent states (session data, extensions, cache) using two approaches:
    1. launch_persistent_context mode: Uses Playwright's launch_persistent_context() for native persistence via user data directories.
    2. browser.launch mode: Uses Playwright's browser.launch() with manual persistence via storage state files and profile directories.

    Prerequisites:
    - Install Playwright: Run "pip install playwright" in your terminal.
    - Install Chromium: Run "playwright install chromium" to download the Chromium browser used by Playwright.
    - Ensure Python 3.7+ is installed to support async/await syntax.

    Setting Up Extensions:
    - Update the EXTENSION_PATHS list in the run() function with paths to unpacked extension directories, each containing a manifest.json file.
    - Example:
        EXTENSION_PATHS = [
            Path("./my_extension1"),
            Path("./my_extension2")
        ]
    - To obtain unpacked extensions:
        - On Linux/macOS: Extract a .crx file using "unzip <extension>.crx".
        - On Windows: Use 7-Zip to extract a .crx file (it is a ZIP archive).
    - Alternatively, pre-install extensions in a Chromium browser and copy its profile directory (e.g., ~/.config/chromium on Linux, ~/Library/Application Support/Chromium on macOS, %LocalAppData%\\Chromium\\User Data on Windows) to chromium_profile_<context_id> for each context.
    - Note: Extensions load natively in launch_persistent_context mode but may not load dynamically in browser.launch mode without additional setup (e.g., automating chrome://extensions/).

    Running the Script:
    - Save the script as script.py.
    - Run it from the terminal: "python script.py".
    - The script will:
        - Create chromium_profiles and browser_states directories in the script's working directory.
        - Run two tests:
            1. launch_persistent_context mode: Creates 3 contexts, each with a unique user data directory (chromium_profile_<context_id>) and storage state file (browser_state_<context_id>.json).
            2. browser.launch mode: Creates 3 contexts using a single browser instance, with manual persistence.
        - For each context:
            - Copies all extensions from EXTENSION_PATHS to the context's profile directory (chromium_profile_<context_id>/Extensions).
            - Loads session data (cookies, local storage) from browser_state_<context_id>.json if it exists.
            - Navigates to https://example.com and logging.infos the page title as a sample action.
            - Saves session data and backs up the profile directory (in launch_persistent_context mode) when closing.
    - Output will include logs for directory creation, extension copying, state loading/saving, and page titles.

    Subsequent Runs:
    - The script restores profile directories and session states for each context using unique context IDs (UUIDs), ensuring persistence across runs.
    - Backed-up profile directories (chromium_profile_backup_<context_id>) are used to restore chromium_profile_<context_id> if needed.
    - Session data persists in browser_state_<context_id>.json files.

    Customization:
    - Change the number of contexts by modifying num_contexts in run().
    - Switch modes by setting use_persistent_context=True (launch_persistent_context) or False (browser.launch) in run().
    - Update EXTENSION_PATHS with your extension directories.
    - Add context options (e.g., user_agent, locale) by passing them to new_context(**kwargs) in run_context().

    DEBUGGING TIPS
    --------------
    If you encounter issues while running the script, consider the following debugging strategies:

    Extension Not Loading:
    - Verify each path in EXTENSION_PATHS points to a valid unpacked extension directory containing a manifest.json file.
    - Check the chromium_profile_<context_id>/Extensions directory to ensure extensions were copied correctly.
    - In launch_persistent_context mode:
        - Open the browser (since headless=False) and navigate to chrome://extensions/ to confirm extensions are loaded.
        - Ensure extensions are compatible with Chromium and have a valid manifest.
    - In browser.launch mode:
        - Extensions may not load dynamically due to Playwright's limitations.
        - To enable extensions, pre-install them in a Chromium browser and copy the profile to chromium_profile_<context_id>.
        - Alternatively, automate extension installation by navigating to chrome://extensions/ and enabling developer mode (not implemented in this script).
    - If extensions fail to load, check console logs for errors related to install_extension().

    Session Data Issues:
    - Ensure browser_state_<context_id>.json files in the browser_states directory contain valid JSON.
    - Open these files to verify cookies and local storage data are correctly formatted.
    - Check console output for messages like "Loaded browser state..." or "Saved browser state..." to confirm state handling.
    - If session data does not persist, verify write permissions for the browser_states directory.
    - Test by logging into a website in one run and checking if the login persists in the next run.

    Backup Errors:
    - The copytree_ignore_errors function prevents errors related to SingletonLock, SingletonCookie, and SingletonSocket files by ignoring them during backups.
    - If backup issues persist, check console logs for messages like "Skipped problematic file..." or "Non-critical errors during copy...".
    - Ensure no other processes (e.g., lingering Chromium instances) are accessing chromium_profile_<context_id> directories.
    - Verify write permissions for the chromium_profiles directory.
    - If backups fail, manually inspect chromium_profile_backup_<context_id> to ensure critical files (e.g., extensions, cache) were copied.

    Indentation Errors:
    - The code uses consistent 4-space indentation per PEP 8. If syntax errors occur, ensure no tabs are mixed with spaces.
    - Copy the code directly to avoid indentation issues introduced by text editors or copy-paste operations.
    - Use an editor with Python linting (e.g., VS Code with Pylance) to detect indentation problems.

    Resource Usage:
    - In launch_persistent_context mode, each context runs a separate Chromium instance, increasing memory and CPU usage. Monitor system resources if using many contexts.
    - In browser.launch mode, a single browser instance is used, but extension/cache persistence is limited.
    - If performance is an issue, reduce num_contexts or switch to browser.launch mode.

    Other Errors:
    - If you encounter runtime errors (e.g., network issues, Playwright exceptions), check the full traceback logging.infoed by the script.
    - Share the error message and traceback for specific debugging assistance.
    - Common issues include:
        - Network errors: Ensure internet connectivity and that https://example.com is accessible.
        - Playwright errors: Verify Playwright and Chromium are installed correctly.
        - File permission errors: Run the script with appropriate permissions (e.g., as administrator on Windows if needed).
    - To isolate issues, try running with num_contexts=1 and a single extension.

    Logging:
    - The script includes verbose logging for key actions (e.g., directory creation, extension copying, state loading/saving).
    - Review console output to trace execution and identify where errors occur.
    - Add custom logging.info statements in functions like new_context() or install_extension() for deeper debugging.

    Testing:
    - Test with a small number of contexts (e.g., num_contexts=1) to verify basic functionality.
    - Test each mode separately by commenting out one call in main().
    - Test with and without extensions to isolate extension-related issues.
    - Run on different platforms (Linux, macOS, Windows) to confirm cross-platform compatibility.
    """
    async with async_playwright() as playwright:
        try:
            # Test with launch_persistent_context mode
            logging.info("Running with launch_persistent_context...")
            await run(playwright, num_contexts=3, use_persistent_context=True)
            
            # Test with browser.launch mode
            logging.info("\nRunning with browser.launch...")
            await run(playwright, num_contexts=3, use_persistent_context=False)
        except Exception as e:
            logging.info(f"An error occurred: {e}")

if __name__ == "__main__":
    asyncio.run(main())
