from dataclasses import dataclass
from typing import Optional

from wapiti_arsenic.browsers import Browser
from wapiti_arsenic.errors import UnknownError
from wapiti_arsenic.services import Service
from wapiti_arsenic.session import Session


@dataclass
class SessionContext:
    service: Service
    browser: Browser
    bind: str
    session: Optional[Session] = None

    async def __aenter__(self):
        self.session = await start_session(self.service, self.browser, self.bind)
        return self.session

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await stop_session(self.session)


def get_session(service, browser, bind=""):
    return SessionContext(service, browser, bind)


async def start_session(service: Service, browser: Browser, bind=""):
    driver = await service.start()
    return await driver.new_session(browser, bind=bind)


async def stop_session(session: Session):
    try:
        await session.close()
    except UnknownError as e:
        if "marionette" not in str(e):
            raise
    await session.driver.close()
