#!/usr/bin/env python3
"""
飞书文档表格 - Block API 精准单元格写入脚本
通过飞书文档Block树结构定位每个单元格，逐行写入翻译内容，100%行列对齐。

用法：
  python write_translations.py <起始行索引> <批次大小>
  
示例：
  python write_translations.py 0 10   # 处理第1-10行
  python write_translations.py 10 10  # 处理第11-20行

前置条件：
  1. 修改下方 DOC_ID 和 TABLE_BLOCK_ID 为目标文档参数
  2. 修改 translations 列表为实际翻译内容
  3. 修改 TARGET_COL_INDEX 为目标列的索引号（从0开始）
  4. Bot应用需有目标文档的编辑权限
"""
import requests, json, time, sys

# ===== 配置区域（使用前修改） =====
APP_ID = "cli_a931175d41799cc7"
APP_SECRET = "Iw2vEfbjT6GtV0GhbxbZqfQ4nAPtbR14"
DOC_ID = ""           # 文档的 obj_token（通过 wiki get_node 获取）
TABLE_BLOCK_ID = ""   # 目标表格的 block_id（通过 Step 2 获取）
TOTAL_COLS = 6        # 表格总列数
TARGET_COL_INDEX = 3  # 目标列索引（从0开始，英文翻译列通常为3）
SKIP_HEADER_ROWS = 1  # 跳过的表头行数
# ===================================

# 翻译内容列表（按行顺序，None 表示跳过该行）
translations = []  # 使用前填充实际翻译内容


def get_token():
    """获取 tenant_access_token"""
    resp = requests.post(
        "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal",
        json={"app_id": APP_ID, "app_secret": APP_SECRET})
    return resp.json()["tenant_access_token"]


def get_headers(token):
    return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}


def find_tables(headers):
    """获取文档中所有表格的信息"""
    resp = requests.get(
        f"https://open.feishu.cn/open-apis/docx/v1/documents/{DOC_ID}/blocks?page_size=500",
        headers=headers)
    data = resp.json()
    if data.get("code") != 0:
        print(f"获取blocks失败: {data}")
        return []
    blocks = data["data"]["items"]
    tables = []
    for b in blocks:
        if b.get("block_type") == 31:
            prop = b.get("table", {}).get("property", {})
            tables.append({
                "block_id": b["block_id"],
                "rows": prop.get("row_size", 0),
                "cols": prop.get("column_size", 0)
            })
    return tables


def get_target_cell_ids(headers):
    """获取目标列所有单元格的 block_id"""
    resp = requests.get(
        f"https://open.feishu.cn/open-apis/docx/v1/documents/{DOC_ID}/blocks/{TABLE_BLOCK_ID}/children?page_size=500",
        headers=headers)
    data = resp.json()
    if data.get("code") != 0:
        print(f"获取单元格失败: {data}")
        return []
    
    cells = data["data"]["items"]
    target_cells = []
    for i, cell in enumerate(cells):
        row = i // TOTAL_COLS
        col = i % TOTAL_COLS
        if col == TARGET_COL_INDEX and row >= SKIP_HEADER_ROWS:
            target_cells.append({"row": row, "block_id": cell["block_id"]})
    return target_cells


def cell_has_content(cell_id, headers):
    """检查单元格是否已有文本内容"""
    resp = requests.get(
        f"https://open.feishu.cn/open-apis/docx/v1/documents/{DOC_ID}/blocks/{cell_id}/children?page_size=50",
        headers=headers)
    data = resp.json()
    if data.get("code") != 0:
        return False, []
    
    children = data["data"].get("items", [])
    for child in children:
        if child.get("block_type") == 2:
            elements = child.get("text", {}).get("elements", [])
            for el in elements:
                content = el.get("text_run", {}).get("content", "").strip()
                if content:
                    return True, children
    return False, children


def write_to_cell(cell_id, text, headers):
    """写入文本到单元格"""
    # 先清理现有空子blocks
    _, existing = cell_has_content(cell_id, headers)
    for child in existing:
        requests.delete(
            f"https://open.feishu.cn/open-apis/docx/v1/documents/{DOC_ID}/blocks/{child['block_id']}",
            headers=headers)
    
    # 构造文本元素
    text_elements = []
    for line in text.split("\n"):
        text_elements.append({"text_run": {"content": line + "\n"}})
    
    payload = {
        "children": [{
            "block_type": 2,
            "text": {
                "elements": text_elements,
                "style": {}
            }
        }],
        "index": 0
    }
    
    resp = requests.post(
        f"https://open.feishu.cn/open-apis/docx/v1/documents/{DOC_ID}/blocks/{cell_id}/children",
        headers=headers, json=payload)
    return resp.json().get("code") == 0


def main():
    batch_start = int(sys.argv[1]) if len(sys.argv) > 1 else 0
    batch_size = int(sys.argv[2]) if len(sys.argv) > 2 else 10
    
    if not DOC_ID or not TABLE_BLOCK_ID:
        print("错误：请先设置 DOC_ID 和 TABLE_BLOCK_ID")
        print("\n获取方法：")
        print("1. DOC_ID：通过 lark-cli wiki spaces get_node 获取 obj_token")
        print("2. TABLE_BLOCK_ID：运行以下命令查找表格：")
        print("   python write_translations.py --find-tables")
        return
    
    if len(sys.argv) > 1 and sys.argv[1] == "--find-tables":
        token = get_token()
        headers = get_headers(token)
        tables = find_tables(headers)
        print(f"找到 {len(tables)} 个表格：")
        for i, t in enumerate(tables):
            print(f"  表格{i}: {t['rows']}行 x {t['cols']}列, block_id={t['block_id']}")
        return
    
    token = get_token()
    headers = get_headers(token)
    
    # 获取目标列单元格
    target_cells = get_target_cell_ids(headers)
    print(f"目标列共 {len(target_cells)} 个单元格")
    
    end_idx = min(batch_start + batch_size, len(target_cells))
    print(f"处理第{batch_start+1}到第{end_idx}行（共{end_idx - batch_start}行）")
    
    success = 0
    skipped = 0
    errors = 0
    
    for idx in range(batch_start, end_idx):
        cell_id = target_cells[idx]["block_id"]
        translation = translations[idx] if idx < len(translations) else None
        
        if translation is None:
            skipped += 1
            print(f"  行{idx+1}: 跳过（无翻译内容）")
            continue
        
        # 检查是否已有内容
        has_content, _ = cell_has_content(cell_id, headers)
        if has_content:
            skipped += 1
            print(f"  行{idx+1}: 跳过（已有内容）")
            continue
        
        # 写入翻译
        if write_to_cell(cell_id, translation, headers):
            success += 1
            print(f"  行{idx+1}: ✅ 写入成功")
        else:
            errors += 1
            print(f"  行{idx+1}: ❌ 写入失败")
        
        time.sleep(0.3)
    
    print(f"\n=== 批次完成 ===")
    print(f"成功: {success}, 跳过: {skipped}, 失败: {errors}")


if __name__ == "__main__":
    main()
