#!/usr/bin/env python3
"""Report whether a validated Hermes image has converged through Flux."""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
import time
from dataclasses import asdict, dataclass


NAMESPACE = "hermes"
FLUX_NAMESPACE = "flux-system"
AUTOMATION = "hermes"
FLUX_KUSTOMIZATION = "hermes"
REVISION_PATTERN = re.compile(r"^[0-9a-f]{40}$")
RELEASE_TAG_PATTERN = re.compile(
    r"^git-(?P<revision>[0-9a-f]{40})-build-(?P<build>[1-9][0-9]*)-release$"
)
COMPONENTS = {
    "agent": {
        "policy": "hermes-agent-release",
        "repository": "registry.bstein.dev/bstein/hermes-agent",
        "workloads": (("deployment", "hermes-agent", "hermes-agent"),),
    },
    "webui": {
        "policy": "hermes-webui-release",
        "repository": "registry.bstein.dev/bstein/hermes-webui",
        "workloads": (
            ("deployment", "hermes", "hermes"),
            ("statefulset", "hermes-chat-tenant", "hermes-chat-tenant"),
        ),
    },
    "router": {
        "policy": "hermes-chat-router-release",
        "repository": "registry.bstein.dev/bstein/hermes-chat-router",
        "workloads": (
            ("deployment", "hermes-chat-router", "hermes-chat-router"),
        ),
    },
    "stt": {
        "policy": "hermes-stt-release",
        "repository": "registry.bstein.dev/bstein/hermes-jetson-stt",
        "workloads": (("deployment", "hermes-stt", "hermes-stt"),),
    },
    "tts": {
        "policy": "hermes-tts-release",
        "repository": "registry.bstein.dev/bstein/hermes-jetson-tts",
        "workloads": (("deployment", "hermes-tts", "hermes-tts"),),
    },
}


@dataclass(frozen=True)
class WorkloadStatus:
    """Desired and observed state for one image-consuming workload."""

    kind: str
    name: str
    desired_replicas: int
    ready_replicas: int
    desired_image_matches: bool
    ready_pods: int
    matching_pods: int
    converged: bool


def _kubectl_json(*command: str) -> dict:
    """Read one Kubernetes object as JSON without shell or credential output."""
    result = subprocess.run(
        ["kubectl", *command, "-o", "json"],
        check=True,
        capture_output=True,
        text=True,
        timeout=30,
    )
    value = json.loads(result.stdout)
    if not isinstance(value, dict):
        raise ValueError("kubectl returned a non-object JSON value")
    return value


def _condition_ready(resource: dict) -> bool:
    """Return true only for an explicit current Ready condition."""
    generation = resource.get("metadata", {}).get("generation")
    for condition in resource.get("status", {}).get("conditions", []):
        if condition.get("type") != "Ready" or condition.get("status") != "True":
            continue
        observed = condition.get("observedGeneration")
        return generation is None or observed is None or int(observed) >= int(generation)
    return False


def _release_ref(policy: dict) -> tuple[str | None, str | None, str | None, int | None]:
    """Extract the immutable release tag, digest, source revision, and build."""
    ref = policy.get("status", {}).get("latestRef", {})
    tag = ref.get("tag")
    digest = ref.get("digest")
    match = RELEASE_TAG_PATTERN.fullmatch(tag or "")
    if not match or not re.fullmatch(r"sha256:[0-9a-f]{64}", digest or ""):
        return tag, digest, None, None
    return tag, digest, match.group("revision"), int(match.group("build"))


def _desired_images(workload: dict) -> list[str]:
    """Return pod-template container and init-container image references."""
    spec = workload.get("spec", {}).get("template", {}).get("spec", {})
    containers = list(spec.get("initContainers", [])) + list(spec.get("containers", []))
    return [item.get("image", "") for item in containers]


def _rollout_ready(kind: str, workload: dict) -> tuple[int, int, bool]:
    """Evaluate generation and replica convergence for a workload controller."""
    metadata = workload.get("metadata", {})
    spec = workload.get("spec", {})
    status = workload.get("status", {})
    desired = int(spec.get("replicas", 1))
    ready = int(status.get("readyReplicas", 0))
    generation_ready = int(status.get("observedGeneration", 0)) >= int(
        metadata.get("generation", 1)
    )
    if kind == "deployment":
        controller_ready = (
            int(status.get("updatedReplicas", 0)) >= desired
            and int(status.get("availableReplicas", 0)) >= desired
        )
    else:
        controller_ready = (
            int(status.get("currentReplicas", 0)) >= desired
            and int(status.get("updatedReplicas", 0)) >= desired
            and status.get("currentRevision") == status.get("updateRevision")
        )
    return desired, ready, generation_ready and controller_ready and ready >= desired


def _pod_image_status(label: str, repository: str, digest: str) -> tuple[int, int]:
    """Count ready pods whose relevant running containers use the exact digest."""
    pod_list = _kubectl_json(
        "-n", NAMESPACE, "get", "pods", "-l", f"app={label}"
    )
    matching = 0
    ready = 0
    for pod in pod_list.get("items", []):
        statuses = list(pod.get("status", {}).get("initContainerStatuses", []))
        statuses.extend(pod.get("status", {}).get("containerStatuses", []))
        relevant = [
            item
            for item in statuses
            if item.get("image", "").startswith(repository)
            or f"{repository}@" in (item.get("imageID") or "")
        ]
        if not relevant:
            continue
        if all((item.get("imageID") or "").endswith(f"@{digest}") for item in relevant):
            matching += 1
            conditions = pod.get("status", {}).get("conditions", [])
            if any(
                item.get("type") == "Ready" and item.get("status") == "True"
                for item in conditions
            ):
                ready += 1
    return matching, ready


def _workload_status(
    kind: str, name: str, label: str, repository: str, digest: str
) -> WorkloadStatus:
    """Read desired image, controller rollout, and running-pod digest evidence."""
    workload = _kubectl_json("-n", NAMESPACE, "get", kind, name)
    desired, ready_replicas, controller_ready = _rollout_ready(kind, workload)
    relevant_images = [
        image for image in _desired_images(workload) if image.startswith(repository)
    ]
    desired_matches = bool(relevant_images) and all(
        image.endswith(f"@{digest}") for image in relevant_images
    )
    matching_pods, ready_pods = _pod_image_status(label, repository, digest)
    converged = (
        desired_matches
        and controller_ready
        and matching_pods >= desired
        and ready_pods >= desired
    )
    return WorkloadStatus(
        kind=kind,
        name=name,
        desired_replicas=desired,
        ready_replicas=ready_replicas,
        desired_image_matches=desired_matches,
        ready_pods=ready_pods,
        matching_pods=matching_pods,
        converged=converged,
    )


def inspect_release(component: str, revision: str) -> dict:
    """Return safe end-to-end release state for one exact reviewed revision."""
    if component not in COMPONENTS:
        raise ValueError("component must be agent, router, webui, stt, or tts")
    if not REVISION_PATTERN.fullmatch(revision):
        raise ValueError("revision must be a lowercase full 40-character commit")
    config = COMPONENTS[component]
    policy = _kubectl_json(
        "-n", NAMESPACE, "get", "imagepolicy", config["policy"]
    )
    tag, digest, selected_revision, build = _release_ref(policy)
    result = {
        "component": component,
        "requested_revision": revision,
        "selected_revision": selected_revision,
        "release_tag": tag,
        "digest": digest,
        "build": build,
        "stage": "image_policy_pending",
        "converged": False,
        "workloads": [],
    }
    if selected_revision != revision or digest is None:
        if selected_revision is not None:
            result["stage"] = "different_revision_selected"
        return result
    automation = _kubectl_json(
        "-n", NAMESPACE, "get", "imageupdateautomation", AUTOMATION
    )
    flux = _kubectl_json(
        "-n", FLUX_NAMESPACE, "get", "kustomization", FLUX_KUSTOMIZATION
    )
    result["image_policy_ready"] = _condition_ready(policy)
    result["image_automation_ready"] = _condition_ready(automation)
    result["flux_ready"] = _condition_ready(flux)
    result["flux_revision"] = flux.get("status", {}).get("lastAppliedRevision")
    workloads = [
        _workload_status(kind, name, label, config["repository"], digest)
        for kind, name, label in config["workloads"]
    ]
    result["workloads"] = [asdict(item) for item in workloads]
    if not all((result["image_policy_ready"], result["image_automation_ready"])):
        result["stage"] = "image_automation_pending"
    elif not result["flux_ready"] or not all(item.desired_image_matches for item in workloads):
        result["stage"] = "flux_apply_pending"
    elif not all(item.converged for item in workloads):
        result["stage"] = "rollout_pending"
    else:
        result["stage"] = "converged"
        result["converged"] = True
    return result


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """Parse one bounded release-follow request."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--component", required=True, choices=sorted(COMPONENTS))
    parser.add_argument("--revision", required=True)
    parser.add_argument("--wait", action="store_true")
    parser.add_argument("--timeout", type=int, default=1800)
    parser.add_argument("--poll", type=int, default=15)
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    """Print safe JSON immediately or wait for exact live convergence."""
    args = parse_args(argv)
    if args.timeout < 1 or args.poll < 1:
        print("timeout and poll must be positive", file=sys.stderr)
        return 2
    deadline = time.monotonic() + args.timeout
    try:
        while True:
            status = inspect_release(args.component, args.revision)
            if status["converged"]:
                print(json.dumps(status, indent=2, sort_keys=True))
                return 0
            if not args.wait or time.monotonic() >= deadline:
                status["timed_out"] = bool(args.wait)
                print(json.dumps(status, indent=2, sort_keys=True))
                return 2 if args.wait else 1
            time.sleep(args.poll)
    except (OSError, ValueError, subprocess.SubprocessError, json.JSONDecodeError) as exc:
        print(f"Hermes release status failed: {exc}", file=sys.stderr)
        return 1


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