#!/usr/bin/env python3
"""
将 Markdown 文件通过飞书 Bot API 写入指定的 docx 文档。
流程：
1. 获取 tenant_access_token
2. 清空文档现有 block 子节点
3. 将 Markdown 解析为飞书 docx block 数组
4. 分批写入
5. 验证 blocks 数量
"""

import json
import re
import sys
import time
import urllib.request
import urllib.error

# ── 配置 ──────────────────────────────────────────────
CONFIG_PATH = "/root/.openclaw/credentials/xiaoxi/config.json"

DOCS = [
    {
        "md_path": "/root/.openclaw/workspace/tmp/data_dictionary.md",
        "doc_id": "YEk6dWmsHoT02lxlO3tcNBrnHXe",
        "name": "数据字典",
        "title": "数据字典",
        "expect_existing": False,
    },
    {
        "md_path": "/root/.openclaw/workspace/output/data_map_lite.md",
        "doc_id": "HmxMd58OeoNXlsxpICLccorhnve",
        "name": "数据地图",
        "title": "数据地图",
        "expect_existing": True,
    }
]

BATCH_SIZE = 35
MAX_BATCH_BYTES = 9000
SLEEP_BETWEEN_BATCHES = 0.5  # 秒


# ── 工具函数 ──────────────────────────────────────────

def load_config():
    with open(CONFIG_PATH) as f:
        cfg = json.load(f)
    app = cfg["apps"][0]
    return app["appId"], app["appSecret"]


def api_request(method, url, token=None, body=None, retries=5):
    """发送 HTTP 请求到飞书 API。"""
    headers = {"Content-Type": "application/json; charset=utf-8"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    
    data = None
    if body is not None:
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    
    last_err = None
    for attempt in range(retries):
        try:
            req = urllib.request.Request(url, data=data, headers=headers)
            if method in ("DELETE",):
                req.method = method
            try:
                with urllib.request.urlopen(req, timeout=60) as resp:
                    result = json.loads(resp.read().decode("utf-8"))
                    code = result.get("code")
                    if code != 0:
                        raise Exception(f"API error {code}: {result.get('msg','')}")
                    return result
            except urllib.error.HTTPError as e:
                err_body = e.read().decode("utf-8", errors="replace")
                if e.code == 429:
                    wait = 2 ** attempt
                    print(f"   频率限制，等待 {wait}s...")
                    time.sleep(wait)
                    continue
                last_err = f"HTTP {e.code}: {err_body}"
                break
        except Exception as e:
            last_err = str(e)
            if "CERTIFICATE" in str(e) or "SSL" in str(e):
                break
            time.sleep(1)
            continue
    raise Exception(last_err)


def get_tenant_token(app_id, app_secret):
    """获取 tenant_access_token。"""
    url = "https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal"
    body = {"app_id": app_id, "app_secret": app_secret}
    result = api_request("POST", url, body=body)
    token = result["tenant_access_token"]
    print(f"[OK] 获取 token 成功: {token[:12]}...")
    return token


def check_document_exists(token, doc_id):
    """检查文档是否存在。"""
    url = f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}"
    try:
        api_request("GET", url, token=token)
        return True
    except Exception:
        return False


# ── 文档清理 ──────────────────────────────────────────

def clear_document(token, doc_id):
    """循环删除文档的所有子 block，直到文档为空。"""
    print(f"[→] 清空文档 {doc_id} 的现有内容...")
    batch_delete_url = f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{doc_id}/children/batch_delete"
    
    iteration = 0
    while True:
        iteration += 1
        try:
            api_request("DELETE", batch_delete_url, token=token,
                        body={"start_index": 0, "end_index": 1})
            print(f"   第 {iteration} 次删除成功")
        except Exception as e:
            err_str = str(e)
            if any(kw in err_str.lower() for kw in [
                "no block", "1740006", "empty", "block is empty",
                "not found", "invalid param", "1770001"
            ]):
                if iteration == 1:
                    print(f"   文档暂无内容，无需清空")
                else:
                    print(f"   文档已清空（{iteration - 1} 次删除后）")
                break
            else:
                raise
    
    print(f"[OK] 文档已清空")


# ── Markdown 解析 ──────────────────────────────────────

def make_text_element(content):
    """创建文本元素，处理行内代码标记。"""
    # 分割 `...` 代码片段
    parts = re.split(r'(`[^`]+`)', content)
    elements = []
    for part in parts:
        if part.startswith('`') and part.endswith('`'):
            inner = part[1:-1]
            elements.append({
                "text_run": {
                    "content": inner,
                    "text_element_style": {
                        "bold": False,
                        "inline_code": True,
                        "italic": False,
                        "strikethrough": False,
                        "underline": False
                    }
                }
            })
        elif part:
            # 处理加粗 **...**
            bold_parts = re.split(r'(\*\*[^*]+\*\*)', part)
            for bp in bold_parts:
                if bp.startswith('**') and bp.endswith('**'):
                    inner = bp[2:-2]
                    elements.append({
                        "text_run": {
                            "content": inner,
                            "text_element_style": {
                                "bold": True,
                                "inline_code": False,
                                "italic": False,
                                "strikethrough": False,
                                "underline": False
                            }
                        }
                    })
                elif bp:
                    elements.append({
                        "text_run": {
                            "content": bp,
                            "text_element_style": {
                                "bold": False,
                                "inline_code": False,
                                "italic": False,
                                "strikethrough": False,
                                "underline": False
                            }
                        }
                    })
    # Remove empty elements
    elements = [e for e in elements if e["text_run"]["content"]]
    if not elements:
        elements = [{"text_run": {"content": "", "text_element_style": {"bold": False, "inline_code": False, "italic": False, "strikethrough": False, "underline": False}}}]
    return elements


def clean_content(text):
    """清理 Markdown 标记获取纯文本。"""
    text = re.sub(r'\*\*(.+?)\*\*', r'\1', text)
    text = re.sub(r'\*(.+?)\*', r'\1', text)
    text = re.sub(r'`([^`]+)`', r'\1', text)
    text = re.sub(r'<br/?>', '', text, flags=re.IGNORECASE)
    text = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text)
    text = re.sub(r'__(.+?)__', r'\1', text)
    text = re.sub(r'_(.+?)_', r'\1', text)
    return text


def make_text_block(content):
    """创建普通文本 block (block_type=2)。"""
    return {
        "block_type": 2,
        "text": {
            "elements": make_text_element(content),
            "style": {"align": 1, "folded": False}
        }
    }


def make_heading_block(level, content):
    """创建标题 block。level: 2=h2, 3=h3, 4=h4"""
    block_type_map = {2: 4, 3: 5, 4: 6, 5: 7, 6: 8}
    block_type = block_type_map.get(level, 4)
    heading_key = f"heading{level}"
    return {
        "block_type": block_type,
        heading_key: {
            "elements": make_text_element(content),
            "style": {"align": 1, "folded": False}
        }
    }


def make_list_block(ordered, content):
    """创建列表 block。"""
    return {
        "block_type": 14 if ordered else 15,
        "text": {
            "elements": make_text_element(content),
            "style": {"align": 1, "folded": False}
        }
    }


def make_quote_block(content):
    """创建引用 block。"""
    return {
        "block_type": 31,
        "text": {
            "elements": make_text_element(content),
            "style": {"align": 1, "folded": False}
        }
    }


def make_code_block(content):
    """创建代码块 block。"""
    return {
        "block_type": 32,
        "text": {
            "elements": [{"text_run": {"content": content, "text_element_style": {
                "bold": False, "inline_code": False,
                "italic": False, "strikethrough": False, "underline": False
            }}}],
            "style": {"align": 1, "folded": False, "language": 29}  # 29 = plain text
        }
    }


def make_divider_block():
    """创建分割线 block。"""
    return {"block_type": 34, "divider": {}}


def parse_markdown_to_blocks(md_text):
    """将 Markdown 文本解析为飞书 docx block 数组。"""
    blocks = []
    lines = md_text.split("\n")
    i = 0
    
    while i < len(lines):
        line = lines[i]
        stripped = line.strip()
        
        # 跳过空行
        if stripped == "":
            i += 1
            continue
        
        # 一级标题跳过（文档 title 已处理）
        if stripped.startswith("# ") and not stripped.startswith("## "):
            i += 1
            continue
        
        # 分割线
        if stripped == "---":
            blocks.append(make_divider_block())
            i += 1
            continue
        
        # 代码块
        if stripped.startswith("```"):
            code_lines = []
            # 提取语言标记（如果有）
            i += 1
            while i < len(lines) and not lines[i].strip().startswith("```"):
                code_lines.append(lines[i])
                i += 1
            i += 1  # 跳过结束 ```
            code_text = "\n".join(code_lines)
            blocks.append(make_code_block(code_text))
            continue
        
        # 标题 ## ~ ######
        heading_match = re.match(r'^(#{2,6})\s+(.+)', stripped)
        if heading_match:
            level = len(heading_match.group(1))
            content = heading_match.group(2).strip()
            blocks.append(make_heading_block(level, content))
            i += 1
            continue
        
        # 无序列表 - xxx
        if re.match(r'^-\s+', stripped):
            content = re.sub(r'^-\s+', '', stripped)
            blocks.append(make_list_block(False, content))
            i += 1
            continue
        
        # 有序列表 N. xxx
        ol_match = re.match(r'^(\d+)\.\s+(.+)', stripped)
        if ol_match:
            content = ol_match.group(2)
            blocks.append(make_list_block(True, content))
            i += 1
            continue
        
        # 引用 > xxx
        if stripped.startswith("> "):
            content = stripped[2:].strip()
            blocks.append(make_quote_block(content))
            i += 1
            continue
        if stripped.startswith(">") and len(stripped) > 1:
            content = stripped[1:].strip()
            blocks.append(make_quote_block(content))
            i += 1
            continue
        
        # 表格行 | a | b |
        if stripped.startswith("|") and stripped.endswith("|"):
            # 跳过纯分隔行
            if re.match(r'^\|[\s\-:|]+\|$', stripped):
                i += 1
                continue
            content = clean_content(stripped.strip("|").strip())
            if content:
                blocks.append(make_text_block(content))
            i += 1
            continue
        
        # 默认：普通文本
        content = stripped.strip()
        if content:
            blocks.append(make_text_block(content))
        i += 1
    
    return blocks


# ── 分批写入 ──────────────────────────────────────────

def write_blocks_batch(token, doc_id, blocks, batch_size=BATCH_SIZE):
    """将 blocks 分批写入文档。"""
    children_url = f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/blocks/{doc_id}/children"
    total = len(blocks)
    written = 0
    batch_num = 0
    
    for start in range(0, total, batch_size):
        batch_num += 1
        batch = blocks[start:start + batch_size]
        body = {"children": batch, "index": -1}
        
        body_json = json.dumps(body, ensure_ascii=False)
        est_bytes = len(body_json.encode("utf-8"))
        
        while est_bytes > MAX_BATCH_BYTES and len(batch) > 5:
            batch = batch[:len(batch)//2]
            body = {"children": batch, "index": -1}
            body_json = json.dumps(body, ensure_ascii=False)
            est_bytes = len(body_json.encode("utf-8"))
        
        try:
            api_request("POST", children_url, token=token, body=body)
            written += len(batch)
            print(f"   批次 {batch_num}: 写入 {len(batch)} blocks ({written}/{total})")
        except Exception as e:
            err_str = str(e)
            # 如果含有 rate limit 错误，等待后重试
            if "99991400" in err_str or "too many" in err_str.lower():
                print(f"   频率限制，等待 5s 后重试...")
                time.sleep(5)
                try:
                    api_request("POST", children_url, token=token, body=body)
                    written += len(batch)
                    print(f"   批次 {batch_num}: 写入 {len(batch)} blocks ({written}/{total})")
                    time.sleep(SLEEP_BETWEEN_BATCHES)
                    continue
                except Exception as e2:
                    print(f"   ✗ 重试失败: {e2}")
            
            print(f"   ✗ 批次 {batch_num} 失败: {e}")
            # 尝试逐个写入
            for block in batch:
                body_single = {"children": [block], "index": -1}
                try:
                    api_request("POST", children_url, token=token, body=body_single)
                    written += 1
                except Exception as e2:
                    print(f"      ✗ 单 block 写入失败: {e2}")
            print(f"   批次 {batch_num} 重试后: {written}/{total}")
        
        time.sleep(SLEEP_BETWEEN_BATCHES)
    
    return written


def verify_document(token, doc_id):
    """验证文档内容并返回 block 数量和统计。"""
    url = f"https://open.feishu.cn/open-apis/docx/v1/documents/{doc_id}/raw_content"
    try:
        result = api_request("GET", url, token=token)
        content = result.get("data", {}).get("content", "")
        block_count = content.count('"block_type"')
        heading2 = content.count('"block_type":4')
        heading3 = content.count('"block_type":5')
        heading4 = content.count('"block_type":6')
        text_blocks = content.count('"block_type":2')
        list_blocks = content.count('"block_type":15') + content.count('"block_type":14')
        quote_blocks = content.count('"block_type":31')
        divider_blocks = content.count('"block_type":34')
        code_blocks = content.count('"block_type":32')
        
        return {
            "total_blocks": block_count,
            "heading2": heading2,
            "heading3": heading3,
            "heading4": heading4,
            "text": text_blocks,
            "list": list_blocks,
            "quote": quote_blocks,
            "divider": divider_blocks,
            "code": code_blocks
        }
    except Exception as e:
        print(f"   ⚠ 验证失败: {e}")
        return None


# ── 主流程 ────────────────────────────────────────────

def process_document(token, doc_config):
    """处理单个文档：检查存在性 → 清理 → 解析 → 写入 → 验证。"""
    name = doc_config["name"]
    doc_id = doc_config["doc_id"]
    md_path = doc_config["md_path"]
    expect_existing = doc_config.get("expect_existing", True)
    
    print(f"\n{'='*60}")
    print(f"📄 处理文档: {name}")
    print(f"   doc_id: {doc_id}")
    print(f"   md_path: {md_path}")
    print(f"{'='*60}")
    
    # 0. 检查文档是否存在
    exists = check_document_exists(token, doc_id)
    if not exists:
        if expect_existing:
            print(f"   ⚠ 文档 {name} 不存在！跳过。")
            return 0, None, "not_found"
        else:
            print(f"   ⚠ 文档 {name} 不存在，正在创建...")
            create_url = "https://open.feishu.cn/open-apis/docx/v1/documents"
            try:
                result = api_request("POST", create_url, token=token,
                    body={"title": doc_config.get("title", name)})
                new_doc = result.get("data", {}).get("document", {})
                new_doc_id = new_doc.get("document_id", "")
                print(f"   创建新文档成功! 新 doc_id: {new_doc_id}")
                doc_config["doc_id"] = new_doc_id
                doc_id = new_doc_id
            except Exception as e:
                print(f"   ✗ 创建文档失败: {e}")
                return 0, None, "create_failed"
    
    # 1. 读取 Markdown
    with open(md_path, "r", encoding="utf-8") as f:
        md_text = f.read()
    print(f"[→] 读取 Markdown: {len(md_text)} 字符")
    
    # 2. 解析为 blocks
    blocks = parse_markdown_to_blocks(md_text)
    print(f"[→] 解析完成: {len(blocks)} blocks")
    
    # 3. 清空文档
    clear_document(token, doc_id)
    
    # 4. 写入
    print(f"[→] 开始写入 {len(blocks)} blocks...")
    written = write_blocks_batch(token, doc_id, blocks)
    print(f"[OK] 写入完成: {written}/{len(blocks)} blocks")
    
    # 5. 验证
    time.sleep(2)
    print(f"[→] 验证文档内容...")
    stats = verify_document(token, doc_id)
    if stats:
        print(f"[验证结果] 总计 {stats['total_blocks']} blocks:")
        print(f"   H2={stats['heading2']}, H3={stats['heading3']}, H4={stats['heading4']}")
        print(f"   text={stats['text']}, list={stats['list']}, quote={stats['quote']}")
        print(f"   divider={stats['divider']}, code={stats['code']}")
    
    return written, stats, doc_id


def main():
    app_id, app_secret = load_config()
    print(f"App ID: {app_id}")
    
    token = get_tenant_token(app_id, app_secret)
    
    results = {}
    for doc_config in DOCS:
        written, stats, final_doc_id = process_document(token, doc_config)
        name = doc_config["name"]
        results[name] = {
            "written": written,
            "stats": stats,
            "doc_id": final_doc_id,
        }
    
    print(f"\n{'='*60}")
    print("📊 总结")
    print(f"{'='*60}")
    for doc_config in DOCS:
        name = doc_config["name"]
        r = results.get(name, {})
        stats = r.get("stats")
        doc_id = r.get("doc_id", doc_config["doc_id"])
        if stats:
            print(f"  ✅ {name}: doc_id={doc_id} 写入 {r['written']} blocks, 验证 {stats['total_blocks']} blocks")
        elif r.get("doc_id") == "not_found":
            print(f"  ❌ {name}: 文档不存在，已跳过")
        elif r.get("doc_id") == "create_failed":
            print(f"  ❌ {name}: 创建文档失败")
        else:
            print(f"  ⚠ {name}: 写入 {r.get('written', 0)} blocks，验证失败")


if __name__ == "__main__":
    main()
