atc-agents/api/storage_s3.py

"""ObjectScale / S3 storage API for Command Center."""

from __future__ import annotations

import os
import time
from collections import deque
from datetime import datetime, timezone
from typing import Any

import boto3
from botocore.client import Config
from botocore.exceptions import ClientError
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse, StreamingResponse

S3_ENDPOINT = os.getenv("S3_ENDPOINT", "http://10.0.20.111:9020")
S3_ACCESS_KEY = os.getenv("S3_ACCESS_KEY", "object_admin1")
S3_SECRET_KEY = os.getenv("S3_SECRET_KEY", "ChangeMeChangeMeChangeMeChangeMeChangeMe")
S3_REGION = os.getenv("S3_REGION", "us-east-1")

router = APIRouter(prefix="/api/storage/s3", tags=["storage"])

# Bounded scan so the dashboard stays responsive on big / fast-growing buckets.
SCAN_MAX_OBJECTS = int(os.getenv("S3_SCAN_MAX_OBJECTS", "80000"))
SCAN_DEADLINE_S = float(os.getenv("S3_SCAN_DEADLINE_S", "20"))
# Per top-level prefix cap so a single huge prefix (e.g. kafka/ CDC json) cannot
# starve the others — this keeps the type/size composition representative.
SCAN_PER_PREFIX = int(os.getenv("S3_SCAN_PER_PREFIX", "20000"))

# In-process activity log — every S3 op routed through this API is recorded so
# the dashboard can show live "what is happening on the storage" timelines.
_activity: deque[tuple[float, str]] = deque(maxlen=20000)
_analytics_cache: dict[str, Any] = {"ts": 0.0, "data": None}
_ANALYTICS_TTL = 45.0


def _client():
    return boto3.client(
        "s3",
        endpoint_url=S3_ENDPOINT,
        aws_access_key_id=S3_ACCESS_KEY,
        aws_secret_access_key=S3_SECRET_KEY,
        region_name=S3_REGION,
        config=Config(signature_version="s3v4", s3={"addressing_style": "path"}),
    )


def _human_size(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if n < 1024:
            return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} PB"


def _track(op: str) -> None:
    try:
        _activity.append((time.time(), op))
    except Exception:
        pass


def _ext(key: str) -> str:
    base = key.rsplit("/", 1)[-1]
    if "." in base:
        e = base.rsplit(".", 1)[-1].lower()
        if 1 <= len(e) <= 8 and e.isalnum():
            return e
    return "(none)"


def _size_class(n: int) -> str:
    kb, mb, gb = 1024, 1024 ** 2, 1024 ** 3
    if n < kb:
        return "<1 KB"
    if n < mb:
        return "1 KB–1 MB"
    if n < 10 * mb:
        return "1–10 MB"
    if n < 100 * mb:
        return "10–100 MB"
    if n < gb:
        return "100 MB–1 GB"
    return ">1 GB"


@router.get("/health")
async def s3_health():
    try:
        _track("health")
        s3 = _client()
        buckets = s3.list_buckets()
        names = [b["Name"] for b in buckets.get("Buckets", [])]
        return {
            "ok": True,
            "endpoint": S3_ENDPOINT,
            "buckets": len(names),
            "bucket_names": names,
        }
    except Exception as exc:
        return JSONResponse({"ok": False, "endpoint": S3_ENDPOINT, "error": str(exc)}, status_code=502)


@router.get("/buckets")
async def list_buckets():
    try:
        _track("list_buckets")
        s3 = _client()
        resp = s3.list_buckets()
        items = []
        for b in resp.get("Buckets", []):
            name = b["Name"]
            try:
                loc = s3.list_objects_v2(Bucket=name, MaxKeys=1)
                count_hint = loc.get("KeyCount", 0)
            except ClientError:
                count_hint = None
            items.append({
                "name": name,
                "created": b.get("CreationDate", "").isoformat() if b.get("CreationDate") else None,
                "has_objects": bool(count_hint),
            })
        return {"ok": True, "buckets": items, "endpoint": S3_ENDPOINT}
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)


@router.get("/buckets/{bucket}/objects")
async def list_objects(
    bucket: str,
    prefix: str = Query("", alias="prefix"),
    max_keys: int = Query(200, le=1000),
    recursive: bool = Query(False),
    q: str = Query(""),
    ext: str = Query(""),
    token: str = Query(""),
):
    """Browse a bucket.

    - Folder mode (default): one level, returns sub-folders + objects.
    - recursive / q / ext: flat search across the whole prefix subtree,
      paginated via `token` (returned as `next_token`).
    """
    try:
        _track("list_objects")
        s3 = _client()
        ql = q.lower().strip()
        ext_l = ext.lower().lstrip(".").strip()
        flat = bool(recursive or ql or ext_l)
        folders: list[dict[str, Any]] = []
        objects: list[dict[str, Any]] = []
        cont = token or None
        pages = 0
        scanned = 0
        next_token: str | None = None
        PAGE_BUDGET = 40

        while True:
            kw: dict[str, Any] = {"Bucket": bucket, "Prefix": prefix, "MaxKeys": 1000}
            if not flat:
                kw["Delimiter"] = "/"
            if cont:
                kw["ContinuationToken"] = cont
            r = s3.list_objects_v2(**kw)
            pages += 1

            for p in r.get("CommonPrefixes", []):
                folders.append({
                    "type": "prefix",
                    "name": p["Prefix"][len(prefix):].rstrip("/"),
                    "prefix": p["Prefix"],
                })

            for o in r.get("Contents", []):
                key = o["Key"]
                if key == prefix:
                    continue
                scanned += 1
                if ql and ql not in key.lower():
                    continue
                if ext_l and _ext(key) != ext_l:
                    continue
                objects.append({
                    "type": "object",
                    "key": key,
                    "name": key[len(prefix):] if key.startswith(prefix) else key,
                    "ext": _ext(key),
                    "size": o.get("Size", 0),
                    "size_human": _human_size(o.get("Size", 0)),
                    "modified": o.get("LastModified", "").isoformat() if o.get("LastModified") else None,
                })

            cont = r.get("NextContinuationToken")
            more = r.get("IsTruncated")

            if not flat:
                next_token = cont if more else None
                break
            if len(objects) >= max_keys or not more or pages >= PAGE_BUDGET:
                next_token = cont if more else None
                break

        return {
            "ok": True,
            "bucket": bucket,
            "prefix": prefix,
            "flat": flat,
            "folders": folders,
            "objects": objects[:max_keys] if flat else objects,
            "scanned": scanned,
            "next_token": next_token,
            "truncated": bool(next_token),
        }
    except ClientError as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=403)
    except Exception as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=502)


@router.get("/buckets/{bucket}/download")
async def download_object(bucket: str, key: str = Query(...)):
    try:
        _track("download")
        s3 = _client()
        obj = s3.get_object(Bucket=bucket, Key=key)
        body = obj["Body"]
        filename = key.split("/")[-1] or "download"
        media = obj.get("ContentType") or "application/octet-stream"

        def stream():
            while chunk := body.read(1024 * 256):
                yield chunk

        return StreamingResponse(
            stream(),
            media_type=media,
            headers={"Content-Disposition": f'attachment; filename="{filename}"'},
        )
    except ClientError as exc:
        return JSONResponse({"ok": False, "error": str(exc)}, status_code=404)


# ── Analytics ────────────────────────────────────────────────────────────────


def _accumulate(agg: dict[str, Any], bucket: str, o: dict[str, Any]) -> None:
    sz = int(o.get("Size", 0) or 0)
    key = o["Key"]
    lm = o.get("LastModified")
    agg["objects"] += 1
    agg["bytes"] += sz
    bb = agg["by_bucket"].setdefault(bucket, [0, 0])
    bb[0] += 1
    bb[1] += sz
    e = agg["by_ext"].setdefault(_ext(key), [0, 0])
    e[0] += 1
    e[1] += sz
    sc = _size_class(sz)
    agg["by_size"][sc] = agg["by_size"].get(sc, 0) + 1
    if lm:
        d = lm.astimezone(timezone.utc).strftime("%Y-%m-%d")
        dd = agg["by_day"].setdefault(d, [0, 0])
        dd[0] += 1
        dd[1] += sz
        agg["recent"].append((lm.isoformat(), bucket, key, sz))
        if len(agg["recent"]) > 400:
            agg["recent"] = sorted(agg["recent"], reverse=True)[:60]
    rest = key[len(bucket) + 1:] if key.startswith(bucket + "/") else key
    top = rest.split("/", 1)[0] if "/" in rest else "(root)"
    pp = agg["by_prefix"].setdefault(f"{bucket}/{top}", [0, 0])
    pp[0] += 1
    pp[1] += sz
    agg["largest"].append((sz, bucket, key, lm.isoformat() if lm else None))
    if len(agg["largest"]) > 400:
        agg["largest"] = sorted(agg["largest"], reverse=True)[:60]


def _scan_prefix(s3, bucket: str, prefix: str, cap: int, deadline: float, agg: dict[str, Any]) -> bool:
    """Scan one prefix subtree (flat). Returns True if capped/cut short."""
    tok = None
    n = 0
    while True:
        if time.time() > deadline or n >= cap or agg["objects"] >= SCAN_MAX_OBJECTS:
            return True
        kw: dict[str, Any] = {"Bucket": bucket, "MaxKeys": 1000}
        if prefix:
            kw["Prefix"] = prefix
        if tok:
            kw["ContinuationToken"] = tok
        try:
            r = s3.list_objects_v2(**kw)
        except Exception:
            return False
        for o in r.get("Contents", []):
            _accumulate(agg, bucket, o)
            n += 1
            if n >= cap or agg["objects"] >= SCAN_MAX_OBJECTS:
                return True
        if r.get("IsTruncated"):
            tok = r.get("NextContinuationToken")
            if not tok:
                return False
        else:
            return False


def _scan_analytics() -> dict[str, Any]:
    s3 = _client()
    deadline = time.time() + SCAN_DEADLINE_S
    resp = s3.list_buckets()
    bucket_objs = resp.get("Buckets", [])
    created_map = {
        b["Name"]: (b.get("CreationDate").isoformat() if b.get("CreationDate") else None)
        for b in bucket_objs
    }
    names = [b["Name"] for b in bucket_objs]

    agg: dict[str, Any] = {
        "objects": 0, "bytes": 0, "by_bucket": {}, "by_ext": {}, "by_size": {},
        "by_day": {}, "by_prefix": {}, "largest": [], "recent": [],
    }
    truncated = False

    for name in names:
        if time.time() > deadline:
            truncated = True
            break
        try:
            r = s3.list_objects_v2(Bucket=name, Delimiter="/", MaxKeys=1000)
        except Exception:
            continue
        # root-level objects
        for o in r.get("Contents", []):
            _accumulate(agg, name, o)
        units = [p["Prefix"] for p in r.get("CommonPrefixes", [])]
        agg["by_bucket"].setdefault(name, [0, 0])
        if not units:
            if _scan_prefix(s3, name, "", SCAN_MAX_OBJECTS, deadline, agg):
                truncated = True
        else:
            # breadth-first: every top-level prefix gets its own budget so a
            # single huge prefix can't hide the rest of the data.
            for pre in units:
                if time.time() > deadline:
                    truncated = True
                    break
                if _scan_prefix(s3, name, pre, SCAN_PER_PREFIX, deadline, agg):
                    truncated = True

    total_bytes = agg["bytes"]
    total_objects = agg["objects"]

    growth = []
    cum_b = cum_o = 0
    for d in sorted(agg["by_day"]):
        c, b = agg["by_day"][d]
        cum_o += c
        cum_b += b
        growth.append({"date": d, "objects": c, "bytes": b, "cum_objects": cum_o, "cum_bytes": cum_b})

    buckets_out = sorted(
        [{"name": n, "objects": v[0], "bytes": v[1], "size_human": _human_size(v[1]),
          "pct": round(100 * v[1] / total_bytes, 1) if total_bytes else 0,
          "created": created_map.get(n)} for n, v in agg["by_bucket"].items()],
        key=lambda x: x["bytes"], reverse=True)

    types_out = sorted(
        [{"ext": k, "objects": v[0], "bytes": v[1], "size_human": _human_size(v[1])}
         for k, v in agg["by_ext"].items()], key=lambda x: x["objects"], reverse=True)[:12]

    size_order = ["<1 KB", "1 KB–1 MB", "1–10 MB", "10–100 MB", "100 MB–1 GB", ">1 GB"]
    size_out = [{"label": k, "count": agg["by_size"].get(k, 0)} for k in size_order]

    prefixes_out = sorted(
        [{"prefix": k, "objects": v[0], "bytes": v[1], "size_human": _human_size(v[1])}
         for k, v in agg["by_prefix"].items()], key=lambda x: x["bytes"], reverse=True)[:12]

    largest = sorted(agg["largest"], reverse=True)[:10]
    largest_out = [{"bucket": b, "key": k, "bytes": s, "size_human": _human_size(s), "modified": m}
                   for s, b, k, m in largest]

    recent = sorted(agg["recent"], reverse=True)[:15]
    recent_out = [{"modified": m, "bucket": b, "key": k, "bytes": s, "size_human": _human_size(s)}
                  for m, b, k, s in recent]

    return {
        "summary": {
            "buckets": len(names),
            "objects": total_objects,
            "bytes": total_bytes,
            "size_human": _human_size(total_bytes),
            "avg_object_bytes": int(total_bytes / total_objects) if total_objects else 0,
            "avg_object_human": _human_size(int(total_bytes / total_objects) if total_objects else 0),
            "largest_human": largest_out[0]["size_human"] if largest_out else "0 B",
            "newest": recent_out[0]["modified"] if recent_out else None,
            "oldest": growth[0]["date"] if growth else None,
            "truncated": truncated,
            "scanned": total_objects,
        },
        "buckets": buckets_out,
        "types": types_out,
        "size_histogram": size_out,
        "growth": growth,
        "top_prefixes": prefixes_out,
        "largest_objects": largest_out,
        "recent": recent_out,
    }


def _activity_view() -> dict[str, Any]:
    now = time.time()
    cutoff = now - 3600
    mins = [0] * 60
    by_op: dict[str, int] = {}
    total = 0
    for ts, op in list(_activity):
        if ts < cutoff:
            continue
        total += 1
        by_op[op] = by_op.get(op, 0) + 1
        idx = int((now - ts) // 60)
        if 0 <= idx < 60:
            mins[59 - idx] += 1
    last_ts = _activity[-1][0] if _activity else None
    return {
        "per_minute": mins,
        "by_op": [{"op": k, "count": v} for k, v in sorted(by_op.items(), key=lambda x: -x[1])],
        "total_last_hour": total,
        "last_activity": datetime.fromtimestamp(last_ts, tz=timezone.utc).isoformat() if last_ts else None,
    }


@router.get("/analytics")
async def analytics(refresh: bool = Query(False)):
    """Aggregated storage analytics for the dashboard (cached ~45s)."""
    _track("analytics")
    now = time.time()
    if not refresh and _analytics_cache["data"] is not None and now - _analytics_cache["ts"] < _ANALYTICS_TTL:
        data = _analytics_cache["data"]
    else:
        try:
            data = _scan_analytics()
            _analytics_cache["data"] = data
            _analytics_cache["ts"] = now
        except Exception as exc:
            return JSONResponse({"ok": False, "error": str(exc), "endpoint": S3_ENDPOINT}, status_code=502)
    return {"ok": True, "endpoint": S3_ENDPOINT, "generated_at": datetime.now(timezone.utc).isoformat(),
            "activity": _activity_view(), **data}