C.E.L_Slide_test2/scripts/fetch_all_frame_screenshots.py

"""모든 Figma 프레임의 스크린샷을 번호 붙은 단일 폴더로 정리.

결과:
  data/figma_previews/01.png, 02.png, ..., 32.png
  data/figma_previews/index.json  ({number: {frame_id, node_id, title}})
"""
from __future__ import annotations

import base64
import json
import sys
from pathlib import Path
from urllib import error, request

sys.path.insert(0, str(Path(__file__).parent.parent))

from src.frame_extractor import extract_all_frames

MCP_URL = "http://127.0.0.1:3845/mcp"
OUT_DIR = Path("data/figma_previews")

# frame_id → node_id (32개, metadata에서 추출)
FRAME_NODE_MAP = {
    "1171281172": "145:8352",
    "1171281173": "182:2870",
    "1171281174": "182:2810",
    "1171281175": "182:2829",
    "1171281176": "182:3046",
    "1171281177": "182:3053",
    "1171281178": "145:8394",
    "1171281179": "182:3024",
    "1171281180": "112:87",
    "1171281181": "182:2572",
    "1171281182": "182:2523",
    "1171281189": "100:65",
    "1171281190": "51:99",
    "1171281191": "100:132",
    "1171281192": "182:2602",
    "1171281193": "106:205",
    "1171281194": "112:7",
    "1171281195": "106:252",
    "1171281197": "182:2727",
    "1171281198": "182:2766",
    "1171281201": "145:8310",
    "1171281202": "112:49",
    "1171281203": "145:8266",
    "1171281204": "145:8223",
    "1171281205": "182:2668",
    "1171281206": "182:2643",
    "1171281208": "145:8504",
    "1171281209": "145:8523",
    "1171281210": "181:2519",
    "1171281211": "181:2520",
    "1171281212": "181:2521",
    "1171281213": "181:2522",
}


def parse_sse(body: str) -> dict:
    for line in body.splitlines():
        if line.startswith("data: "):
            return json.loads(line[6:])
    raise RuntimeError(f"No data line in response: {body[:200]}")


def post(payload: dict, session_id: str | None = None) -> tuple[dict, str | None]:
    data = json.dumps(payload).encode()
    req = request.Request(
        MCP_URL,
        data=data,
        method="POST",
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        },
    )
    if session_id:
        req.add_header("mcp-session-id", session_id)
    with request.urlopen(req, timeout=60) as resp:
        body = resp.read().decode()
        sid = resp.headers.get("mcp-session-id")
        return (parse_sse(body) if body.strip() else {}, sid)


def initialize() -> str:
    payload = {
        "jsonrpc": "2.0",
        "id": 1,
        "method": "initialize",
        "params": {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "frame-dumper", "version": "1.0"},
        },
    }
    _, sid = post(payload)
    notify = {"jsonrpc": "2.0", "method": "notifications/initialized", "params": {}}
    data = json.dumps(notify).encode()
    req = request.Request(
        MCP_URL,
        data=data,
        method="POST",
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
            "mcp-session-id": sid or "",
        },
    )
    try:
        request.urlopen(req, timeout=10).read()
    except error.HTTPError:
        pass
    return sid or ""


def get_screenshot(session_id: str, node_id: str, call_id: int) -> bytes:
    payload = {
        "jsonrpc": "2.0",
        "id": call_id,
        "method": "tools/call",
        "params": {"name": "get_screenshot", "arguments": {"nodeId": node_id}},
    }
    resp, _ = post(payload, session_id=session_id)
    if "error" in resp:
        raise RuntimeError(f"MCP error for {node_id}: {resp['error']}")
    for item in resp.get("result", {}).get("content", []):
        if item.get("type") == "image":
            return base64.b64decode(item["data"])
    raise RuntimeError(f"No image in response for {node_id}")


def main() -> int:
    # frame_id → title_text 맵
    frames = extract_all_frames("figma_to_html_agent/blocks")
    title_map = {f["frame_id"]: (f.get("title_text") or "").replace("\n", " ")[:80] for f in frames}

    # 정렬된 frame_id 목록에 1부터 번호 매김
    frame_ids = sorted(FRAME_NODE_MAP.keys())

    OUT_DIR.mkdir(parents=True, exist_ok=True)

    print("[init] MCP session...")
    sid = initialize()
    print(f"[init] session-id={sid}")

    index: dict[str, dict] = {}
    for i, fid in enumerate(frame_ids, start=1):
        node_id = FRAME_NODE_MAP[fid]
        num = f"{i:02d}"
        out_path = OUT_DIR / f"{num}.png"
        title = title_map.get(fid, "")

        if out_path.exists():
            print(f"[{num}] {fid} (node {node_id}) — 이미 있음, skip")
        else:
            print(f"[{num}] {fid} (node {node_id}) fetching...")
            try:
                png = get_screenshot(sid, node_id, 100 + i)
                out_path.write_bytes(png)
                print(f"      saved {len(png)} bytes → {out_path}")
            except Exception as e:
                print(f"      FAILED: {e}")
                continue

        index[num] = {
            "frame_id": fid,
            "node_id": node_id,
            "title_text": title,
            "png": f"{num}.png",
        }

    (OUT_DIR / "index.json").write_text(
        json.dumps(index, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    print(f"\n[done] {len(index)}개 저장, index: {OUT_DIR/'index.json'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())