---

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

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

DEBUG = False


def log(msg):
    if DEBUG:
        print(f"[debug] {msg}", file=sys.stderr)


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


class Client:
    def __init__(self, args):
        self.host = args.host.rstrip("/")
        self.timeout = args.timeout
        self.auth = None
        if args.user:
            self.auth = base64.b64encode(f"{args.user}:{args.password or ''}".encode()).decode()
        self.ctx = ssl.create_default_context(cafile=args.cacert) if args.cacert else ssl.create_default_context()
        if args.insecure:
            self.ctx.check_hostname = False
            self.ctx.verify_mode = ssl.CERT_NONE

    def get(self, path, params=None):
        url = f"{self.host}{path}"
        if params:
            url += "?" + urllib.parse.urlencode(params)
        log(f"GET {url}")
        req = urllib.request.Request(url)
        if self.auth:
            req.add_header("Authorization", f"Basic {self.auth}")
        try:
            with urllib.request.urlopen(req, context=self.ctx, timeout=self.timeout) as resp:
                body = resp.read().decode(errors="replace")
                log(f"HTTP {resp.status}, {len(body)} байт: {body[:500]}")
                return resp.status, body
        except urllib.error.HTTPError as e:
            body = e.read().decode(errors="replace")
            log(f"HTTP {e.code}: {body[:500]}")
            return e.code, body
        except urllib.error.URLError as e:
            sys.exit(f"Ошибка подключения к {self.host}: {e.reason}")


def parse_json(status, body, what):
    try:
        return json.loads(body)
    except json.JSONDecodeError:
        sys.exit(f"{what}: ответ не JSON (HTTP {status}). Проверьте --host — "
                 f"возможно, это Dashboards/прокси, а не REST API.\nНачало ответа: {body[:300]}")


def fetch_indices(client, pattern, expand):
    status, body = client.get(
        f"/_cat/indices/{urllib.parse.quote(pattern, safe='*,-')}",
        {"format": "json", "bytes": "b", "expand_wildcards": expand,
         "h": "index,health,status,pri,rep,docs.count,store.size,pri.store.size"},
    )
    if status == 401:
        sys.exit("HTTP 401: неверный логин/пароль или не переданы -u/-p")
    if status == 403:
        sys.exit(f"HTTP 403: у пользователя нет прав на _cat/indices.\n{body[:500]}")
    if status == 404:
        data = parse_json(status, body, "_cat/indices")
        err_type = (data.get("error") or {}).get("type") if isinstance(data, dict) else None
        if err_type == "index_not_found_exception":
            return []
        sys.exit(f"HTTP 404, но это не 'индекс не найден' — вероятно, неверный URL.\n{body[:500]}")
    if status != 200:
        sys.exit(f"HTTP {status}: {body[:500]}")

    data = parse_json(status, body, "_cat/indices")
    if not isinstance(data, list):
        sys.exit(f"Неожиданный формат ответа: {body[:500]}")
    return data


def diagnose(client, args):
    """Вызывается, если по паттерну пусто: пытаемся понять почему."""
    print(f"Индексы по паттерну '{args.pattern}' не найдены. Диагностика:", file=sys.stderr)

    status, body = client.get("/")
    if status == 200:
        info = parse_json(status, body, "/")
        ver = info.get("version", {})
        print(f"  • Кластер: {info.get('cluster_name')}, "
              f"{ver.get('distribution', 'elasticsearch')} {ver.get('number')}", file=sys.stderr)
    else:
        print(f"  • GET / вернул HTTP {status}", file=sys.stderr)

    if not args.all:
        hidden = fetch_indices(client, args.pattern, "all")
        if hidden:
            names = ", ".join(r["index"] for r in hidden[:5])
            print(f"  • С учётом закрытых/скрытых индексов найдено {len(hidden)} "
                  f"(напр. {names}). Запустите с --all", file=sys.stderr)
            return

    status, body = client.get("/_plugins/_security/authinfo")
    if status == 200:
        info = parse_json(status, body, "authinfo")
        print(f"  • Пользователь: {info.get('user_name')}, роли: {info.get('roles')}", file=sys.stderr)
        print("  • Если индексы точно есть — у этих ролей, скорее всего, нет прав "
              "indices:monitor/stats на нужные индексы (security plugin скрывает недоступные)",
              file=sys.stderr)
    print("  • Для подробностей запустите с --debug", file=sys.stderr)


def main():
    global DEBUG
    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)
    p.add_argument("--debug", action="store_true", help="Печатать запросы и ответы")
    args = p.parse_args()
    DEBUG = args.debug

    client = Client(args)
    rows = fetch_indices(client, args.pattern, "all" if args.all else "open")
    if not rows:
        diagnose(client, args)
        sys.exit(1)

    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("INDEX"), *(len(r["index"]) for r in shown))
    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()