import asyncio
from functools import partial
from typing import Any, Dict, Optional, Union, overload

from playwright.async_api import (
    Browser,
    BrowserContext,
    Playwright,
    PlaywrightContextManager,
)
from typing_extensions import Literal

from camoufox.virtdisplay import VirtualDisplay

from .utils import async_attach_vd, launch_options


class AsyncCamoufox(PlaywrightContextManager):
    """
    Wrapper around playwright.async_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

    async def __aenter__(self) -> Union[Browser, BrowserContext]:
        _playwright = await super().__aenter__()
        self.browser = await AsyncNewBrowser(_playwright, **self.launch_options)
        return self.browser

    async def __aexit__(self, *args: Any):
        if self.browser:
            await self.browser.close()
        await super().__aexit__(*args)


@overload
async def AsyncNewBrowser(
    playwright: Playwright,
    *,
    from_options: Optional[Dict[str, Any]] = None,
    persistent_context: Literal[False] = False,
    **kwargs,
) -> Browser: ...


@overload
async def AsyncNewBrowser(
    playwright: Playwright,
    *,
    from_options: Optional[Dict[str, Any]] = None,
    persistent_context: Literal[True],
    **kwargs,
) -> BrowserContext: ...


async def AsyncNewBrowser(
    playwright: Playwright,
    *,
    headless: Optional[Union[bool, Literal['virtual']]] = None,
    from_options: Optional[Dict[str, Any]] = None,
    persistent_context: bool = False,
    debug: Optional[bool] = None,
    **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()`.
    """
    if headless == 'virtual':
        virtual_display = VirtualDisplay(debug=debug)
        kwargs['virtual_display'] = virtual_display.get()
        headless = False
    else:
        virtual_display = None

    if not from_options:
        from_options = await asyncio.get_event_loop().run_in_executor(
            None,
            partial(launch_options, headless=headless, debug=debug, **kwargs),
        )

    # Persistent context
    if persistent_context:
        context = await playwright.firefox.launch_persistent_context(**from_options)
        return await async_attach_vd(context, virtual_display)

    # Browser
    browser = await playwright.firefox.launch(**from_options)
    return await async_attach_vd(browser, virtual_display)