#!/usr/bin/python3
"""Stream one cdx Kubernetes Job and return its agent container exit code."""
from __future__ import annotations

import argparse
import json
import subprocess
import sys
import time


def kubectl(namespace: str, kubeconfig: str | None, *args: str) -> list[str]:
    command = ["kubectl"]
    if kubeconfig:
        command += ["--kubeconfig", kubeconfig]
    return [*command, "--namespace", namespace, *args]


def get_pod(namespace: str, kubeconfig: str | None, job_name: str) -> dict | None:
    completed = subprocess.run(
        kubectl(namespace, kubeconfig, "get", "pods", "-l", f"job-name={job_name}", "-o", "json"),
        text=True, capture_output=True, check=False, timeout=30,
    )
    if completed.returncode:
        raise RuntimeError((completed.stderr or completed.stdout).strip() or "pod lookup failed")
    items = json.loads(completed.stdout).get("items", [])
    if len(items) > 1:
        raise RuntimeError(f"Kubernetes Job {job_name} has multiple Pods")
    return items[0] if items else None


def startup_failure(pod: dict) -> str | None:
    status = pod.get("status", {})
    if status.get("phase") == "Failed":
        return str(status.get("message") or status.get("reason") or "Pod failed")
    fatal_waits = {"ImagePullBackOff", "ErrImagePull", "CreateContainerConfigError", "InvalidImageName", "CrashLoopBackOff"}
    for item in [*status.get("initContainerStatuses", []), *status.get("containerStatuses", [])]:
        name = item.get("name", "container")
        state = item.get("state", {})
        terminated = state.get("terminated")
        if isinstance(terminated, dict) and terminated.get("exitCode") not in {None, 0}:
            return f"{name} exited {terminated.get('exitCode')}"
        waiting = state.get("waiting")
        if isinstance(waiting, dict) and waiting.get("reason") in fatal_waits:
            return f"{name}: {waiting.get('reason')}"
    return None


def agent_status(pod: dict) -> tuple[str, int | None]:
    statuses = pod.get("status", {}).get("containerStatuses", [])
    agent = next((item for item in statuses if item.get("name") == "agent"), None)
    state = agent.get("state", {}) if agent else {}
    if "terminated" in state:
        code = state["terminated"].get("exitCode")
        return "terminated", code if isinstance(code, int) else 1
    if "running" in state:
        return "running", None
    return "waiting", None


def run(namespace: str, kubeconfig: str | None, job_name: str, timeout: int) -> int:
    deadline = time.monotonic() + timeout
    pod_name = ""
    while time.monotonic() < deadline:
        pod = get_pod(namespace, kubeconfig, job_name)
        if pod is not None:
            if failure := startup_failure(pod):
                raise RuntimeError(f"Kubernetes Job {job_name} could not start: {failure}")
            pod_name = pod.get("metadata", {}).get("name", "")
            state, code = agent_status(pod)
            if pod_name and state in {"running", "terminated"}:
                break
        time.sleep(1)
    else:
        raise TimeoutError(f"Kubernetes Job {job_name} did not start before its deadline")
    logs = subprocess.run(
        kubectl(namespace, kubeconfig, "logs", "-f", pod_name, "-c", "agent"),
        check=False, timeout=max(1, int(deadline - time.monotonic()) + 30),
    )
    if logs.returncode:
        print(
            f"cdx-k3s-job: log stream detached for {job_name}; continuing status supervision",
            file=sys.stderr,
        )
    while time.monotonic() < deadline:
        pod = get_pod(namespace, kubeconfig, job_name)
        if pod is not None:
            state, code = agent_status(pod)
            if state == "terminated":
                return code if code is not None else 1
        time.sleep(1)
    raise TimeoutError(f"Kubernetes Job {job_name} did not finish before its deadline")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--namespace", required=True)
    parser.add_argument("--job", required=True)
    parser.add_argument("--kubeconfig")
    parser.add_argument("--timeout", type=int, default=7200)
    args = parser.parse_args(argv)
    return run(args.namespace, args.kubeconfig, args.job, args.timeout)


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as exc:
        print(f"cdx-k3s-job: {exc}", file=sys.stderr)
        raise SystemExit(1)