#!/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}}"
DETAILS=false
MAX_EXAMPLES="${LAB_RESOURCE_BUDGET_EXAMPLES:-12}"

failures=0
warnings=0

usage() {
    cat <<'EOF'
Usage: ./jeannie resource-budget [--details]

Print host and Kubernetes resource budget checks.

Options:
  --details  Show every pod/container missing requests or limits.
EOF
}

while [ "$#" -gt 0 ]; do
    case "$1" in
        --details|--verbose)
            DETAILS=true
            ;;
        -h|--help|help)
            usage
            exit 0
            ;;
        *)
            usage >&2
            exit 2
            ;;
    esac
    shift
done

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" "$DETAILS" "$MAX_EXAMPLES" <<'PY'
import json
import subprocess
import sys
from collections import Counter, defaultdict

budget_file, kubeconfig, details_arg, max_examples_arg = sys.argv[1:5]
show_details = details_arg.lower() == "true"
max_examples = int(max_examples_arg)

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 = []
missing_by_namespace = defaultdict(Counter)

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")
            missing_by_namespace[ns]["requests"] += 1
        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")
            missing_by_namespace[ns]["limits"] += 1

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(f"missing_resource_policy: total={len(missing)}")
    print("by_namespace:")
    for ns, counts in sorted(
        missing_by_namespace.items(),
        key=lambda item: (item[1]["requests"] + item[1]["limits"], item[0]),
        reverse=True,
    ):
        print(f"  {ns}: requests_missing={counts['requests']} limits_missing={counts['limits']}")

    if show_details:
        print("details:")
        for item in missing:
            print(f"  {item}")
    else:
        print("examples:")
        for item in missing[:max_examples]:
            print(f"  {item}")
        if len(missing) > max_examples:
            print(f"  ... {len(missing) - max_examples} more")
            print("  run ./jeannie resource-budget --details for the full list")
        print("next_steps:")
        print("  1. Add requests before limits for app workloads you own.")
        print("  2. Leave kube-system/control-plane static pods for a separate pass.")
        print("  3. Use ./jeannie control-plane taint after Pimox workers are Ready.")
    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