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 buckets.
SCAN_MAX_OBJECTS = int(os.getenv("S3_SCAN_MAX_OBJECTS", "50000"))
SCAN_DEADLINE_S = float(os.getenv("S3_SCAN_DEADLINE_S", "12"))

# 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 _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"


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: int) -> 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"


@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=500),
):
    try:
        _track("list_objects")
        s3 = _client()
        resp = s3.list_objects_v2(Bucket=bucket, Prefix=prefix, Delimiter="/", MaxKeys=max_keys)
        folders = [
            {"type": "prefix", "name": p["Prefix"][len(prefix):].rstrip("/"), "prefix": p["Prefix"]}
            for p in resp.get("CommonPrefixes", [])
        ]
        objects = [
            {
                "type": "object",
                "key": o["Key"],
                "name": o["Key"][len(prefix):] if o["Key"].startswith(prefix) else o["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,
            }
            for o in resp.get("Contents", [])
            if o["Key"] != prefix
        ]
        return {
            "ok": True,
            "bucket": bucket,
            "prefix": prefix,
            "folders": folders,
            "objects": objects,
            "truncated": resp.get("IsTruncated", False),
        }
    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)



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]

    total_bytes = 0
    total_objects = 0
    by_bucket: dict[str, list[int]] = {}
    by_ext: dict[str, list[int]] = {}
    by_size: dict[str, int] = {}
    by_day: dict[str, list[int]] = {}
    by_prefix: dict[str, list[int]] = {}
    largest: list[tuple[int, str, str, str | None]] = []
    recent: list[tuple[str, str, str, int]] = []
    truncated = False
    scanned = 0

    for name in names:
        bc = bb = 0
        token = None
        while True:
            if time.time() > deadline or scanned >= SCAN_MAX_OBJECTS:
                truncated = True
                break
            kw: dict[str, Any] = {"Bucket": name, "MaxKeys": 1000}
            if token:
                kw["ContinuationToken"] = token
            try:
                r = s3.list_objects_v2(**kw)
            except Exception:
                break
            for o in r.get("Contents", []):
                sz = int(o.get("Size", 0) or 0)
                key = o["Key"]
                lm = o.get("LastModified")
                bc += 1
                bb += sz
                scanned += 1
                total_objects += 1
                total_bytes += sz
                e = by_ext.setdefault(_ext(key), [0, 0])
                e[0] += 1
                e[1] += sz
                sc = _size_class(sz)
                by_size[sc] = by_size.get(sc, 0) + 1
                if lm:
                    d = lm.astimezone(timezone.utc).strftime("%Y-%m-%d")
                    dd = by_day.setdefault(d, [0, 0])
                    dd[0] += 1
                    dd[1] += sz
                    recent.append((lm.isoformat(), name, key, sz))
                top = key.split("/", 1)[0] if "/" in key else "(root)"
                pp = by_prefix.setdefault(f"{name}/{top}", [0, 0])
                pp[0] += 1
                pp[1] += sz
                largest.append((sz, name, key, lm.isoformat() if lm else None))
                if scanned >= SCAN_MAX_OBJECTS:
                    truncated = True
                    break
            if r.get("IsTruncated") and not (time.time() > deadline or scanned >= SCAN_MAX_OBJECTS):
                token = r.get("NextContinuationToken")
                if not token:
                    break
            else:
                break
        by_bucket[name] = [bc, bb]

    # cumulative growth timeline
    growth = []
    cum_b = cum_o = 0
    for d in sorted(by_day):
        c, b = 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 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 by_ext.items()]), key=lambda x: x["bytes"], 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": 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 by_prefix.items()]), key=lambda x: x["bytes"], reverse=True)[:10]

    largest.sort(key=lambda x: x[0], reverse=True)
    largest_out = [{"bucket": b, "key": k, "bytes": s, "size_human": _human_size(s), "modified": m}
                   for s, b, k, m in largest[:10]]

    recent.sort(key=lambda x: x[0], reverse=True)
    recent_out = [{"modified": m, "bucket": b, "key": k, "bytes": s, "size_human": _human_size(s)}
                  for m, b, k, s in recent[:15]]

    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": scanned,
        },
        "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}