app/core/service/playwright/playwright_context.py

from __future__ import annotations

from typing import AsyncIterator, Awaitable, ClassVar

from playwright.async_api import (  # noqa: F401
    Browser,
    BrowserContext,
    Page,
    TimeoutError,
    async_playwright,
)

from .models import GetContentModel, PageModel, ScreenshotModel  # noqa: TCH001


class AsyncMixin:
    """Experimental: making awaitable class."""

    async def __ainit__(self) -> None:
        """Initialize the class."""

    def __await__(self) -> AsyncIterator[Awaitable]:
        """Make the class awaitable."""
        return self.__ainit__().__await__()


class PlaywrightInstance(AsyncMixin):
    """This class is designed to keep playwright browser instance open for reusability and scalability handling api requests."""  # noqa: E501

    HEADERS = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36"  # noqa: E501

    FIREFOX_USER_PREFS: ClassVar[dict[str, int | str]] = {
        "extensions.enabledScopes": 1,
        "extensions.autoDisableScopes": 1,
        "dom.webdriver.enabled": False,
        "useAutomationExtension": False,
        "general.useragent.override": HEADERS,
    }

    def __init__(self) -> None:
        """Initialize the class."""
        self.playwright: async_playwright | None = None
        self.browser: Browser | None = None

    async def __ainit__(self) -> PlaywrightInstance:
        if not self.playwright:
            self.playwright = await async_playwright().start()
            self.browser = await self.playwright.firefox.launch(
                firefox_user_prefs=self.FIREFOX_USER_PREFS,
            )
        return self

    async def new_context_page(
        self,
        browser: Browser,
        screenshot_model: GetContentModel,
        page_model: PageModel,
    ) -> tuple[BrowserContext | None, Page]:
        """create a brwoser or new browser context page.

        Parameters:
            browser (Browser):
                The Playwright Browser instance to create a new context in.
            screenshot_model (GetContentModel):
                A pydantic BaseModel instance containing the configuration for the screenshot.
            page_model (PageModel):
                A pydantic BaseModel instance containing the configuration for the page.

        Returns:
            tuple: BrowserContext and Page
        """
        params = {
            "color_scheme": page_model.color_scheme,
            "java_script_enabled": page_model.java_script_enabled,
            "no_viewport": page_model.no_viewport,
            "proxy": page_model.proxy.model_dump() if page_model.proxy else None,
            "viewport": page_model.viewport.model_dump() if page_model.viewport else None,
        }

        if not screenshot_model.new_browser:
            return None, await browser.new_page(**params)

        new_context = await browser.new_context(**params)
        return new_context, await new_context.new_page()

    async def screenshot(
        self,
        screenshot_model: ScreenshotModel,
        page_model: PageModel,
    ) -> bytes:
        """Take a screenshot of a webpage url.

        Parameters:
            screenshot_model (ScreenshotModel):
                A pydantic BaseModel instance containing the configuration for the screenshot.
            page_model (PageModel):
                A pydantic BaseModel instance containing the configuration for the page.

        Returns:
            bytes: The screenshot data in bytes.
        """
        context, page = await self.new_context_page(
            screenshot_model=screenshot_model,
            browser=self.browser,
            page_model=page_model,
        )

        await page.goto(str(screenshot_model.url))
        await page.wait_for_timeout(screenshot_model.ms_delay)

        screenshot_locator = (
            page.locator(screenshot_model.query_selector)
            if screenshot_model.query_selector
            else None
        )

        if screenshot_locator:
            screenshot_data: bytes = await screenshot_locator.screenshot(
                type=screenshot_model.image_type,
            )
        else:
            screenshot_data: bytes = await page.screenshot(
                full_page=screenshot_model.full_page,
                type=screenshot_model.image_type,
            )

        await page.close()

        if context:
            await context.close()

        return screenshot_data

    async def close_instance(self) -> None:
        """For manual closing of playwright if needed."""
        if self.playwright:
            await self.browser.close()
            await self.playwright.stop()