#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import subprocess
import time
from pathlib import Path


def run_nlm(args: list[str], *, json_output: bool = False) -> object:
    cmd = ["uvx", "--from", "notebooklm-mcp-cli", "nlm", *args]
    result = subprocess.run(cmd, check=False, text=True, capture_output=True)
    if result.returncode:
        detail = (result.stderr or result.stdout).strip()
        raise RuntimeError(f"nlm command failed ({result.returncode}): {detail}")
    if json_output:
        return json.loads(result.stdout)
    return result.stdout.strip()


def main() -> int:
    parser = argparse.ArgumentParser(description="Create an owned NotebookLM notebook and download a main-flow infographic.")
    parser.add_argument("--manifest", required=True, type=Path)
    parser.add_argument("--source-dir", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    parser.add_argument("--profile", required=True)
    parser.add_argument("--notebook-id", help="Reuse an existing empty notebook owned by the intended profile")
    parser.add_argument("--timeout", type=int, default=600)
    args = parser.parse_args()
    data = json.loads(args.manifest.read_text(encoding="utf-8"))

    run_nlm(["login", "switch", args.profile])
    run_nlm(["login", "--profile", args.profile, "--check"])
    if args.notebook_id:
        notebook_id = args.notebook_id
        created = run_nlm(["notebook", "get", notebook_id, "--profile", args.profile, "--json"], json_output=True)
        existing = run_nlm(["source", "list", notebook_id, "--profile", args.profile, "--json"], json_output=True)
        if existing:
            raise RuntimeError("existing notebook is not empty; replace stale sources explicitly before reuse")
    else:
        created = run_nlm(["notebook", "create", f"{data['sku']}｜{data['name']}｜销售主流程图", "--profile", args.profile, "--json"], json_output=True)
        notebook_id = created["notebook_id"]

    source_ids: list[str] = []
    source_files = sorted(args.source_dir.glob("0[1-3]-*.md"))
    if len(source_files) != 3:
        raise RuntimeError(f"expected 3 source files, found {len(source_files)}")
    for source in source_files:
        added = run_nlm(["source", "add", notebook_id, "--file", str(source), "--wait", "--profile", args.profile, "--json"], json_output=True)
        source_ids.append(added["source_id"])

    flow = " → ".join(data["main_flow"])
    query = run_nlm(["notebook", "query", notebook_id, f"只依据当前来源回答：主流程是否严格为{flow}？是否包含任何额外步骤？", "--profile", args.profile, "--json"], json_output=True)
    focus = (
        f"生成{data['sku']}{data['name']}主流程信息图。流程必须且只能是：{flow}。"
        f"步骤数量固定为{len(data['main_flow'])}，从左到右排列，不得增加其他步骤。"
        f"禁止出现：{'、'.join(data['forbidden_terms'])}。"
        "视觉必须年轻、时尚、明亮：白色或暖白底，康比特橙黄色为主色，使用橙色到暖黄色的轻渐变、"
        "圆角Bento卡片、轻量图标、充足留白和有动感的流程箭头。避免深蓝大色块、传统企业培训图、"
        "厚重阴影和陈旧报表感。"
        "每个步骤标题必须直接使用给定主流程原文，不得自行添加营销副标题。"
        "不得添加来源未明确给出的速度、准确率、吞吐、节省、提升比例、毫秒级、极速、高精度或效果承诺。"
        f"{data['infographic_focus']}"
    )
    create_args = [
        "create", "infographic", notebook_id,
        "--orientation", "landscape", "--detail", "concise", "--style", "bento_grid",
        "--language", "zh-CN", "--source-ids", ",".join(source_ids), "--focus", focus,
        "--confirm", "--profile", args.profile, "--json",
    ]
    for attempt in range(1, 6):
        try:
            artifact = run_nlm(create_args, json_output=True)
            break
        except RuntimeError as exc:
            message = str(exc)
            if "RESOURCE_EXHAUSTED" not in message and "Rate limited" not in message:
                raise
            if attempt == 5:
                raise
            time.sleep(min(15 * attempt, 60))
    artifact_id = artifact["artifact_id"]

    deadline = time.time() + args.timeout
    final_status = None
    while time.time() < deadline:
        statuses = run_nlm(["studio", "status", notebook_id, "--artifact-id", artifact_id, "--full", "--profile", args.profile, "--json"], json_output=True)
        if statuses:
            final_status = statuses[0]
            if final_status.get("status") == "completed":
                break
            if final_status.get("status") in {"failed", "error"}:
                raise RuntimeError(f"infographic failed: {final_status}")
        time.sleep(20)
    else:
        raise TimeoutError(f"infographic did not complete within {args.timeout}s")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    run_nlm(["download", "infographic", notebook_id, "--id", artifact_id, "--output", str(args.output), "--no-progress"])
    receipt = {
        "sku": data["sku"],
        "notebook_id": notebook_id,
        "notebook_url": created.get("url", f"https://notebook.google.com/notebook/{notebook_id}"),
        "source_ids": source_ids,
        "artifact_id": artifact_id,
        "artifact_status": final_status.get("status") if final_status else None,
        "output": str(args.output),
        "grounded_query": query.get("answer"),
    }
    receipt_path = args.source_dir / "notebooklm-result.json"
    receipt_path.write_text(json.dumps(receipt, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(receipt, ensure_ascii=False))
    return 0


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