gyxx-flow/src/gyxx_flow/adapters/browser.py

"""Portable, cross-process browser profile ownership."""

from __future__ import annotations

import re
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator
from urllib.parse import urlparse

from gyxx_flow.core.artifacts import atomic_write_json
from gyxx_flow.core.layout import DataLayout
from gyxx_flow.core.locks import LockManager

_PROFILE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")


@dataclass(frozen=True, slots=True)
class BrowserProfileLease:
    profile_id: str
    profile_path: Path


class BrowserProfileManager:
    def __init__(
        self, data_root: Path, *, lock_manager: LockManager | None = None
    ) -> None:
        self._layout = DataLayout(data_root)
        self._locks = lock_manager or LockManager(
            self._layout.root / "state" / "locks"
        )

    @contextmanager
    def acquire(
        self,
        profile_id: str,
        *,
        owner: str,
        timeout_seconds: float = 30.0,
    ) -> Iterator[BrowserProfileLease]:
        if not isinstance(profile_id, str) or not _PROFILE_ID.fullmatch(profile_id):
            raise ValueError("invalid browser profile id")
        path = self._layout.state("browser_profiles", profile_id)
        with self._locks.acquire(
            f"browser:{profile_id}", owner=owner, timeout_seconds=timeout_seconds
        ):
            path.mkdir(parents=True, exist_ok=True)
            yield BrowserProfileLease(profile_id, path)


@dataclass(frozen=True, slots=True)
class BrowserCookieStore:
    """Playwright-compatible cookie/storage state persisted with atomic replace."""

    cookie_file: Path
    storage_state_file: Path

    def __post_init__(self) -> None:
        object.__setattr__(self, "cookie_file", Path(self.cookie_file).expanduser().resolve())
        object.__setattr__(
            self,
            "storage_state_file",
            Path(self.storage_state_file).expanduser().resolve(),
        )

    def load_cookies(self) -> list[dict[str, Any]]:
        payload = self._read(self.cookie_file, default=[])
        # A few legacy collectors persisted ``{"cookies": [...], ...}``
        # envelopes instead of the canonical Playwright list.  Read those
        # envelopes compatibly, while all new writes remain canonical lists.
        if isinstance(payload, dict) and isinstance(payload.get("cookies"), list):
            payload = payload["cookies"]
        if not isinstance(payload, list) or not all(isinstance(item, dict) for item in payload):
            raise ValueError("browser cookie file must contain a list of objects")
        return payload

    def save_cookies(self, cookies: list[dict[str, Any]]) -> Path:
        if not isinstance(cookies, list) or not all(isinstance(item, dict) for item in cookies):
            raise ValueError("browser cookies must be a list of objects")
        return atomic_write_json(self.cookie_file, cookies)

    def load_storage_state(self) -> dict[str, Any] | None:
        payload = self._read(self.storage_state_file, default=None)
        if payload is not None and not isinstance(payload, dict):
            raise ValueError("browser storage state must contain an object")
        return payload

    def save_storage_state(self, state: dict[str, Any]) -> Path:
        if not isinstance(state, dict):
            raise ValueError("browser storage state must be an object")
        return atomic_write_json(self.storage_state_file, state)

    @staticmethod
    def _read(path: Path, *, default: Any) -> Any:
        if not path.exists():
            return default
        import json

        try:
            return json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            raise ValueError(f"cannot read browser state file: {path.name}") from exc


def restore_browser_state(
    context: Any,
    *,
    cookie_file: str | Path | None = None,
    storage_state_file: str | Path | None = None,
    restore_origins: bool = True,
) -> None:
    """Restore an account or binding snapshot into an existing context.

    Persistent Scrapling contexts keep their own Profile, while account vaults
    are the shared login authority.  Direct Scrapling scripts therefore need a
    small explicit restore step before their first authenticated navigation.
    """

    store = BrowserCookieStore(
        Path(cookie_file) if cookie_file else Path(),
        Path(storage_state_file) if storage_state_file else Path(),
    ) if cookie_file and storage_state_file else None
    if store is None:
        return

    state = store.load_storage_state() if store.storage_state_file.exists() else None
    setter = getattr(context, "set_storage_state", None)
    adder = getattr(context, "add_cookies", None)
    if state and callable(setter):
        state_to_restore = state
        if not restore_origins:
            state_to_restore = {
                "cookies": list(state.get("cookies", []))
                if isinstance(state.get("cookies"), list)
                else []
            }
        setter(state_to_restore)
    elif state and callable(adder) and isinstance(state.get("cookies"), list):
        adder(state["cookies"])

    cookies = store.load_cookies() if store.cookie_file.exists() else []
    if cookies and callable(adder):
        adder(cookies)


def persist_browser_state(
    context: Any,
    *,
    cookie_file: str | Path,
    storage_state_file: str | Path,
) -> int:
    """Persist a verified context snapshot and return its cookie count."""

    cookies = context.cookies()
    storage_state = context.storage_state()
    store = BrowserCookieStore(Path(cookie_file), Path(storage_state_file))
    store.save_cookies(cookies)
    store.save_storage_state(storage_state)
    return len(cookies)


def looks_like_login_url(url: str) -> bool:
    """Whether a page URL looks like a platform login/passport redirect."""

    parsed = urlparse(str(url))
    hostname = (parsed.hostname or "").casefold()
    path = parsed.path.casefold()
    markers = ("login", "passport", "signin", "sign-in")
    return any(marker in hostname or marker in path for marker in markers)