#!/usr/bin/env python3
"""
文档抽取脚本 — 通过 ab-render HTTP 接口实现，本地不需要任何 Office/解析依赖。

调用链：
  POST $RENDER_API_URL/extract        → { taskId }
  POST $RENDER_API_URL/extractStatus  → { status, progress, result?, error? }

两层产物（--fields）：
  content   title / text[] / rawNotes —— 正文与备注页，轻量，不含任何 base64
  layout    slide{nodes,beats,mode,…} —— 坐标形状，带内嵌图标，**体积很大**

默认**不把完整结果打到 stdout**：一页 PPT 的 layout 带上内嵌图标可能几百 KB，
直接进对话正文会撑爆上下文，还会让用户在聊天记录里滚过一屏 base64。默认写文件 +
打一份人能读的摘要；确实要管道传全量时用 --json-output（或只要 content 层）。

必要环境变量（由 ab-agent 自动注入）：
  RENDER_API_URL   ab-render 服务地址（默认 https://api-render.remixmate.ai）
  PRIV_TOKEN       ab-api 私有 token

可选环境变量：
  CONVERSATION_ID  对话 ID，用于留痕与产物归属
"""

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.request

# ─── 环境变量 ─────────────────────────────────────────────────────────────────

RENDER_API_URL = (
    os.environ.get("RENDER_API_URL")
    or os.environ.get("REMOTION_RENDER_API_URL")
    or "https://api-render.remixmate.ai"
).rstrip("/")
PRIV_TOKEN = os.environ.get("PRIV_TOKEN", "")
CONVERSATION_ID = os.environ.get("CONVERSATION_ID", "")

# ─── HTTP 工具 ────────────────────────────────────────────────────────────────


def api_post(path: str, body: dict, timeout: int = 30) -> dict:
    """向 ab-render 发送 POST 请求，返回解析后的 JSON。"""
    url = f"{RENDER_API_URL}{path}"
    data = json.dumps(body).encode()

    headers: dict[str, str] = {
        "Content-Type": "application/json",
        "X-Priv-Token": PRIV_TOKEN,
    }
    if CONVERSATION_ID:
        headers["x-conversation-id"] = CONVERSATION_ID

    req = urllib.request.Request(url, data=data, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return json.loads(resp.read())
    except urllib.error.HTTPError as e:
        body_text = e.read().decode(errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {body_text}") from e


def log(msg: str, silent: bool = False) -> None:
    if not silent:
        print(msg, flush=True)


# ─── 摘要 ─────────────────────────────────────────────────────────────────────


def summarize(result: dict, out_path: str | None) -> str:
    """
    人能读的一页纸。**这是这个 skill 存在的一半价值** —— 端点给的是全量结构，
    决定"喂多少进上下文"是入口层的责任（设计文档 §4.4）。

    每页一行：页码、档位、拍数、标题、讲稿来源。档位与拍数决定后面怎么写旁白，
    讲稿来源决定走严格改写还是扩写，这两件是下一步真正要用的信息。
    """
    src = result.get("source", {})
    lines = [
        f"共 {src.get('slideCount', '?')} 页，本次抽到 {len(result.get('slides', []))} 页"
        f"（fields={'+'.join(src.get('fields', []))}）",
        "",
    ]

    for s in result.get("slides", []):
        idx = s.get("slideIndex", "?")
        bits = [f"第 {idx} 页"]

        if s.get("mode"):
            beats = len(s.get("beats") or [])
            bits.append(f"{s['mode']}/{beats} 拍")
        if s.get("title"):
            bits.append(f"「{s['title']}」")

        notes = (s.get("rawNotes") or "").strip()
        if notes:
            head = notes.replace("\n", " ")[:40]
            bits.append(f"备注 {len(notes)} 字：{head}…")
        elif "rawNotes" in s:
            bits.append("无备注（旁白需按正文扩写）")

        lines.append("  " + " · ".join(bits))

        if s.get("modeReason"):
            lines.append(f"      判定：{s['modeReason']}")
        flags = s.get("flags") or []
        if flags:
            # flag 是显式的降级记录，不是日志 —— 画面上少一块，这里不说就没人知道为什么。
            lines.append(f"      ⚠ 降级：{'; '.join(flags[:3])}")

    if out_path:
        lines += ["", f"完整结果已写入：{out_path}"]
    return "\n".join(lines)


# ─── 主逻辑 ───────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Document extraction — calls the ab-render /extract endpoint"
    )
    parser.add_argument("--url", required=True, help="Direct document URL (.pptx)")
    parser.add_argument("--kind", default="pptx", choices=["pptx"], help="Document kind (default pptx)")
    parser.add_argument("--slide", type=int, help="Only extract this 1-based page")
    parser.add_argument(
        "--fields", default="content,layout",
        help="Comma-separated: content / layout (default both)",
    )
    parser.add_argument(
        "--out", help="Write the full JSON result here (default: ./<name>.extract.json)",
    )
    parser.add_argument(
        "--json-output", action="store_true",
        help="Pipeline mode: emit the full JSON on stdout (careful: layout is large)",
    )
    args = parser.parse_args()

    if not PRIV_TOKEN:
        print(
            "❌ PRIV_TOKEN env var is not configured — /extract requires auth.",
            file=sys.stderr,
        )
        sys.exit(1)

    fields = [f.strip() for f in args.fields.split(",") if f.strip()]

    # ── 1. Submit ────────────────────────────────────────────────────────────
    log(f"[doc-parse] Submitting extraction: {args.url[:80]}", args.json_output)

    body: dict = {"documentUrl": args.url, "kind": args.kind, "fields": fields}
    if args.slide is not None:
        body["slide"] = args.slide

    resp = api_post("/extract", body)
    if resp.get("code") != 0:
        print(f"❌ submission failed: {resp.get('msg', 'unknown error')}", file=sys.stderr)
        sys.exit(1)

    task_id: str = resp["data"]["taskId"]
    log(f"[doc-parse] task accepted, taskId={task_id}", args.json_output)

    # ── 2. Poll ──────────────────────────────────────────────────────────────
    max_wait_sec = 600
    poll_interval = 2
    elapsed = 0

    while elapsed < max_wait_sec:
        time.sleep(poll_interval)
        elapsed += poll_interval

        try:
            status_resp = api_post("/extractStatus", {"taskId": task_id})
        except Exception as e:  # noqa: BLE001
            log(f"[doc-parse] polling error (will retry): {e}", args.json_output)
            continue

        if status_resp.get("code") != 0:
            log(f"[doc-parse] polling error code (will retry): {status_resp.get('msg')}", args.json_output)
            continue

        task = status_resp["data"]
        status = task.get("status", "unknown")

        if status == "succeeded":
            result = task.get("result") or {}

            if args.json_output:
                print(json.dumps(result, ensure_ascii=False))
                return

            # 默认落盘 + 摘要：layout 层直接打进 stdout 会撑爆上下文（见文件头）。
            out_path = args.out or f"./{(result.get('source') or {}).get('name', 'document')}.extract.json"
            with open(out_path, "w", encoding="utf-8") as f:
                json.dump(result, f, ensure_ascii=False, indent=2)

            print("\n✅ 抽取完成\n")
            print(summarize(result, out_path))
            return

        if status == "failed":
            # 端点的错误已经是人话（格式不对/超页数/超体积），原样透出，不要再包一层。
            print(f"❌ extraction failed: {task.get('error', 'unknown error')}", file=sys.stderr)
            sys.exit(1)

        log(f"[doc-parse] status={status} elapsed={elapsed}s", args.json_output)

    print(f"❌ extraction timed out (over {max_wait_sec}s)", file=sys.stderr)
    sys.exit(1)


if __name__ == "__main__":
    main()