---

#!/usr/bin/env python3
"""
Подсчёт размера индексов OpenSearch по паттерну.

Примеры:
  ./os_index_size.py 'logs-*'
  ./os_index_size.py 'logs-2026.09.*' --host https://os.local:9200 -u admin -p secret --insecure
  ./os_index_size.py 'app-*,audit-*' --sort size --top 20
  OS_USER=admin OS_PASS=secret ./os_index_size.py 'metrics-*' --json
"""
import argparse
import base64
import json
import os
import ssl
import sys
import urllib.error
import urllib.parse
import urllib.request


def human(n: int) -> str:
    size = float(n)
    for unit in ("B", "KB", "MB", "GB", "TB", "PB"):
        if size < 1024 or unit == "PB":
            return f"{size:.2f} {unit}"
        size /= 1024


def to_int(v) -> int:
    try:
        return int(v)
    except (TypeError, ValueError):
        return 0  # закрытые индексы возвращают null


def fetch_indices(args):
    pattern = urllib.parse.quote(args.pattern, safe="*,-")
    params = urllib.parse.urlencode({
        "format": "json",
        "bytes": "b",
        "h": "index,health,status,pri,rep,docs.count,store.size,pri.store.size",
        "expand_wildcards": "all" if args.all else "open",
    })
    url = f"{args.host.rstrip('/')}/_cat/indices/{pattern}?{params}"

    req = urllib.request.Request(url)
    if args.user:
        token = base64.b64encode(f"{args.user}:{args.password or ''}".encode()).decode()
        req.add_header("Authorization", f"Basic {token}")

    ctx = ssl.create_default_context(cafile=args.cacert) if args.cacert else ssl.create_default_context()
    if args.insecure:
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

    try:
        with urllib.request.urlopen(req, context=ctx, timeout=args.timeout) as resp:
            return json.load(resp)
    except urllib.error.HTTPError as e:
        if e.code == 404:
            return []  # по паттерну ничего не найдено
        sys.exit(f"HTTP {e.code}: {e.read().decode(errors='replace')}")
    except urllib.error.URLError as e:
        sys.exit(f"Ошибка подключения: {e.reason}")


def main():
    p = argparse.ArgumentParser(description="Размер индексов OpenSearch по паттерну")
    p.add_argument("pattern", help="Паттерн индексов, напр. 'logs-*' или 'a-*,b-*,-b-old*'")
    p.add_argument("--host", default=os.getenv("OS_HOST", "https://localhost:9200"))
    p.add_argument("-u", "--user", default=os.getenv("OS_USER"))
    p.add_argument("-p", "--password", default=os.getenv("OS_PASS"))
    p.add_argument("--cacert", help="Путь к CA-сертификату")
    p.add_argument("-k", "--insecure", action="store_true", help="Не проверять TLS-сертификат")
    p.add_argument("--all", action="store_true", help="Включать закрытые/скрытые индексы")
    p.add_argument("--sort", choices=["name", "size", "docs"], default="size")
    p.add_argument("--top", type=int, default=0, help="Показать только N первых")
    p.add_argument("--json", action="store_true", help="Вывод в JSON")
    p.add_argument("--timeout", type=int, default=30)
    args = p.parse_args()

    rows = fetch_indices(args)
    if not rows:
        print(f"Индексы по паттерну '{args.pattern}' не найдены")
        return

    for r in rows:
        r["_size"] = to_int(r.get("store.size"))
        r["_pri"] = to_int(r.get("pri.store.size"))
        r["_docs"] = to_int(r.get("docs.count"))

    key = {"name": lambda r: r["index"], "size": lambda r: r["_size"], "docs": lambda r: r["_docs"]}[args.sort]
    rows.sort(key=key, reverse=args.sort != "name")

    total_size = sum(r["_size"] for r in rows)
    total_pri = sum(r["_pri"] for r in rows)
    total_docs = sum(r["_docs"] for r in rows)
    shown = rows[: args.top] if args.top else rows

    if args.json:
        print(json.dumps({
            "pattern": args.pattern,
            "count": len(rows),
            "total_store_bytes": total_size,
            "total_primary_bytes": total_pri,
            "total_docs": total_docs,
            "indices": [
                {"index": r["index"], "health": r.get("health"), "status": r.get("status"),
                 "store_bytes": r["_size"], "primary_bytes": r["_pri"], "docs": r["_docs"]}
                for r in shown
            ],
        }, indent=2, ensure_ascii=False))
        return

    w = max(len(r["index"]) for r in shown)
    w = max(w, len("INDEX"))
    print(f"{'INDEX':<{w}}  {'HEALTH':<7} {'STATUS':<6} {'DOCS':>14} {'PRIMARY':>12} {'TOTAL':>12}")
    for r in shown:
        print(f"{r['index']:<{w}}  {str(r.get('health')):<7} {str(r.get('status')):<6} "
              f"{r['_docs']:>14,} {human(r['_pri']):>12} {human(r['_size']):>12}")

    print("-" * (w + 58))
    if args.top and args.top < len(rows):
        print(f"(показано {args.top} из {len(rows)})")
    print(f"Индексов: {len(rows)}")
    print(f"Документов: {total_docs:,}")
    print(f"Размер primary-шардов: {human(total_pri)} ({total_pri} B)")
    print(f"Полный размер (с репликами): {human(total_size)} ({total_size} B)")


if __name__ == "__main__":
    main()