#!/usr/bin/env python3
"""Wasabi Storage Utilization special agent for Checkmk.

Queries the Wasabi Stats API for the utilization of every bucket in the
account and emits one Checkmk agent section line (JSON) per bucket. Buckets
are discovered entirely from the API response - no bucket list is configured
anywhere.
"""

import argparse
import json
import sys
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone

SECTION_HEADER = "<<<wasabi_bucket_utilization:sep(0)>>>"
PAGE_SIZE = 100  # Wasabi Stats API caps pageSize at 100; larger values return HTTP 400


def api_get(base_url: str, path: str, params: dict, auth_header: str, timeout: int) -> list:
    """Perform a GET request against the Wasabi Stats API and return the JSON body."""
    url = f"{base_url}{path}?{urllib.parse.urlencode(params)}"
    req = urllib.request.Request(url, headers={"Authorization": auth_header})
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode("utf-8")
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP error {exc.code} from Wasabi Stats API: {body}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"Connection error to Wasabi Stats API: {exc.reason}") from exc

    try:
        data = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Invalid JSON response from Wasabi Stats API: {exc}") from exc

    if isinstance(data, dict):
        # The list-all endpoint wraps the list as {"PageInfo": {...}, "Records": [...]}.
        for key in ("Records", "data", "records", "result", "results"):
            if key in data and isinstance(data[key], list):
                return data[key]
        raise RuntimeError(f"Unexpected response shape from Wasabi Stats API: {raw[:500]}")

    if not isinstance(data, list):
        raise RuntimeError(f"Unexpected response type from Wasabi Stats API: {raw[:500]}")

    return data


def fetch_all_records(base_url: str, auth_header: str, date_from: str, date_to: str, timeout: int) -> list:
    """Fetch all bucket utilization records across all pages."""
    records: list = []
    page_num = 1
    while True:
        page = api_get(
            base_url,
            "/v1/standalone/utilizations/bucket",
            {
                "pageSize": PAGE_SIZE,
                "pageNum": page_num,
                "from": date_from,
                "to": date_to,
                "latest": "true",
            },
            auth_header,
            timeout,
        )
        if not page:
            break
        records.extend(page)
        if len(page) < PAGE_SIZE:
            break
        page_num += 1

    return records


def newest_per_bucket(records: list) -> dict:
    """Reduce a list of utilization records to the newest record per bucket name."""
    newest: dict = {}
    for record in records:
        bucket = record.get("Bucket")
        if not bucket:
            continue
        end_time = record.get("EndTime", "")
        current = newest.get(bucket)
        if current is None or end_time > current.get("EndTime", ""):
            newest[bucket] = record
    return newest


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Wasabi Storage Utilization special agent for Checkmk"
    )
    parser.add_argument("--access-key", required=True, help="Wasabi Access Key")
    parser.add_argument("--secret-key", required=True, help="Wasabi Secret Key")
    parser.add_argument(
        "--base-url", default="https://stats.wasabisys.com", help="Wasabi Stats API base URL"
    )
    parser.add_argument(
        "--lookback-days", type=int, default=30, help="Days to look back for the 'from' date"
    )
    parser.add_argument("--timeout", type=int, default=30, help="Request timeout in seconds")
    args = parser.parse_args()

    today = datetime.now(timezone.utc).date()
    date_from = (today - timedelta(days=args.lookback_days)).isoformat()
    date_to = today.isoformat()
    auth_header = f"{args.access_key}:{args.secret_key}"

    try:
        records = fetch_all_records(
            args.base_url, auth_header, date_from, date_to, args.timeout
        )
    except RuntimeError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1

    buckets = newest_per_bucket(records)

    print(SECTION_HEADER)
    for bucket_name, record in buckets.items():
        output = {
            "bucket": bucket_name,
            "region": record.get("Region"),
            "active_bytes": record.get("RawStorageSizeBytes", 0),
            "deleted_bytes": record.get("DeletedStorageSizeBytes", 0),
            "padded_bytes": record.get("PaddedStorageSizeBytes"),
            "end_time": record.get("EndTime"),
        }
        print(json.dumps(output, ensure_ascii=False))

    return 0


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