# -*- coding: utf-8 -*-
"""
CDP 截图脚本
============
用途：截取当前 Chrome 页面截图，保存为 PNG 文件。
若当前页面包含二维码登录界面，截图后打印保存路径，
由 batch_publish.py 读取并展示给用户扫码。

输出：
  - 截图保存到 OUT_PATH（默认同目录的 qrcode_login.png）
  - stdout 打印 SCREENSHOT_SAVED:<路径>，供 batch_publish.py 解析
"""
import sys, socket, json, base64, urllib.request, os, time
sys.stdout.reconfigure(encoding='utf-8', errors='replace')

HOST = "localhost"
PORT = 9222
OUT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qrcode_login.png")

# ── WebSocket helpers（与其他脚本相同的极简实现）────────────────────────────
def get_first_page_id():
    pages = json.loads(urllib.request.urlopen(f'http://{HOST}:{PORT}/json', timeout=5).read())
    for p in pages:
        if p.get('type') == 'page':
            return p['id']
    return None

def ws_connect(page_id):
    import urllib.parse
    ws_url = f"ws://{HOST}:{PORT}/devtools/page/{page_id}"
    parsed = urllib.parse.urlparse(ws_url)
    sock = socket.create_connection((parsed.hostname, parsed.port or 80), timeout=10)
    key = 'dGhlIHNhbXBsZSBub25jZQ=='
    hs = (
        f'GET {parsed.path} HTTP/1.1\r\n'
        f'Host: {parsed.hostname}:{parsed.port}\r\n'
        'Upgrade: websocket\r\nConnection: Upgrade\r\n'
        f'Sec-WebSocket-Key: {key}\r\nSec-WebSocket-Version: 13\r\n\r\n'
    )
    sock.send(hs.encode())
    resp = b''
    while b'\r\n\r\n' not in resp:
        resp += sock.recv(4096)
    return sock

def send_frame(sock, data):
    data = data.encode('utf-8')
    n = len(data)
    mask = b'\x01\x02\x03\x04'
    if n <= 125:
        hdr = bytes([0x81, n | 0x80]) + mask
    elif n <= 65535:
        hdr = bytes([0x81, 0xFE]) + n.to_bytes(2, 'big') + mask
    else:
        hdr = bytes([0x81, 0xFF]) + n.to_bytes(8, 'big') + mask
    sock.send(hdr + bytes(data[i] ^ mask[i % 4] for i in range(n)))

def recv_frame(sock):
    def rx(n):
        buf = b''
        while len(buf) < n:
            buf += sock.recv(n - len(buf))
        return buf
    h = rx(2)
    n = h[1] & 0x7f
    if n == 126:
        n = int.from_bytes(rx(2), 'big')
    elif n == 127:
        n = int.from_bytes(rx(8), 'big')
    raw = rx(n)
    # 服务端帧可能有 mask（理论不常见），但先不处理，直接解码
    return raw.decode('utf-8')

_mid = 0
def cdp_call(sock, method, params=None, timeout=15):
    global _mid
    _mid += 1
    mid = _mid
    msg = {'id': mid, 'method': method}
    if params:
        msg['params'] = params
    send_frame(sock, json.dumps(msg))
    sock.settimeout(timeout)
    while True:
        raw = recv_frame(sock)
        obj = json.loads(raw)
        if obj.get('id') == mid:
            return obj.get('result', {})

# ── 主逻辑 ──────────────────────────────────────────────────────────────────
page_id = get_first_page_id()
if not page_id:
    print("[ERROR] 未找到 Chrome 页面，请先启动 Chrome 调试模式")
    sys.exit(1)

sock = ws_connect(page_id)

# 截图（format=png, quality 忽略，clip 不传则全页）
result = cdp_call(sock, 'Page.captureScreenshot', {'format': 'png'}, timeout=20)
sock.close()

data_b64 = result.get('data', '')
if not data_b64:
    print("[ERROR] 截图返回空数据")
    sys.exit(1)

img_bytes = base64.b64decode(data_b64)
with open(OUT_PATH, 'wb') as f:
    f.write(img_bytes)

size_kb = len(img_bytes) // 1024
print(f"[OK] 截图已保存（{size_kb} KB）")
print(f"SCREENSHOT_SAVED:{OUT_PATH}")
