"""修复 wiki 解析 + 重拉失败的文档"""
import json, os, subprocess, sys, time
from datetime import datetime

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from paths import LINKS_DIR, DOCS_DIR, OUTPUT_DIR

LARK = r'C:\Users\admin\AppData\Roaming\npm\node_modules\@larksuite\cli\bin\lark-cli.exe'
IMAGES_DIR = os.path.join(OUTPUT_DIR, 'feishu-images')

def lark_api(path, params=None, timeout=30):
    cmd = [LARK, 'api', 'GET', path, '--format', 'json']
    if params: cmd += ['--params', json.dumps(params)]
    r = subprocess.run(cmd, capture_output=True, timeout=timeout)
    out = r.stdout.decode('utf-8', errors='replace')
    try: return json.loads(out)
    except: return None

def download_media(token, local_path, timeout=30):
    if os.path.isfile(local_path): return True
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    filename = os.path.basename(local_path)
    try:
        r = subprocess.run(
            [LARK, 'api', 'GET', f'/open-apis/drive/v1/medias/{token}/download', '-o', filename],
            capture_output=True, timeout=timeout, cwd=os.path.dirname(local_path)
        )
        return r.returncode == 0 and os.path.isfile(local_path)
    except: return False

def resolve_wiki(node_id):
    data = lark_api('/open-apis/wiki/v2/spaces/get_node', params={'token': node_id})
    if data and data.get('code') == 0:
        node = data.get('data', {}).get('node', {})
        return node.get('obj_token', ''), node.get('title', '')
    return '', ''

def extract_text(text_obj):
    if not text_obj: return ''
    parts = []
    for elem in text_obj.get('elements', []):
        if 'text_run' in elem:
            content = elem['text_run'].get('content', '')
            style = elem['text_run'].get('text_element_style', {})
            if style.get('bold'): content = f'**{content}**'
            if style.get('italic'): content = f'*{content}*'
            parts.append(content)
        elif 'mention_user' in elem:
            parts.append(f'@{elem["mention_user"].get("user_id", "?")}')
    return ''.join(parts).strip()

def blocks_to_markdown(blocks, doc_id, doc_url):
    lines = []
    image_map = {}
    block_map = {b['block_id']: b for b in blocks}
    root_blocks = [b for b in blocks if b.get('parent_id') == doc_id or not b.get('parent_id')]
    
    def process(b, depth=0):
        bt = b.get('block_type', 0)
        children = b.get('children', [])
        
        if bt == 2:
            t = extract_text(b.get('text', {}))
            if t: lines.append(t)
        elif 3 <= bt <= 9:
            t = extract_text(b.get(f'heading{bt-2}', {}) or b.get('text', {}))
            if t: lines.append(f'{"#"*(bt-2)} {t}')
        elif bt == 10:
            t = extract_text(b.get('bullet', {}) or b.get('text', {}))
            if t: lines.append(f'{"  "*depth}- {t}')
        elif bt == 11:
            t = extract_text(b.get('ordered', {}) or b.get('text', {}))
            if t: lines.append(f'{"  "*depth}1. {t}')
        elif bt == 12:
            if 'bullet' in b:
                t = extract_text(b.get('bullet', {}))
                if t: lines.append(f'{"  "*depth}- {t}')
            else:
                t = extract_text(b.get('code', {}) or b.get('text', {}))
                if t: lines.append(f'```\n{t}\n```')
        elif bt == 13:
            t = extract_text(b.get('quote', {}) or b.get('text', {}))
            if t: lines.append(f'> {t}')
        elif bt == 14:
            t = extract_text(b.get('todo', {}) or b.get('text', {}))
            if t: lines.append(f'- [ ] {t}')
        elif bt == 19:
            if 'callout' in b:
                emoji = b.get('callout', {}).get('emoji_id', '')
                if emoji: lines.append(f'\n> :{emoji}:')
            else: lines.append('---')
        elif bt == 27:
            img = b.get('image', {})
            token = img.get('token', '')
            if token:
                local_name = f'{doc_id}_{token[:12]}.png'
                local_path = os.path.join(IMAGES_DIR, local_name)
                image_map[token] = local_path
                lines.append(f'![图片](../feishu-images/{local_name})')
        elif bt == 30:
            sheet_token = b.get('sheet', {}).get('token', '')
            if sheet_token: lines.append(f'\n> [嵌入电子表格: {sheet_token[:20]}...]\n')
        elif bt == 28:
            lines.append('\n| (嵌入表格) |')
        elif bt in (21, 22, 24, 25, 31, 34): pass
        else:
            for key in ['text', 'bullet', 'heading1', 'heading2', 'heading3', 'quote', 'code', 'todo']:
                if key in b:
                    t = extract_text(b[key])
                    if t: lines.append(t)
                    break
        
        for cid in children:
            if cid in block_map: process(block_map[cid], depth + 1)
    
    for b in root_blocks: process(b, depth=0)
    header = [f'# {doc_id}\n', f'Source: {doc_url}\n', '']
    return '\n'.join(header + lines), image_map

# 加载已有索引
content_path = os.path.join(LINKS_DIR, 'all_feishu_content.json')
with open(content_path, 'r', encoding='utf-8') as f:
    content_data = json.load(f)
existing = {d['doc_id']: d for d in content_data.get('documents', [])}

# 找出需要重试的文档（fetched=False 或不存在的）
links_path = os.path.join(LINKS_DIR, 'all_feishu_links.json')
with open(links_path, 'r', encoding='utf-8') as f:
    all_links = json.load(f).get('links', [])

retry = []
for link in all_links:
    did = link.get('doc_id', '')
    if not did: continue
    old = existing.get(did, {})
    if not old.get('fetched', False) or old.get('content_length', 0) <= 200:
        retry.append(link)

print(f'需要重试: {len(retry)} 篇\n')

new = err = 0
for i, link in enumerate(retry):
    doc_id = link.get('doc_id', '')
    doc_type = link.get('type', 'docx')
    url = link.get('url', '')
    sender = link.get('sender', '?')
    
    print(f'[{i+1}/{len(retry)}] {doc_id} ({sender})...', end=' ', flush=True)
    
    real_id = doc_id
    title = ''
    if doc_type == 'wiki':
        real_id, wiki_title = resolve_wiki(doc_id)
        if not real_id:
            print('wiki FAIL'); err += 1; continue
        title = wiki_title
    
    data = lark_api(f'/open-apis/docx/v1/documents/{real_id}/blocks')
    if not data or data.get('code') != 0:
        print('API FAIL'); err += 1; continue
    
    blocks = data.get('data', {}).get('items', [])
    if not blocks:
        print('空'); err += 1; continue
    
    md, image_map = blocks_to_markdown(blocks, doc_id, url)
    
    downloaded = 0
    for token, local_path in image_map.items():
        if download_media(token, local_path): downloaded += 1
        time.sleep(0.1)
    
    with open(os.path.join(DOCS_DIR, f'{doc_id}.md'), 'w', encoding='utf-8') as f: f.write(md)
    
    if not title:
        for line in md.split('\n'):
            if line.startswith('# ') and len(line) > 2: title = line[2:].strip(); break
    
    existing[doc_id] = {
        'url': url, 'doc_id': doc_id, 'title': title,
        'sender': sender, 'group': link.get('group', ''), 'time': link.get('time', ''),
        'content_length': len(md), 'image_blocks': len(image_map),
        'images_downloaded': downloaded, 'fetched': True,
        'fetched_at': datetime.now().isoformat(),
    }
    new += 1
    print(f'OK ({len(md)}c, {len(image_map)}img, {downloaded}dl)')
    time.sleep(0.15)

# 保存索引
content_data['documents'] = list(existing.values())
content_data['total'] = len(existing)
with open(content_path, 'w', encoding='utf-8') as f:
    json.dump(content_data, f, ensure_ascii=False, indent=2)

imgs = len([f for f in os.listdir(IMAGES_DIR) if f.endswith('.png')]) if os.path.isdir(IMAGES_DIR) else 0
print(f'\n完成! 新拉取:{new} 失败:{err} 总计:{len(existing)} 图片:{imgs}')