#!/usr/bin/env python3
"""
从 local.settings.json 加载环境变量，将 ww 目录下所有内容上传到知识库：
- .md：正文写入 content，入索引
- 图片（.png/.jpg/.jpeg/.gif/.webp）：先上传到 Blob，再在索引中建文档，metadata 含 image_url
- .docx/.pptx：在索引中建文档（标题+文件名可搜），可选上传到 Blob 存 file_url
"""
import asyncio
import json
import os
import re
import sys

# 加载 local.settings.json
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
SETTINGS_PATH = os.path.join(SCRIPT_DIR, "local.settings.json")
if not os.path.exists(SETTINGS_PATH):
    print("未找到 local.settings.json", file=sys.stderr)
    sys.exit(1)
with open(SETTINGS_PATH, "r", encoding="utf-8") as f:
    for k, v in json.load(f).get("Values", {}).items():
        os.environ.setdefault(k, v)

sys.path.insert(0, SCRIPT_DIR)

WW_ROOT = os.path.normpath(os.path.join(SCRIPT_DIR, "../../../../ww"))
BLOB_CONTAINER = os.environ.get("INDEXER_BLOB_CONTAINER", "documents")
BLOB_PREFIX = "kb/ww/"
INDEX_NAME = "openclaw-resources"

IMAGE_EXT = {".png", ".jpg", ".jpeg", ".gif", ".webp"}
OFFICE_EXT = {".docx", ".pptx"}


def _safe_id(path: str) -> str:
    """相对路径转成唯一 id，仅保留字母数字、下划线、横线、等号（Azure Search key 要求，不含点）。"""
    rel = path.replace(os.path.sep, "-").replace(".", "_")
    rel = re.sub(r"[^a-zA-Z0-9_\-=]", "_", rel)
    # 去掉连续下划线，避免过长
    rel = re.sub(r"_+", "_", rel).strip("_")
    return "ww-" + (rel[:200] if rel else path.encode("utf-8").hex()[:64])


def _upload_file_to_blob(local_path: str, blob_path: str) -> str:
    """上传本地文件到 Blob，返回可访问的 URL（不含 SAS）。"""
    conn_str = (
        os.environ.get("INDEXER_BLOB_CONNECTION_STRING")
        or os.environ.get("BLOB_STORAGE_CONNECTION_STRING")
        or os.environ.get("AzureWebJobsStorage")
    )
    if not conn_str:
        raise RuntimeError("未配置 Blob 连接字符串")
    from azure.storage.blob import BlobServiceClient

    client = BlobServiceClient.from_connection_string(conn_str)
    container = client.get_container_client(BLOB_CONTAINER)
    blob = container.get_blob_client(blob_path)
    with open(local_path, "rb") as f:
        blob.upload_blob(f, overwrite=True)
    return blob.url


def _collect_ww_docs():
    """遍历 ww 目录，收集要写入索引的文档列表（含图片先上传到 Blob 的 URL）。"""
    docs = []
    if not os.path.isdir(WW_ROOT):
        return docs

    for root, _dirs, files in os.walk(WW_ROOT):
        for name in files:
            local_path = os.path.join(root, name)
            rel_path = os.path.relpath(local_path, WW_ROOT)
            ext = os.path.splitext(name)[1].lower()
            doc_id = _safe_id(rel_path)

            if ext == ".md":
                with open(local_path, "r", encoding="utf-8", errors="ignore") as f:
                    content = f.read()
                # 对接示例下 API 文档保留原 id 便于兼容
                if "对接示例" in rel_path and "API_DOCUMENTATION" in name:
                    doc_id = "ww-api-doc-intelligent-search-agent"
                docs.append({
                    "id": doc_id,
                    "title": name.replace(ext, "") or name,
                    "content": content,
                    "project": "openclaw",
                    "category": "对接示例" if "对接示例" in rel_path else ("产品" if "产品" in rel_path else "PO图"),
                    "tags": "API,文档" if ext == ".md" else "",
                    "source": "human",
                    "author": "ww",
                })
                continue

            if ext in IMAGE_EXT:
                blob_path = BLOB_PREFIX + rel_path.replace(os.path.sep, "/")
                try:
                    url = _upload_file_to_blob(local_path, blob_path)
                except Exception as e:
                    print(f"  [跳过] Blob 上传失败 {rel_path}: {e}", file=sys.stderr)
                    continue
                title = name.replace(ext, "") or name
                docs.append({
                    "id": doc_id,
                    "title": title,
                    "content": f"图片: {title}",
                    "project": "openclaw",
                    "category": "产品" if "产品" in rel_path else "PO图",
                    "tags": "图片,ww",
                    "source": "human",
                    "author": "ww",
                    "metadata": json.dumps({"image_url": url, "file_path": rel_path}, ensure_ascii=False),
                })
                continue

            if ext in OFFICE_EXT:
                docs.append({
                    "id": doc_id,
                    "title": name.replace(ext, "") or name,
                    "content": name,
                    "project": "openclaw",
                    "category": "产品" if "产品" in rel_path else "PO图",
                    "tags": "文档,Office",
                    "source": "human",
                    "author": "ww",
                    "metadata": json.dumps({"file_path": rel_path}, ensure_ascii=False),
                })
                continue

    return docs


async def main():
    from src.server.mcp_server import upload_documents

    docs = _collect_ww_docs()
    if not docs:
        print("ww 目录下未发现可上传的 .md / 图片 / .docx/.pptx 文件")
        return

    print(f"共 {len(docs)} 条文档（含图片等）待上传到索引 {INDEX_NAME}。")
    batch_size = 20
    total_ok, total_fail = 0, 0
    for i in range(0, len(docs), batch_size):
        batch = docs[i : i + batch_size]
        result = await upload_documents(documents=json.dumps(batch), index_name=INDEX_NAME)
        out = json.loads(result)
        total_ok += out.get("uploaded", 0)
        total_fail += out.get("failed", 0)
        if out.get("errors"):
            for e in out["errors"]:
                print(f"  错误: {e}", file=sys.stderr)
    print(json.dumps({"success": total_fail == 0, "uploaded": total_ok, "failed": total_fail}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    asyncio.run(main())