#!/usr/bin/env python3
"""Read-only scheduling-headroom preflight for State Hub Helm promotion.

STATE-WP-0091-T01. Does not start Helm, does not lower requests, and does not
claim a scheduling guarantee from an aggregate remaining-millicores sum.
"""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


CPU_FACTORS = {"n": 0.000001, "u": 0.001, "m": 1.0, "": 1000.0}
BYTE_FACTORS = {
    "": 1,
    "Ki": 1024,
    "Mi": 1024**2,
    "Gi": 1024**3,
    "Ti": 1024**4,
    "K": 1000,
    "M": 1000**2,
    "G": 1000**3,
    "T": 1000**4,
}

# Fallback surge figures, used only when no values file can be read. With
# maxSurge=1 / maxUnavailable=0 the surge pod carries the *target* release's
# API request, so the real figure is derived from the Helm values the release
# will deploy (STATE-WP-0099). The historical 100m constant under-modelled the
# 250m API request; tests keep these fallbacks equal to the shipped values.
DEFAULT_API_SURGE_CPU_M = 390
DEFAULT_MCP_SURGE_CPU_M = 10
DEFAULT_MIGRATE_CPU_M = 50
DEFAULT_API_SURGE_MEMORY = 512 * 1024**2
DEFAULT_FRESHNESS_SECONDS = 15 * 60

REPO_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_VALUES_FILES = (
    REPO_ROOT / "deploy/railiance/apps/charts/state-hub/values.yaml",
    REPO_ROOT / "deploy/railiance/apps/helm/state-hub-values.yaml",
)


def cpu_millicores(value: str | None) -> int:
    import re

    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)(n|u|m)?", str(value or "0"))
    if not match:
        raise ValueError(f"unsupported CPU quantity: {value}")
    return round(float(match.group(1)) * CPU_FACTORS[match.group(2) or ""])


def bytes_value(value: str | None) -> int:
    import re

    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)(Ki|Mi|Gi|Ti|K|M|G|T)?", str(value or "0"))
    if not match:
        raise ValueError(f"unsupported byte quantity: {value}")
    return round(float(match.group(1)) * BYTE_FACTORS[match.group(2) or ""])


def pod_effective_requests(pod: dict[str, Any]) -> tuple[int, int]:
    """Pod request is max(sum(containers), max(init containers))."""
    spec = pod.get("spec") or {}
    containers = spec.get("containers") or []
    inits = spec.get("initContainers") or []
    cpu = sum(cpu_millicores((c.get("resources") or {}).get("requests", {}).get("cpu")) for c in containers)
    mem = sum(bytes_value((c.get("resources") or {}).get("requests", {}).get("memory")) for c in containers)
    init_cpu = max((cpu_millicores((c.get("resources") or {}).get("requests", {}).get("cpu")) for c in inits), default=0)
    init_mem = max((bytes_value((c.get("resources") or {}).get("requests", {}).get("memory")) for c in inits), default=0)
    return max(cpu, init_cpu), max(mem, init_mem)


def _deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
    merged = dict(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _deep_merge(merged[key], value)
        else:
            merged[key] = value
    return merged


def _request(section: dict[str, Any] | None, key: str) -> Any:
    return (((section or {}).get("resources") or {}).get("requests") or {}).get(key)


def release_from_values(values: dict[str, Any]) -> dict[str, Any]:
    """Derive surge figures from merged Helm values (chart defaults + overrides).

    Helm merges `-f` files in order; pass them already merged. Only keys that are
    present are returned, so missing sections fall back to the defaults.
    """
    release: dict[str, Any] = {"surge_source": "values"}
    api_cpu = _request(values, "cpu")
    if api_cpu is not None:
        release["api_surge_cpu_m"] = cpu_millicores(api_cpu)
    api_mem = _request(values, "memory")
    if api_mem is not None:
        release["api_surge_memory_bytes"] = bytes_value(api_mem)
    mcp = values.get("mcp") or {}
    if mcp.get("enabled") is False:
        release["mcp_surge_cpu_m"] = 0
    elif _request(mcp, "cpu") is not None:
        release["mcp_surge_cpu_m"] = cpu_millicores(_request(mcp, "cpu"))
    migrations = values.get("migrations") or {}
    if migrations.get("enabled") is False:
        release["migrate_cpu_m"] = 0
    elif _request(migrations, "cpu") is not None:
        release["migrate_cpu_m"] = cpu_millicores(_request(migrations, "cpu"))
    return release


def load_values(paths: list[Path] | tuple[Path, ...]) -> dict[str, Any]:
    import yaml

    merged: dict[str, Any] = {}
    for path in paths:
        merged = _deep_merge(merged, yaml.safe_load(Path(path).read_text()) or {})
    return merged


def evaluate(observation: dict[str, Any], release: dict[str, Any] | None = None) -> dict[str, Any]:
    """Return a preflight report. `ok` is whether Helm may start, not a guarantee."""
    release = release or {}
    api_cpu = int(release.get("api_surge_cpu_m", DEFAULT_API_SURGE_CPU_M))
    mcp_cpu = int(release.get("mcp_surge_cpu_m", DEFAULT_MCP_SURGE_CPU_M))
    migrate_cpu = int(release.get("migrate_cpu_m", DEFAULT_MIGRATE_CPU_M))
    api_mem = int(release.get("api_surge_memory_bytes", DEFAULT_API_SURGE_MEMORY))
    freshness_limit = int(release.get("freshness_seconds", DEFAULT_FRESHNESS_SECONDS))
    atomic = bool(release.get("atomic", True))
    surge_source = str(release.get("surge_source", "release" if "api_surge_cpu_m" in release else "defaults"))

    reasons: list[str] = []
    notes: list[str] = []
    nodes = observation.get("nodes") or []
    if not nodes:
        reasons.append("no node capacity evidence")
    remaining_cpu = 0
    remaining_mem = 0
    for node in nodes:
        if node.get("unschedulable") or not node.get("ready", True):
            reasons.append(f"node {node.get('name')} is not schedulable")
            continue
        remaining_cpu += int(node["allocatable_cpu_m"]) - int(node["allocated_cpu_m"])
        remaining_mem += int(node["allocatable_memory_bytes"]) - int(node["allocated_memory_bytes"])

    pending = observation.get("pending_unrelated") or []
    pending_cpu = sum(int(p.get("cpu_m", 0)) for p in pending)
    pending_mem = sum(int(p.get("memory_bytes", 0)) for p in pending)

    observed_at = observation.get("observed_at")
    freshness_seconds = observation.get("freshness_seconds")
    if freshness_seconds is None and observed_at:
        try:
            stamp = datetime.fromisoformat(str(observed_at).replace("Z", "+00:00"))
            freshness_seconds = int((datetime.now(timezone.utc) - stamp).total_seconds())
        except ValueError:
            freshness_seconds = None
            reasons.append("observed_at is not parseable")
    if freshness_seconds is not None and freshness_seconds > freshness_limit:
        reasons.append(
            f"capacity evidence is {freshness_seconds}s old (limit {freshness_limit}s)"
        )

    if remaining_cpu < migrate_cpu:
        reasons.append(
            f"migration hook needs {migrate_cpu}m CPU; remaining {remaining_cpu}m"
        )
    if remaining_cpu < api_cpu:
        reasons.append(
            f"API maxSurge=1 needs {api_cpu}m CPU; remaining {remaining_cpu}m "
            f"(Insufficient CPU class; revision 61 had 65m remaining)"
        )
    if remaining_mem < api_mem:
        reasons.append(
            f"API maxSurge=1 needs {api_mem} bytes memory; remaining {remaining_mem}"
        )
    concurrent = api_cpu + mcp_cpu
    if remaining_cpu < concurrent:
        notes.append(
            f"API+MCP concurrent surge would need {concurrent}m; remaining {remaining_cpu}m. "
            "Helm may roll both deployments together; this is not a scheduling guarantee."
        )
    effective_cpu = remaining_cpu - pending_cpu
    effective_mem = remaining_mem - pending_mem
    if pending and effective_cpu < api_cpu:
        reasons.append(
            f"unrelated pending demand {pending_cpu}m leaves {effective_cpu}m "
            f"against API surge {api_cpu}m"
        )
    if pending and effective_mem < api_mem:
        reasons.append(
            f"unrelated pending memory {pending_mem} leaves {effective_mem} against API surge"
        )

    live_api_cpu = observation.get("live_api_request_cpu_m")
    remaining_after_rollout = None
    if live_api_cpu is not None:
        remaining_after_rollout = remaining_cpu - (api_cpu - int(live_api_cpu))
        if int(live_api_cpu) != api_cpu:
            notes.append(
                f"release changes the API request {int(live_api_cpu)}m -> {api_cpu}m; "
                f"steady remaining after rollout is {remaining_after_rollout}m"
            )
        if remaining_after_rollout < api_cpu:
            notes.append(
                f"after this rollout {remaining_after_rollout}m remains, below the {api_cpu}m "
                "surge; the next release will be refused until headroom is freed"
            )

    if not atomic:
        reasons.append("promotion path must keep helm --atomic; preflight will not start a non-atomic upgrade")

    ok = not reasons
    if ok and remaining_cpu - api_cpu < 50:
        notes.append(
            f"remaining {remaining_cpu}m is only {remaining_cpu - api_cpu}m above the "
            f"{api_cpu}m API surge (revision 63 was 105m). Not factory capacity admission."
        )
    notes.append("Aggregate remaining millicores is not a kube-scheduler guarantee.")
    notes.append("Preflight does not lower requests and does not start Helm.")

    return {
        "schema": "state-hub.release-headroom-preflight.v1",
        "ok": ok,
        "observed_at": observed_at,
        "freshness_seconds": freshness_seconds,
        "remaining_cpu_m": remaining_cpu,
        "remaining_memory_bytes": remaining_mem,
        "pending_unrelated_cpu_m": pending_cpu,
        "api_surge_cpu_m": api_cpu,
        "surge_source": surge_source,
        "live_api_request_cpu_m": live_api_cpu,
        "remaining_after_rollout_cpu_m": remaining_after_rollout,
        "mcp_surge_cpu_m": mcp_cpu,
        "migrate_cpu_m": migrate_cpu,
        "atomic": atomic,
        "reasons": reasons,
        "notes": notes,
        "hook_order": [
            "pre-upgrade migrate job (helm.sh/hook-weight -5)",
            "API RollingUpdate maxSurge=1 maxUnavailable=0",
            "MCP RollingUpdate maxSurge=1 maxUnavailable=0",
        ],
    }


def collect_live(kubectl: str = "kubectl") -> dict[str, Any]:
    nodes = json.loads(subprocess.check_output([kubectl, "get", "nodes", "-o", "json"], timeout=60))
    pods = json.loads(subprocess.check_output([kubectl, "get", "pods", "-A", "-o", "json"], timeout=60))
    allocated: dict[str, dict[str, int]] = {}
    pending_unrelated: list[dict[str, Any]] = []
    now = datetime.now(timezone.utc)
    for node in nodes.get("items") or []:
        name = node["metadata"]["name"]
        alloc = node.get("status", {}).get("allocatable") or {}
        allocated[name] = {
            "allocatable_cpu_m": cpu_millicores(alloc.get("cpu")),
            "allocatable_memory_bytes": bytes_value(alloc.get("memory")),
            "allocated_cpu_m": 0,
            "allocated_memory_bytes": 0,
            "ready": any(
                c.get("type") == "Ready" and c.get("status") == "True"
                for c in (node.get("status") or {}).get("conditions") or []
            ),
            "unschedulable": bool((node.get("spec") or {}).get("unschedulable")),
            "name": name,
        }
    for pod in pods.get("items") or []:
        phase = (pod.get("status") or {}).get("phase")
        ns = pod.get("metadata", {}).get("namespace")
        name = pod.get("metadata", {}).get("name")
        cpu, mem = pod_effective_requests(pod)
        node_name = (pod.get("spec") or {}).get("nodeName")
        if phase == "Pending" and not node_name:
            if ns != "state-hub":
                pending_unrelated.append({"name": f"{ns}/{name}", "cpu_m": cpu, "memory_bytes": mem})
            continue
        if phase in {"Succeeded", "Failed"}:
            continue
        if node_name in allocated:
            allocated[node_name]["allocated_cpu_m"] += cpu
            allocated[node_name]["allocated_memory_bytes"] += mem
    live_api_cpu = None
    try:
        deploy = json.loads(subprocess.check_output(
            [kubectl, "-n", "state-hub", "get", "deployment", "state-hub", "-o", "json"], timeout=60
        ))
        for container in deploy["spec"]["template"]["spec"]["containers"]:
            if container.get("name") == "state-hub":
                live_api_cpu = cpu_millicores(_request(container, "cpu"))
    except (subprocess.SubprocessError, KeyError, ValueError, OSError):
        live_api_cpu = None
    return {
        "observed_at": now.isoformat(),
        "freshness_seconds": 0,
        "live_api_request_cpu_m": live_api_cpu,
        "nodes": list(allocated.values()),
        "pending_unrelated": pending_unrelated,
        "source": "kubectl",
    }


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--fixture", type=Path, help="JSON observation instead of live kubectl")
    parser.add_argument("--release", type=Path, help="JSON release surge description (overrides values)")
    parser.add_argument(
        "--values", type=Path, action="append",
        help="Helm values file, in -f order (repeatable). Default: chart values + production values",
    )
    parser.add_argument("--kubectl", default="kubectl")
    args = parser.parse_args(argv)
    if args.fixture:
        observation = json.loads(args.fixture.read_text())
    else:
        observation = collect_live(args.kubectl)
    values_files = args.values or [p for p in DEFAULT_VALUES_FILES if p.exists()]
    release = release_from_values(load_values(values_files)) if values_files else {}
    if args.release:
        explicit = json.loads(args.release.read_text())
        release.update(explicit)
        if "api_surge_cpu_m" in explicit:
            release["surge_source"] = "release"
    report = evaluate(observation, release)
    json.dump(report, sys.stdout, indent=2)
    sys.stdout.write("\n")
    return 0 if report["ok"] else 1


if __name__ == "__main__":
    raise SystemExit(main())