#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
飞书云文档抓取脚本 —— 抓取决策链画板(docx/whiteboard)的文本内容。
用途：把飞书「决策链」文档链接转成结构化文本，供工作台 AI 助手引用。

用法：
    python3 fetch_feishu_doc.py <document_id>
    python3 fetch_feishu_doc.py FmYDdYkq1oDubgx5kQLcLKBnnXg

前置条件（飞书开放平台需人工完成一次）：
    1. 工作台 App (cli_aaee6a2f90b89bc1) 在「权限管理」中开通云文档读取权限：
       - docx:document:readonly  （读取新版文档）
       - wiki:wiki:readonly     （若是知识库 wiki 文档）
    2. 目标决策链文档需对 App 可见：把文档「分享/协作者」加入该 App，或
       确保文档在 App 的可用范围内（企业自建应用默认可见同企业文档时需确认）。
    3. 若文档是「知识库 wiki」，document_id 要用 wiki 节点 token，而非普通 docx id。

注意：
    - 画板(whiteboard)块结构复杂，本脚本递归 dump 所有子块并提取 text 类型内容。
    - 若返回 code=99991663（无权限）或 code=106（token 无效），先检查上面权限。
"""
import json
import sys
import urllib.request

APP_ID = 'cli_aaee6a2f90b89bc1'
APP_SECRET = 'arht5fd8wyppZysFvOjUcb2bV0UMVVYF'
BASE = 'https://open.feishu.cn/open-apis'


def http_get(url, token):
    req = urllib.request.Request(
        url, headers={'Authorization': 'Bearer ' + token}, method='GET')
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode('utf-8'))


def http_post(url, body, token=None):
    data = json.dumps(body).encode()
    h = {'Content-Type': 'application/json'}
    if token:
        h['Authorization'] = 'Bearer ' + token
    req = urllib.request.Request(url, data=data, headers=h, method='POST')
    with urllib.request.urlopen(req, timeout=30) as r:
        return json.loads(r.read().decode('utf-8'))


def get_tenant_token():
    r = http_post(f'{BASE}/auth/v3/tenant_access_token/internal',
                  {'app_id': APP_ID, 'app_secret': APP_SECRET})
    if r.get('code') != 0:
        raise RuntimeError('获取 tenant_access_token 失败: ' + json.dumps(r, ensure_ascii=False))
    return r['tenant_access_token']


def get_blocks(token, doc_id):
    out = []
    page_token = ''
    while True:
        url = f'{BASE}/docx/v1/documents/{doc_id}/blocks?page_size=500'
        if page_token:
            url += '&page_token=' + page_token
        r = http_get(url, token)
        if r.get('code') != 0:
            raise RuntimeError('拉取 blocks 失败: ' + json.dumps(r, ensure_ascii=False))
        data = r.get('data', {})
        items = data.get('items') or []
        out.extend(items)
        if not data.get('has_more'):
            break
        page_token = data.get('page_token', '')
    return out


def extract_texts(blocks, depth=0):
    """递归提取块中的文字，保留层级。"""
    lines = []
    for b in blocks:
        bt = b.get('block_type')
        if bt == 2:  # text 块
            els = (b.get('text') or {}).get('elements') or []
            txt = ''
            for e in els:
                tr = e.get('text_run') or {}
                txt += tr.get('content', '')
            if txt.strip():
                lines.append('  ' * depth + txt.strip())
        elif bt == 3:  # heading 块
            for n in (1, 2, 3, 4, 5, 6, 7, 8, 9):
                key = f'heading{n}'
                if key in b:
                    els = (b[key].get('elements') or [])
                    txt = ''.join((e.get('text_run') or {}).get('content', '') for e in els)
                    if txt.strip():
                        lines.append('  ' * depth + '[H] ' + txt.strip())
                    break
        # 递归子块（whiteboard / 列表 / 表格 / 图片等都可能带 children）
        children = b.get('children') or []
        if children:
            lines.extend(extract_texts(children, depth + 1))
    return lines


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    doc_id = sys.argv[1]
    token = get_tenant_token()
    blocks = get_blocks(token, doc_id)
    print(f'== 文档 {doc_id} 共 {len(blocks)} 个块 ==\n')
    lines = extract_texts(blocks)
    if not lines:
        print('(未提取到文本块，打印原始结构前3块)')
        print(json.dumps(blocks[:3], ensure_ascii=False, indent=2))
    else:
        print('\n'.join(lines))


if __name__ == '__main__':
    main()