"""
导入知识文件到 Dify 知识库
用法: python import_to_dify.py <知识库ID> <文件路径1> [文件路径2 ...]
"""
import os
import sys
import json
import requests

# Fix Windows console encoding for Unicode
import io
if sys.platform == "win32":
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")

DIFY_BASE_URL = "http://dify.bianwoyou.cn"
API_KEY = os.environ.get("DIFY_API_KEY", "")

def import_document(dataset_id, file_path):
    """通过文件上传方式导入文档到知识库"""
    headers = {
        "Authorization": f"Bearer {API_KEY}"
    }

    file_name = os.path.splitext(os.path.basename(file_path))[0]
    data = {
        "name": file_name,
        "indexing_technique": "high_quality",
        "process_rule": {"mode": "automatic"}
    }

    url = f"{DIFY_BASE_URL}/v1/datasets/{dataset_id}/document/create_by_file"

    with open(file_path, "r", encoding="utf-8") as f:
        text_content = f.read()

    # Use create_by_text instead to avoid multipart encoding issues
    text_data = {
        "name": file_name,
        "text": text_content,
        "indexing_technique": "high_quality",
        "process_rule": {"mode": "automatic"}
    }

    text_url = f"{DIFY_BASE_URL}/v1/datasets/{dataset_id}/document/create_by_text"

    print(f"正在导入: {file_name}...")
    resp = requests.post(
        text_url,
        headers={**headers, "Content-Type": "application/json"},
        json=text_data,
        timeout=120
    )

    if resp.status_code in (200, 201):
        result = resp.json()
        doc_id = result.get("document", {}).get("id", "unknown")
        print(f"  [OK] 导入成功! 文档ID: {doc_id}")
        return True
    else:
        print(f"  [FAIL] 导入失败: {resp.status_code} - {resp.text}")
        return False

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("用法: python import_to_dify.py <dataset_id> <file_path1> [file_path2 ...]")
        sys.exit(1)

    if not API_KEY:
        print("错误: 请设置环境变量 DIFY_API_KEY")
        sys.exit(1)

    dataset_id = sys.argv[1]
    file_paths = sys.argv[2:]

    success = 0
    for fp in file_paths:
        if import_document(dataset_id, fp):
            success += 1

    print(f"\n导入完成: {success}/{len(file_paths)} 个文档导入成功")