from typing import Any, Dict, Optional, Union

from playwright.sync_api import (
    Browser,
    BrowserContext,
    Playwright,
    PlaywrightContextManager,
)

from .utils import launch_options


class Camoufox(PlaywrightContextManager):
    """
    Wrapper around playwright.sync_api.PlaywrightContextManager that automatically
    launches a browser and closes it when the context manager is exited.
    """

    def __init__(self, **launch_options):
        super().__init__()
        self.launch_options = launch_options
        self.browser: Optional[Union[Browser, BrowserContext]] = None

    def __enter__(self) -> Union[Browser, BrowserContext]:
        super().__enter__()
        self.browser = NewBrowser(self._playwright, **self.launch_options)
        return self.browser

    def __exit__(self, *args: Any):
        if self.browser:
            self.browser.close()
        super().__exit__(*args)


def NewBrowser(
    playwright: Playwright,
    *,
    from_options: Optional[Dict[str, Any]] = None,
    persistent_context: bool = False,
    **kwargs,
) -> Union[Browser, BrowserContext]:
    """
    Launches a new browser instance for Camoufox given a set of launch options.

    Parameters:
        from_options (Dict[str, Any]):
            A set of launch options generated by `launch_options()` to use
        persistent_context (bool):
            Whether to use a persistent context.
        **kwargs:
            All other keyword arugments passed to `launch_options()`.
    """
    opt = from_options or launch_options(**kwargs)
    if persistent_context:
        return playwright.firefox.launch_persistent_context(**opt)

    return playwright.firefox.launch(**opt)