#!/usr/bin/env python3
"""SPA 渲染抓取 + 数字提取验证脚本

策略：
1. 用 Playwright 启动 headless chromium
2. 打开目标 SPA 页面，等待关键 DOM 元素（图表 / SVG / canvas）
3. 等待 JS hydration 完成 + 服务包络线散点渲染
4. 从渲染后的 DOM 中提取数字
5. 失败时尝试：
   - 等待更久（lazy load）
   - 触发 lazy event（window.scrollTo / IntersectionObserver）
   - 截图作 fallback
"""
import json
import re
import sys
import time
from pathlib import Path

from playwright.sync_api import sync_playwright, Page, Browser

# Puppeteer 启动参数参考 CLAUDE.md §1.2
LAUNCH_ARGS = [
    '--no-sandbox',
    '--disable-setuid-sandbox',
    '--disable-dev-shm-usage',
]

TARGETS = [
    {
        'name': 'MiniMax M3 vs H200 / GB300 (sub-page)',
        'url': 'https://inferencex.semianalysis.com/zh/compare/minimax-m3-gb300-vs-h200/agentic',
        'expect': 'TCO table + SLO sweep cost numbers',
    },
    {
        'name': 'DeepSeek V4 Pro B200 vs MI355X (sub-page)',
        'url': 'https://inferencex.semianalysis.com/zh/compare/deepseek-v4-b200-vs-mi355x/agentic',
        'expect': 'B200 $0.007 vs MI355X $0.011 comparison',
    },
    {
        'name': 'Qwen 3.5 B200 vs MI355X (sub-page)',
        'url': 'https://inferencex.semianalysis.com/zh/compare/qwen-3-5-b200-vs-mi355x/agentic',
        'expect': 'B200 FP8 $0.018 vs MI355X FP8 $0.008 (-58%)',
    },
    {
        'name': 'Overview (zh) — main dashboard',
        'url': 'https://inferencex.semianalysis.com/zh/overview',
        'expect': 'Full table with all model × hardware cells',
    },
    {
        'name': 'PR #3274 unofficial run',
        'url': 'https://inferencex.semianalysis.com/inference?unofficialRun=35434571418',
        'expect': 'DeepSeek V4 Pro B200 vLLM AgentX chart',
    },
]


def wait_for_spa_data(page: Page, max_wait_ms=15000):
    """等待 SPA 完成数据加载：等待 service envelope 等关键 DOM 出现"""
    selectors = [
        'svg .point',          # d3 散点
        'svg circle',          # d3 circles
        'svg .datum',          # recharts 等
        '[class*="envelope"]',
        '[class*="point"]',
        '[class*="scatter"]',
        'td[class*="cell"]',
        'table tbody tr',
        'button[aria-label]',  # data chip
    ]
    end = time.time() + max_wait_ms / 1000.0
    found = set()
    while time.time() < end:
        # 滚动触发 lazy load
        page.evaluate('window.scrollTo(0, document.body.scrollHeight)')
        time.sleep(0.3)
        for sel in selectors:
            try:
                if page.locator(sel).count() > 0:
                    found.add(sel)
            except Exception:
                pass
        if len(found) >= 2:
            return found, time.time()
    return found, time.time()


def extract_numbers_from_dom(page: Page, name: str, debug_dir: Path):
    """从已渲染的 DOM 提取所有数字 / 货币 / 百分比"""
    text = page.evaluate('document.body.innerText')
    numbers = []

    # 货币 $X.XXX  /  $X  /  $X.X
    for m in re.finditer(r'\$\s*([\d.]+)\s*(?:/M\s*tok(?:ens)?)?', text, re.I):
        numbers.append(('cost', m.group(0)))

    # 百分比  ±XX%
    for m in re.finditer(r'[+\-−]\s*(\d+)%', text):
        numbers.append(('pct', m.group(0)))

    # tok/s/user 与 tok/s/chip
    for m in re.finditer(r'(\d+)\s*tok/s(?:/user|/chip)?', text, re.I):
        numbers.append(('tok-rate', m.group(0)))

    # N 个数据点
    for m in re.finditer(r'(\d+)\s*个数据点', text):
        numbers.append(('data-points', m.group(0)))

    # 并发范围
    for m in re.finditer(r'并发\s*(\d+)[–\-~](\d+)', text):
        numbers.append(('conc-range', f'{m.group(1)}-{m.group(2)}'))

    # $/chip/hr
    for m in re.finditer(r'\$([\d.]+)/chip/hr', text):
        numbers.append(('tco', m.group(0)))

    # SLO 数字 126/169/212
    for k in ['126', '169', '212']:
        if k in text:
            numbers.append(('slo', k + ' tok/s/user'))

    # 保存截图
    debug_dir.mkdir(exist_ok=True, parents=True)
    safe = re.sub(r'[^\w\-]', '_', name)[:60]
    screenshot = debug_dir / f'{safe}.png'
    page.screenshot(path=str(screenshot), full_page=True)
    return numbers, screenshot, text[:2000]


def main():
    debug_dir = Path('/home/ubuntu/spa_screenshots')
    summary = []

    with sync_playwright() as p:
        # 用系统 google-chrome 而非 Playwright 自带 chromium（已知 chromium_headless_shell 缺失）
        browser: Browser = p.chromium.launch(
            headless=True,
            executable_path='/opt/google/chrome/chrome',
            args=LAUNCH_ARGS,
        )
        context = browser.new_context(
            viewport={'width': 1600, 'height': 1200},
            user_agent='Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/146.0.0.0 Safari/537.36',
        )
        page = context.new_page()

        for target in TARGETS:
            print(f"\n{'='*70}\n[{target['name']}] {target['url']}")
            try:
                page.goto(target['url'], wait_until='domcontentloaded', timeout=30000)
                print(f"  期望: {target['expect']}")
                print('  等待 SPA hydration...')
                selectors, _ = wait_for_spa_data(page, max_wait_ms=20000)
                print(f"  找到 selectors: {selectors or '(无)'}")
                time.sleep(2)  # 让图表 canvas / svg 完成绘制

                numbers, screenshot, sample_text = extract_numbers_from_dom(page, target['name'], debug_dir)
                print(f"  数字提取: {len(numbers)} 条")
                for cat, val in numbers[:30]:
                    print(f"    [{cat}] {val}")

                # 文本预览
                print(f"  文本前 800 字: {sample_text[:800]!r}")
                print(f"  截图保存: {screenshot}")

                summary.append({
                    'name': target['name'],
                    'url': target['url'],
                    'numbers_found': len(numbers),
                    'sample_numbers': numbers[:10],
                    'screenshot': str(screenshot),
                    'selectors_hit': list(selectors),
                })
            except Exception as e:
                print(f"  ERROR: {e}")
                summary.append({'name': target['name'], 'url': target['url'], 'error': str(e)})

        browser.close()

    Path('/home/ubuntu/spa_summary.json').write_text(
        json.dumps(summary, ensure_ascii=False, indent=2, default=str)
    )
    print(f"\n摘要已保存: /home/ubuntu/spa_summary.json")


if __name__ == '__main__':
    main()