"""Playwright 页面抓取工具。

对外暴露三个动作：
- open_page(url, wait_ms): 打开页面，返回 page_id
- query(page_id, selector): 选中节点并提取文字、图片 URL 与属性
- close_page(page_id): 关闭页面，释放浏览器上下文

实现上内部常驻一个 Browser；每次 open_page 创建独立 BrowserContext 和 Page。
Playwright 用 async API 跑在一个后台事件循环线程里，多个调用方线程可以真正并发
（页面加载彼此重叠，不会互相排队）；外层 API 仍保持同步。
"""

from __future__ import annotations

import asyncio
import atexit
import threading
import uuid
from collections.abc import Coroutine
from dataclasses import dataclass
from typing import Any

from playwright.async_api import Browser, BrowserContext, Page, async_playwright

# 页面打开超时（毫秒）
_GOTO_TIMEOUT = 30000

# Chromium 启动参数：
#   --no-sandbox           容器内以 root 运行时 Chrome 沙箱不可用，不加会直接启动失败
#   --disable-dev-shm-usage 容器 /dev/shm 偏小时改用磁盘，避免渲染进程崩溃
_LAUNCH_ARGS = ["--no-sandbox", "--disable-dev-shm-usage"]

# 浏览器上下文固定参数：不显式指定的话，UA 会跟着宿主系统变
# （Windows 开发机是 Windows UA，Linux 容器是 X11 UA），目标站会按 UA 返回不同页面模板，
# 导致同一套选择器在容器里全部落空
_CONTEXT_OPTIONS: dict[str, Any] = {
    "user_agent": (
        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
        "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
    ),
    "viewport": {"width": 1920, "height": 1080},
    "locale": "zh-CN",
    "timezone_id": "Asia/Shanghai",
}


@dataclass
class NodeSnapshot:
    """页面节点快照。"""

    tag: str
    text: str
    image_urls: list[str]
    attrs: dict[str, str]


@dataclass
class _PageSession:
    context: BrowserContext
    page: Page


# 节点提取脚本：取标签名、可见文字、图片地址（含懒加载）与全部属性
_EXTRACT_JS = r"""(elements) => elements.map((el) => {
    const attrs = {};
    for (const attr of Array.from(el.attributes || [])) {
        attrs[attr.name] = attr.value;
    }

    const imageUrls = [];
    const addImageUrl = (url) => {
        if (!url) return;
        try {
            const absUrl = new URL(url, location.href).href;
            if (
                absUrl.startsWith('data:') ||
                absUrl.startsWith('blob:') ||
                absUrl === 'about:blank'
            ) {
                return;
            }
            imageUrls.push(absUrl);
        } catch {
            if (
                url.startsWith('data:') ||
                url.startsWith('blob:') ||
                url === 'about:blank'
            ) {
                return;
            }
            imageUrls.push(url);
        }
    };

    const resolveImageUrl = (img) => {
        const candidates = [
            img.getAttribute('data-src'),
            img.getAttribute('data-original'),
            img.getAttribute('data-croporisrc'),
            img.currentSrc,
            img.src,
            img.getAttribute('src'),
        ];

        const srcset = img.getAttribute('srcset');
        if (srcset) {
            for (const part of srcset.split(',')) {
                const candidate = part.trim().split(/\s+/)[0];
                candidates.push(candidate);
            }
        }

        for (const candidate of candidates) {
            if (!candidate) continue;
            if (
                candidate.startsWith('data:') ||
                candidate.startsWith('blob:') ||
                candidate === 'about:blank'
            ) {
                continue;
            }
            try {
                const absUrl = new URL(candidate, location.href).href;
                if (
                    absUrl.startsWith('data:') ||
                    absUrl.startsWith('blob:') ||
                    absUrl === 'about:blank'
                ) {
                    continue;
                }
                return absUrl;
            } catch {
                return candidate;
            }
        }

        return null;
    };

    if (el.tagName === 'IMG') {
        addImageUrl(resolveImageUrl(el));
    }

    for (const img of Array.from(el.querySelectorAll('img'))) {
        addImageUrl(resolveImageUrl(img));
    }

    const text = (el.innerText || el.textContent || '').trim();

    return {
        tag: (el.tagName || '').toLowerCase(),
        text,
        image_urls: Array.from(new Set(imageUrls)),
        attrs,
    };
})"""


class _BrowserRuntime:
    """在后台事件循环线程中管理 Playwright 生命周期。

    所有 Playwright 调用都以协程形式提交到同一个事件循环，
    因此多个业务线程的页面操作可以并发交叠执行。
    """

    def __init__(self) -> None:
        self._thread: threading.Thread | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        self._launch_lock: asyncio.Lock | None = None
        self._ready = threading.Event()
        self._lock = threading.RLock()
        self._playwright: Any = None
        self._browser: Browser | None = None
        self._pages: dict[str, _PageSession] = {}
        atexit.register(self.close)

    # ──────────── 事件循环线程 ────────────

    def _thread_main(self) -> None:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        self._loop = loop
        # 在循环所属线程内创建，避免多协程各自新建锁
        self._launch_lock = asyncio.Lock()
        self._ready.set()
        try:
            loop.run_forever()
        finally:
            loop.close()

    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
        with self._lock:
            if self._thread is None or not self._thread.is_alive():
                self._ready.clear()
                self._thread = threading.Thread(
                    target=self._thread_main,
                    name="offerpie-playwright",
                    daemon=True,
                )
                self._thread.start()
                self._ready.wait()

        loop = self._loop
        if loop is None:
            raise RuntimeError("Browser runtime loop is not ready.")
        return loop

    def _submit(self, coro: Coroutine[Any, Any, Any]) -> Any:
        """把协程提交到事件循环并同步等待结果。"""
        loop = self._ensure_loop()
        return asyncio.run_coroutine_threadsafe(coro, loop).result()

    # ──────────── Playwright 操作 ────────────

    async def _ensure_browser(self) -> Browser:
        if self._browser is not None:
            return self._browser

        assert self._launch_lock is not None
        async with self._launch_lock:
            if self._browser is None:
                self._playwright = await async_playwright().start()
                self._browser = await self._playwright.chromium.launch(
                    headless=True, args=_LAUNCH_ARGS
                )
        return self._browser

    async def _open_page_impl(self, url: str, wait_ms: int) -> str:
        browser = await self._ensure_browser()
        context = await browser.new_context(**_CONTEXT_OPTIONS)
        page = await context.new_page()

        try:
            await page.goto(url, wait_until="domcontentloaded", timeout=_GOTO_TIMEOUT)
            if wait_ms > 0:
                await page.wait_for_timeout(wait_ms)
        except BaseException:
            # 打开失败也要释放上下文，避免泄漏
            await self._dispose(context, page)
            raise

        page_id = uuid.uuid4().hex
        self._pages[page_id] = _PageSession(context=context, page=page)
        return page_id

    async def _query_impl(self, page_id: str, selector: str) -> list[NodeSnapshot]:
        session = self._pages.get(page_id)
        if session is None:
            raise KeyError(f"Unknown page_id: {page_id}")

        locator = session.page.locator(selector)
        payload = await locator.evaluate_all(_EXTRACT_JS)
        return [
            NodeSnapshot(
                tag=item.get("tag", ""),
                text=item.get("text", ""),
                image_urls=list(item.get("image_urls", [])),
                attrs=dict(item.get("attrs", {})),
            )
            for item in payload
        ]

    @staticmethod
    async def _dispose(context: BrowserContext, page: Page | None) -> None:
        if page is not None:
            try:
                await page.close()
            except Exception:
                pass
        try:
            await context.close()
        except Exception:
            pass

    async def _close_page_impl(self, page_id: str) -> None:
        session = self._pages.pop(page_id, None)
        if session is None:
            return
        await self._dispose(session.context, session.page)

    async def _shutdown_impl(self) -> None:
        for page_id in list(self._pages):
            await self._close_page_impl(page_id)

        if self._browser is not None:
            try:
                await self._browser.close()
            except Exception:
                pass
            self._browser = None

        if self._playwright is not None:
            try:
                await self._playwright.stop()
            except Exception:
                pass
            self._playwright = None

    # ──────────── 生命周期 ────────────

    def close(self) -> None:
        """关闭全部页面与浏览器，停止事件循环。"""
        with self._lock:
            thread = self._thread
            loop = self._loop
            if thread is None or not thread.is_alive() or loop is None:
                return

            try:
                asyncio.run_coroutine_threadsafe(
                    self._shutdown_impl(), loop
                ).result(timeout=15)
            except Exception:
                pass

            loop.call_soon_threadsafe(loop.stop)
            thread.join(timeout=10)
            self._thread = None
            self._loop = None
            self._launch_lock = None


_runtime = _BrowserRuntime()


def open_page(url: str, wait_ms: int = 8000) -> str:
    """打开页面并返回 page_id。"""
    return _runtime._submit(_runtime._open_page_impl(url, wait_ms))


def query(page_id: str, selector: str) -> list[NodeSnapshot]:
    """按选择器抓取节点，提取文本、图片 URL 与属性。"""
    items = _runtime._submit(_runtime._query_impl(page_id, selector))
    seen_urls: set[str] = set()

    for item in items:
        unique_urls: list[str] = []
        for url in item.image_urls:
            if url in seen_urls:
                continue
            seen_urls.add(url)
            unique_urls.append(url)
        item.image_urls = unique_urls

    return items


def close_page(page_id: str) -> None:
    """关闭页面，释放对应的浏览器上下文。"""
    _runtime._submit(_runtime._close_page_impl(page_id))