playwright_transcript.py

from __future__ import annotations
import asyncio
import html
import os
import re
import shutil
from urllib.parse import quote

try:
    from playwright.async_api import async_playwright, TimeoutError as PlaywrightTimeoutError
except Exception:
    async_playwright = None
    PlaywrightTimeoutError = Exception


def _clean(text):
    return re.sub(r"\s+", " ", html.unescape(text or "")).strip()


def _timestamp(text):
    parts = _clean(text).split(":")
    try:
        if len(parts) == 2:
            return int(parts[0]) * 60 + float(parts[1])
        if len(parts) == 3:
            return int(parts[0]) * 3600 + int(parts[1]) * 60 + float(parts[2])
    except Exception:
        pass
    return 0.0


async def _click_transcript(page):
    candidates = [
        page.get_by_role("button", name=re.compile(r"show transcript|transcript|bản chép lời", re.I)).first,
        page.get_by_text(re.compile(r"^show transcript$|^transcript$|hiện bản chép lời|hiển thị bản chép lời", re.I)).first,
        page.locator("ytd-video-description-transcript-section-renderer button").first,
        page.locator("ytd-video-description-transcript-section-renderer tp-yt-paper-button").first,
        page.locator("button[aria-label*='transcript' i]").first,
    ]
    for loc in candidates:
        try:
            if await loc.count() and await loc.is_visible(timeout=1200):
                await loc.scroll_into_view_if_needed(timeout=2000)
                await loc.click(timeout=5000)
                return True
        except Exception:
            pass
    return False


async def _read_segments(page):
    loc = page.locator("ytd-transcript-segment-renderer")
    count = await loc.count()
    out = []
    for i in range(count):
        item = loc.nth(i)
        try:
            text = _clean(await item.inner_text(timeout=2000))
        except Exception:
            continue
        if not text:
            continue
        try:
            ts = await item.locator(".segment-timestamp").inner_text(timeout=700)
        except Exception:
            try:
                ts = await item.locator("div[class*='timestamp']").first.inner_text(timeout=700)
            except Exception:
                ts = ""
        out.append((text, _timestamp(ts)))

    fixed = []
    for i, (text, start) in enumerate(out):
        end = out[i + 1][1] if i + 1 < len(out) and out[i + 1][1] > start else start + 2.0
        fixed.append((text, start, end))
    return fixed


def _chromium_path():
    # HF Spaces installs apt packages from packages.txt. Prefer system Chromium
    # so Playwright does not depend on the per-user browser cache.
    configured = os.getenv("PLAYWRIGHT_EXECUTABLE_PATH", "").strip()
    candidates = [configured, "/usr/bin/chromium", "/usr/bin/chromium-browser", shutil.which("chromium"), shutil.which("chromium-browser")]
    for path in candidates:
        if path and os.path.exists(path):
            return path
    return None


async def _fetch_transcript_async(video_id, timeout_ms=45000):
    if not re.fullmatch(r"[A-Za-z0-9_-]{11}", str(video_id or "")):
        return [], "Playwright: video ID không hợp lệ"
    if async_playwright is None:
        return [], "Playwright chưa được cài đặt"

    pw = await async_playwright().start()
    browser = None
    context = None
    try:
        launch_kwargs = dict(
            headless=True,
            args=[
                "--no-sandbox",
                "--disable-dev-shm-usage",
                "--disable-gpu",
                "--disable-blink-features=AutomationControlled",
            ],
        )
        chromium = _chromium_path()
        if chromium:
            launch_kwargs["executable_path"] = chromium
        browser = await pw.chromium.launch(**launch_kwargs)
        context = await browser.new_context(
            locale="en-US",
            viewport={"width": 1280, "height": 900},
            user_agent="Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
        )
        page = await context.new_page()
        url = f"https://www.youtube.com/watch?v={quote(video_id)}&hl=en&gl=US"
        await page.goto(url, wait_until="domcontentloaded", timeout=timeout_ms)

        for name in ("Accept all", "I agree", "Reject all"):
            try:
                btn = page.get_by_role("button", name=re.compile(re.escape(name), re.I)).first
                if await btn.count() and await btn.is_visible(timeout=1000):
                    await btn.click(timeout=2500)
                    await page.wait_for_timeout(700)
                    break
            except Exception:
                pass

        await page.wait_for_timeout(2500)
        opened = await _click_transcript(page)
        for _ in range(4):
            if opened:
                break
            await page.mouse.wheel(0, 1600)
            await page.wait_for_timeout(1200)
            opened = await _click_transcript(page)

        if not opened:
            return [], "Playwright: không tìm thấy nút Show transcript trên YouTube"

        segments = page.locator("ytd-transcript-segment-renderer")
        try:
            await segments.first.wait_for(timeout=15000)
        except Exception:
            pass

        previous = -1
        stable = 0
        for _ in range(20):
            current = await segments.count()
            if current == previous:
                stable += 1
            else:
                stable = 0
            previous = current
            try:
                panel = page.locator("ytd-transcript-renderer").first
                await panel.evaluate("el => { const e = el.querySelector('#segments-container'); if (e) e.scrollTop = e.scrollHeight; else el.scrollTop = el.scrollHeight; }")
            except Exception:
                await page.mouse.wheel(0, 1200)
            await page.wait_for_timeout(400)
            if stable >= 4:
                break

        pieces = await _read_segments(page)
        if pieces:
            return pieces, f"YouTube UI / Playwright Async · {len(pieces)} đoạn"
        return [], "Playwright: transcript panel mở nhưng không có segment"
    except PlaywrightTimeoutError as e:
        return [], f"Playwright timeout: {e}"
    except Exception as e:
        return [], f"Playwright {type(e).__name__}: {e}"
    finally:
        try:
            if context:
                await context.close()
        except Exception:
            pass
        try:
            if browser:
                await browser.close()
        except Exception:
            pass
        try:
            await pw.stop()
        except Exception:
            pass


def fetch_transcript(video_id, timeout_ms=45000):
    """Sync-compatible wrapper; runs Playwright Async in a worker thread."""
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(_fetch_transcript_async(video_id, timeout_ms))

    result = {}
    error = []

    def runner():
        try:
            result["value"] = asyncio.run(_fetch_transcript_async(video_id, timeout_ms))
        except Exception as exc:
            error.append(exc)

    import threading
    t = threading.Thread(target=runner, daemon=True)
    t.start()
    t.join(timeout=(timeout_ms / 1000.0) + 20)
    if error:
        return [], f"Playwright worker {type(error[0]).__name__}: {error[0]}"
    if "value" not in result:
        return [], "Playwright worker timeout"
    return result["value"]