#!/usr/bin/env python3
"""目标 WeKnora ES 索引里某个（些）知识库的文档导出（只读，scroll + gzip JSONL）。

用途：发布/导入之前把 ES 那一层也留档到本地，与 PG 行级快照配成"导入前完整现场"。
只发 _count / _search / _search/scroll，不写 ES 任何数据。

用法：
  python3 es_kb_export.py --endpoint http://172.18.0.3:9200 --index weknora_test_chunks \\
      --kb <kb_id> [--kb <kb_id> ...] --output /path/es_documents.jsonl.gz [--limit N] [--batch 2000]

输出：
  <output>                 gzip JSONL，一行一条 {"_index":..,"_id":..,"_source":{...}}
  <output>.manifest.json   每个库的文档数、字节数、gz 的 sha256/size、索引与 endpoint
"""
import argparse
import gzip
import hashlib
import json
import sys
import urllib.error
import urllib.request


def call(endpoint, path, payload=None, method=None, timeout=300):
    url = endpoint.rstrip("/") + path
    data = json.dumps(payload).encode() if payload is not None else None
    request = urllib.request.Request(url, data=data, method=method or ("POST" if data else "GET"),
                                     headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.load(response)
    except urllib.error.HTTPError as error:
        body = error.read().decode("utf-8", "replace")[:300]
        raise SystemExit(f"ES 请求失败 {path}: HTTP {error.code} {body}")


def count_for(endpoint, index, knowledge_base_id):
    result = call(endpoint, f"/{index}/_count",
                  {"query": {"term": {"knowledge_base_id.keyword": knowledge_base_id}}})
    return int(result.get("count", 0))


def export(endpoint, index, knowledge_base_ids, output, batch, limit):
    counts = {}
    written = 0
    uncompressed_bytes = 0
    for kb in knowledge_base_ids:
        counts[kb] = count_for(endpoint, index, kb)

    # scroll：每次 batch 条，_doc 排序，只取 _source
    with gzip.open(output, "wt", encoding="utf-8", compresslevel=6) as handle:
        for kb in knowledge_base_ids:
            if counts[kb] == 0:
                continue
            scroll_id = None
            try:
                response = call(endpoint, f"/{index}/_search?scroll=5m", {
                    "size": min(batch, limit) if limit else batch,
                    "sort": ["_doc"],
                    "query": {"term": {"knowledge_base_id.keyword": kb}},
                    "_source": True,
                })
                while True:
                    scroll_id = response.get("_scroll_id") or scroll_id
                    hits = response.get("hits", {}).get("hits", [])
                    if not hits:
                        break
                    for hit in hits:
                        line = json.dumps({"_index": hit.get("_index"), "_id": hit.get("_id"),
                                           "_source": hit.get("_source")}, ensure_ascii=False)
                        handle.write(line + "\n")
                        written += 1
                        uncompressed_bytes += len(line.encode("utf-8"))
                        if limit and written >= limit:
                            break
                    if limit and written >= limit:
                        break
                    response = call(endpoint, "/_search/scroll",
                                    {"scroll": "5m", "scroll_id": scroll_id})
            finally:
                if scroll_id:
                    try:
                        call(endpoint, "/_search/scroll", {"scroll_id": [scroll_id]}, method="DELETE")
                    except SystemExit:
                        pass
            if limit and written >= limit:
                break

    digest = hashlib.sha256()
    with open(output, "rb") as handle:
        for block in iter(lambda: handle.read(1 << 20), b""):
            digest.update(block)
    manifest = {
        "kind": "weknora-es-kb-export",
        "endpoint": endpoint,
        "index": index,
        "knowledge_base_document_counts": counts,
        "documents_written": written,
        "uncompressed_bytes": uncompressed_bytes,
        "file": output.split("/")[-1],
        "bytes": __import__("os").path.getsize(output),
        "sha256": digest.hexdigest(),
    }
    with open(output + ".manifest.json", "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, ensure_ascii=False, indent=2)
        handle.write("\n")
    return manifest


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--endpoint", required=True)
    parser.add_argument("--index", default="weknora_test_chunks")
    parser.add_argument("--kb", action="append", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--batch", type=int, default=2000)
    parser.add_argument("--limit", type=int, default=0, help="仅用于冒烟测试：最多导出多少条")
    arguments = parser.parse_args()
    manifest = export(arguments.endpoint, arguments.index, arguments.kb, arguments.output,
                      arguments.batch, arguments.limit)
    print(json.dumps({"documents": manifest["documents_written"],
                      "counts": manifest["knowledge_base_document_counts"],
                      "bytes": manifest["bytes"], "sha256": manifest["sha256"][:16] + "…"},
                     ensure_ascii=False))


if __name__ == "__main__":
    sys.exit(main())