#!/usr/bin/env bash
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
BUDGET_FILE="${HOMELAB_RESOURCE_BUDGET_FILE:-${REPO_ROOT}/infra/resource-budgets.yml}"
KUBECONFIG_PATH="${KUBECONFIG:-${LAB_KUBECONFIG_PATH:-/home/jv/.kube/config}}"

failures=0
warnings=0

section() {
    printf '\n== %s ==\n' "$1"
}

warn() {
    warnings=$((warnings + 1))
    printf 'warn - %s\n' "$1"
}

fail() {
    failures=$((failures + 1))
    printf 'fail - %s\n' "$1"
}

have() {
    command -v "$1" >/dev/null 2>&1
}

host_budget() {
    section "Host Budget"

    if have free; then
        free -m | awk '/^Mem:/ {
          used_pct = int(($3 / $2) * 100)
          printf "Debian memory used        %s%%\n", used_pct
          if (used_pct > 80) exit 2
        }' || warn "Debian memory usage is above the default 80% budget"
    else
        warn "free command unavailable"
    fi

    if have df; then
        df -P / /data 2>/dev/null | awk 'NR > 1 { gsub("%", "", $5); free = 100 - $5; printf "%-25s free=%s%%\n", $6, free }' || true
    else
        warn "df command unavailable"
    fi
}

kubernetes_budget() {
    section "Kubernetes Budget"

    if ! have kubectl || [ ! -s "$KUBECONFIG_PATH" ]; then
        warn "kubectl or kubeconfig unavailable; skipping Kubernetes budget checks"
        return 0
    fi
    if ! kubectl --kubeconfig "$KUBECONFIG_PATH" get --raw=/readyz >/dev/null 2>&1; then
        warn "Kubernetes API unavailable; skipping Kubernetes budget checks"
        return 0
    fi

    python3 - "$BUDGET_FILE" "$KUBECONFIG_PATH" <<'PY'
import json
import subprocess
import sys

budget_file, kubeconfig = sys.argv[1:3]

try:
    import yaml
except ImportError:
    yaml = None

def load_budget(path):
    if yaml:
        with open(path, encoding="utf-8") as handle:
            return yaml.safe_load(handle) or {}
    return {
        "cluster": {
            "max_total_cpu_millicores": 9000,
            "max_total_memory_mib": 22000,
            "require_requests": True,
            "require_limits": True,
        },
        "namespaces": {},
    }

def cpu_to_m(value):
    if not value:
        return 0
    text = str(value)
    if text.endswith("m"):
        return int(text[:-1])
    return int(float(text) * 1000)

def mem_to_mib(value):
    if not value:
        return 0
    text = str(value)
    units = {
        "Ki": 1 / 1024,
        "Mi": 1,
        "Gi": 1024,
        "Ti": 1024 * 1024,
        "K": 1 / 1000,
        "M": 1000 / 1024,
        "G": 1000 * 1000 / 1024,
    }
    for suffix, multiplier in units.items():
        if text.endswith(suffix):
            return int(float(text[: -len(suffix)]) * multiplier)
    return int(float(text) / (1024 * 1024))

budget = load_budget(budget_file)
cluster_budget = budget.get("cluster", {})
namespace_budgets = budget.get("namespaces", {})
raw = subprocess.check_output(
    ["kubectl", "--kubeconfig", kubeconfig, "get", "pods", "-A", "-o", "json"],
    text=True,
)
pods = json.loads(raw).get("items", [])
namespace_totals = {}
missing = []

for pod in pods:
    ns = pod["metadata"]["namespace"]
    pod_name = pod["metadata"]["name"]
    namespace_totals.setdefault(ns, {"cpu": 0, "memory": 0})
    for container in pod.get("spec", {}).get("containers", []):
        name = container["name"]
        resources = container.get("resources", {})
        requests = resources.get("requests", {})
        limits = resources.get("limits", {})
        namespace_totals[ns]["cpu"] += cpu_to_m(requests.get("cpu"))
        namespace_totals[ns]["memory"] += mem_to_mib(requests.get("memory"))
        if cluster_budget.get("require_requests", True) and ("cpu" not in requests or "memory" not in requests):
            missing.append(f"{ns}/{pod_name}:{name} missing requests")
        if cluster_budget.get("require_limits", True) and ("cpu" not in limits or "memory" not in limits):
            missing.append(f"{ns}/{pod_name}:{name} missing limits")

total_cpu = sum(item["cpu"] for item in namespace_totals.values())
total_memory = sum(item["memory"] for item in namespace_totals.values())
print(f"cluster_requests_cpu_m={total_cpu}")
print(f"cluster_requests_memory_mib={total_memory}")

exit_code = 0
if total_cpu > int(cluster_budget.get("max_total_cpu_millicores", 9000)):
    print("fail - cluster CPU requests exceed budget")
    exit_code = 1
if total_memory > int(cluster_budget.get("max_total_memory_mib", 22000)):
    print("fail - cluster memory requests exceed budget")
    exit_code = 1

for ns, totals in sorted(namespace_totals.items()):
    ns_budget = namespace_budgets.get(ns)
    if not ns_budget:
        continue
    print(f"{ns}: cpu_m={totals['cpu']} memory_mib={totals['memory']}")
    if totals["cpu"] > int(ns_budget.get("max_cpu_millicores", 999999)):
        print(f"fail - {ns} CPU requests exceed budget")
        exit_code = 1
    if totals["memory"] > int(ns_budget.get("max_memory_mib", 999999)):
        print(f"fail - {ns} memory requests exceed budget")
        exit_code = 1

if missing:
    print("missing_resource_policy:")
    for item in missing[:50]:
        print(f"  {item}")
    if len(missing) > 50:
        print(f"  ... {len(missing) - 50} more")
    exit_code = 1

raise SystemExit(exit_code)
PY
}

section "Policy"
if [ -s "$BUDGET_FILE" ]; then
    printf 'budget_file=%s\n' "$BUDGET_FILE"
else
    fail "budget file is missing: $BUDGET_FILE"
fi

host_budget
kubernetes_budget || failures=$((failures + 1))

section "Summary"
printf 'failures=%s warnings=%s\n' "$failures" "$warnings"

if [ "$failures" -gt 0 ]; then
    exit 1
fi
