#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""作业16 · 恢复演练后抽检查询（Milvus / Weaviate / ES 探测）

在「删测试集合 → 从备份恢复 → 抽检查询」演练末尾运行。

用法：
  # Milvus（隔离 compose milvus profile，默认 localhost:19530）
  python hw16_02_恢复演练抽检查询.py --backend milvus --collection hw16_test

  # Weaviate（207 现网，需 WEAVIATE_ENDPOINT + API KEY）
  python hw16_02_恢复演练抽检查询.py --backend weaviate --class-name Document

  # Elasticsearch kNN 索引
  python hw16_02_恢复演练抽检查询.py --backend es --index hw16_test

输出：data/homework16/spot_check_<backend>_<timestamp>.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path
from typing import Any


def out_dir() -> Path:
    p = Path(os.environ.get("HW16_OUT", Path(__file__).resolve().parents[2] / "data" / "homework16"))
    p.mkdir(parents=True, exist_ok=True)
    return p


def http_json(url: str, method: str = "GET", data: dict | None = None, headers: dict | None = None) -> Any:
    h = {"Content-Type": "application/json", **(headers or {})}
    body = json.dumps(data).encode("utf-8") if data is not None else None
    req = urllib.request.Request(url, data=body, headers=h, method=method)
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            raw = resp.read().decode("utf-8")
            return json.loads(raw) if raw.strip() else {}
    except urllib.error.HTTPError as e:
        err = e.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {e.code} {url}: {err}") from e


def check_milvus(host: str, port: int, collection: str) -> dict[str, Any]:
    try:
        from pymilvus import Collection, connections, utility
    except ImportError:
        return {
            "ok": False,
            "error": "pymilvus 未安装；pip install pymilvus 或在 milvus 容器内执行",
            "hint": f"curl http://{host}:9091/healthz",
        }
    alias = "hw16"
    connections.connect(alias=alias, host=host, port=port)
    exists = utility.has_collection(collection, using=alias)
    result: dict[str, Any] = {"backend": "milvus", "collection": collection, "exists": exists}
    if exists:
        c = Collection(collection, using=alias)
        c.load()
        result["num_entities"] = c.num_entities
        result["ok"] = c.num_entities > 0
    else:
        result["ok"] = False
        result["error"] = "collection 不存在，恢复可能未完成"
    connections.disconnect(alias)
    return result


def check_weaviate(endpoint: str, class_name: str, api_key: str | None) -> dict[str, Any]:
    headers = {}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    meta = http_json(f"{endpoint.rstrip('/')}/v1/meta", headers=headers)
    # aggregate count
    gql = {"query": f"{{ Aggregate {{ {class_name} {{ meta {{ count }} }} }} }}"}
    try:
        agg = http_json(f"{endpoint.rstrip('/')}/v1/graphql", method="POST", data=gql, headers=headers)
        count = (
            agg.get("data", {})
            .get("Aggregate", {})
            .get(class_name, [{}])[0]
            .get("meta", {})
            .get("count", 0)
        )
    except RuntimeError as e:
        return {"backend": "weaviate", "ok": False, "meta": meta, "error": str(e)}
    return {
        "backend": "weaviate",
        "class_name": class_name,
        "meta_version": meta.get("version"),
        "object_count": count,
        "ok": count > 0,
    }


def check_es(base_url: str, index: str, user: str | None, password: str | None) -> dict[str, Any]:
    auth = None
    if user and password:
        import base64

        token = base64.b64encode(f"{user}:{password}".encode()).decode()
        auth = {"Authorization": f"Basic {token}"}
    health = http_json(f"{base_url.rstrip('/')}/_cluster/health", headers=auth or {})
    count_body = http_json(f"{base_url.rstrip('/')}/{index}/_count", headers=auth or {})
    cnt = count_body.get("count", 0)
    return {
        "backend": "elasticsearch",
        "index": index,
        "cluster_status": health.get("status"),
        "doc_count": cnt,
        "ok": health.get("status") in ("green", "yellow") and cnt > 0,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="向量恢复演练抽检查询")
    parser.add_argument("--backend", required=True, choices=["milvus", "weaviate", "es"])
    parser.add_argument("--collection", default="hw16_test", help="Milvus collection 名")
    parser.add_argument("--class-name", default="Document", help="Weaviate class")
    parser.add_argument("--index", default="hw16_test", help="ES index")
    parser.add_argument("--milvus-host", default=os.environ.get("MILVUS_HOST", "127.0.0.1"))
    parser.add_argument("--milvus-port", type=int, default=int(os.environ.get("MILVUS_PORT", "19530")))
    parser.add_argument("--weaviate-endpoint", default=os.environ.get("WEAVIATE_ENDPOINT", "http://127.0.0.1:8080"))
    parser.add_argument("--weaviate-api-key", default=os.environ.get("WEAVIATE_API_KEY"))
    parser.add_argument("--es-url", default=os.environ.get("ES_URL", "http://127.0.0.1:9200"))
    parser.add_argument("--es-user", default=os.environ.get("ES_USER"))
    parser.add_argument("--es-password", default=os.environ.get("ES_PASSWORD"))
    args = parser.parse_args()

    if args.backend == "milvus":
        result = check_milvus(args.milvus_host, args.milvus_port, args.collection)
    elif args.backend == "weaviate":
        result = check_weaviate(args.weaviate_endpoint, args.class_name, args.weaviate_api_key)
    else:
        result = check_es(args.es_url, args.index, args.es_user, args.es_password)

    result["checked_at"] = datetime.now().isoformat(timespec="seconds")
    od = out_dir()
    out = od / f"spot_check_{args.backend}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    out.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(result, ensure_ascii=False, indent=2))
    print(f"\nWrote {out}")
    sys.exit(0 if result.get("ok") else 1)


if __name__ == "__main__":
    main()