#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""IMA 知识库索引抓取（只读元数据：标题/类型/层级，无时间字段）
用法：python ima_kb_index.py
输出：数据/ima_kb_index_诊断原料.json
"""
import json
import os
import time

import requests

KB = 'ima_kb_IktEcZCS'
BASE = 'https://ima.qq.com/openapi/wiki/v1/'
OUT = '/Users/liuxinyuan/Desktop/Hermes输出-工作类/数据/ima_kb_index_诊断原料.json'


def cred():
    cid = open(os.path.expanduser('~/.config/ima/client_id')).read().strip()
    key = open(os.path.expanduser('~/.config/ima/api_key')).read().strip()
    return {'ima-openapi-clientid': cid, 'ima-openapi-apikey': key,
            'Content-Type': 'application/json'}


H = cred()


def call(path, body):
    r = requests.post(BASE + path, headers=H, json=body, timeout=30)
    r.raise_for_status()
    d = r.json()
    if d.get('code') != 0:
        raise RuntimeError('%s: %s' % (path, d.get('msg')))
    return d['data']


def list_dir(folder_id=None):
    out, cursor = [], ''
    while True:
        body = {'knowledge_base_id': KB, 'cursor': cursor, 'limit': 20}
        if folder_id:
            body['folder_id'] = folder_id
        d = call('get_knowledge_list', body)
        out.extend(d.get('knowledge_list', []))
        if d.get('is_end') or not d.get('next_cursor'):
            break
        cursor = d['next_cursor']
        time.sleep(0.2)
    return out


index = []
seen = set()


def walk(folder_id=None, path='', depth=0):
    if depth > 3:
        return
    for it in list_dir(folder_id):
        mid = it.get('media_id', '')
        if it.get('media_type') == 11:  # 个人/业务笔记一律不入索引
            continue
        if mid in seen:
            continue
        seen.add(mid)
        rec = {'title': it.get('title'), 'media_id': mid, 'media_type': it.get('media_type'),
               'path': path, 'tags': it.get('tags', [])}
        index.append(rec)
        if it.get('media_type') == 99:
            walk(mid, path + '/' + str(it.get('title')), depth + 1)


walk()
os.makedirs(os.path.dirname(OUT), exist_ok=True)
json.dump({'kb': KB, 'count': len(index), 'items': index}, open(OUT, 'w'),
          ensure_ascii=False, indent=1)
print('抓取条目：%d ｜输出 %s' % (len(index), OUT))
folders = [x for x in index if x['media_type'] == 99]
print('文件夹 %d 个' % len(folders))
from collections import Counter
print('类型分布:', Counter(x['media_type'] for x in index).most_common())
for f in folders[:15]:
    print('  DIR', f['path'] + '/' + str(f['title']))