prole/knoe/core/ops/registry.py

from __future__ import annotations

from pathlib import Path
import json

from ._services_common import (
    _LogFn,
    _detect_mode,
    _docker,
    _ensure_namespace,
    _exists,
    _k3d,
    _kubectl,
    _log,
    _manifest_path,
    _prune_named_workload_other_namespaces,
    _reconcile_deployment_replicasets,
    _registry_namespace,
    _to_bool,
    _wait_rollout,
)


def _ensure_k3d_registry(*, env: dict | None = None, log: _LogFn | None = None) -> None:
    name = str((env or {}).get("K3D_REGISTRY_NAME") or "prole-registry")
    port = str((env or {}).get("REGISTRY_PORT") or "5000")

    listed = _k3d(["registry", "list"], env=env, timeout=60)
    if listed.returncode == 0 and name in listed.stdout:
        _log(log, f"[REGISTRY] k3d registry {name} already exists")
        return

    _log(log, f"[REGISTRY] Creating k3d registry {name} on port {port}")
    created = _k3d(["registry", "create", name, "--port", f"{port}:{port}"], env=env, timeout=180)
    if created.returncode == 0:
        return

    # fallback docker registry
    _log(log, "[REGISTRY] k3d registry creation failed; trying docker registry fallback")
    running = _docker(["ps", "--format", "{{.Names}}"], env=env, timeout=30)
    if running.returncode == 0 and name in running.stdout.splitlines():
        return
    _docker(["rm", "-f", name], env=env, timeout=30)
    _docker(
        [
            "run",
            "-d",
            "--restart=always",
            "-p",
            f"{port}:5000",
            "--name",
            name,
            "registry:2",
        ],
        env=env,
        timeout=120,
        check=True,
    )


def _manifest(project_root: str | Path) -> Path:
    return _manifest_path(project_root, "k8s", "registry", "deployment.yaml")


def _deployment_is_available(
    *,
    deployment: str,
    namespace: str,
    env: dict | None = None,
) -> bool:
    res = _kubectl(
        ["-n", namespace, "get", "deployment", deployment, "-o", "json"],
        env=env,
        timeout=20,
    )
    if res.returncode != 0 or not res.stdout.strip():
        return False
    try:
        payload = json.loads(res.stdout)
    except Exception:
        return False
    spec = payload.get("spec") or {}
    status = payload.get("status") or {}
    desired = int(spec.get("replicas") or 1)
    available = int(status.get("availableReplicas") or 0)
    return desired > 0 and available >= desired


def initialize(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    project_root: str | Path = ".",
    log: _LogFn | None = None,
    mode: str | None = None,
) -> None:
    update(namespace=namespace, env=env, project_root=project_root, log=log, mode=mode)


def start(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    project_root: str | Path = ".",
    log: _LogFn | None = None,
    mode: str | None = None,
) -> None:
    update(namespace=namespace, env=env, project_root=project_root, log=log, mode=mode)


def update(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    project_root: str | Path = ".",
    log: _LogFn | None = None,
    mode: str | None = None,
) -> None:
    effective_mode = _detect_mode(mode, env)
    if effective_mode == "k3d":
        _ensure_k3d_registry(env=env, log=log)
        return

    target_ns = _registry_namespace(namespace, env)
    _ensure_namespace(target_ns, env)
    _prune_named_workload_other_namespaces(
        kind="deployment",
        name="registry",
        target_namespace=target_ns,
        label_selector="app=registry",
        env=env,
        log=log,
    )
    _prune_named_workload_other_namespaces(
        kind="service",
        name="registry",
        target_namespace=target_ns,
        label_selector="app=registry",
        env=env,
        log=log,
    )

    force_apply = _to_bool((env or {}).get("PROLE_REGISTRY_FORCE_APPLY"), default=False)
    if not force_apply and _exists("deployment", "registry", target_ns, env=env) and _exists(
        "service", "registry", target_ns, env=env
    ):
        _reconcile_deployment_replicasets(
            deployment="registry",
            namespace=target_ns,
            label_selector="app=registry",
            env=env,
            log=log,
        )
        if _deployment_is_available(deployment="registry", namespace=target_ns, env=env):
            _log(
                log,
                f"[REGISTRY] Healthy deployment already present in namespace {target_ns}; skipping re-apply",
            )
            return
        _log(
            log,
            f"[REGISTRY] Existing deployment in namespace {target_ns} is not ready; applying manifest for recovery",
        )

    if force_apply:
        _log(
            log,
            f"[REGISTRY] Force apply enabled; reconciling manifest in namespace {target_ns}",
        )

    manifest = _manifest(project_root)
    if not manifest.exists():
        raise RuntimeError(f"Registry manifest not found: {manifest}")
    _log(log, f"[REGISTRY] Applying {manifest}")
    _kubectl(
        ["-n", target_ns, "apply", "-f", str(manifest)],
        env=env,
        timeout=180,
        check=True,
    )
    _reconcile_deployment_replicasets(
        deployment="registry",
        namespace=target_ns,
        label_selector="app=registry",
        env=env,
        log=log,
    )
    _wait_rollout("deployment", "registry", target_ns, env=env)


def stop(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    mode: str | None = None,
    log: _LogFn | None = None,
) -> None:
    effective_mode = _detect_mode(mode, env)
    if effective_mode == "k3d":
        name = str((env or {}).get("K3D_REGISTRY_NAME") or "prole-registry")
        _log(log, f"[REGISTRY] Removing k3d registry {name}")
        _k3d(["registry", "delete", name], env=env, timeout=120)
        _docker(["rm", "-f", name], env=env, timeout=30)
        return

    target_ns = _registry_namespace(namespace, env)
    if _exists("deployment", "registry", target_ns, env=env):
        _kubectl(
            ["-n", target_ns, "scale", "deployment/registry", "--replicas=0"],
            env=env,
            timeout=90,
            check=True,
        )


def restart(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    mode: str | None = None,
    log: _LogFn | None = None,
) -> None:
    effective_mode = _detect_mode(mode, env)
    if effective_mode == "k3d":
        update(namespace=namespace, env=env, mode=effective_mode, log=log)
        return

    target_ns = _registry_namespace(namespace, env)
    if not _exists("deployment", "registry", target_ns, env=env):
        update(namespace=namespace, env=env, mode=effective_mode, log=log)
        return
    _kubectl(
        ["-n", target_ns, "rollout", "restart", "deployment/registry"],
        env=env,
        timeout=120,
        check=True,
    )
    _wait_rollout("deployment", "registry", target_ns, env=env)


def status(
    *,
    namespace: str | None = None,
    env: dict | None = None,
    mode: str | None = None,
) -> bool:
    effective_mode = _detect_mode(mode, env)
    if effective_mode == "k3d":
        name = str((env or {}).get("K3D_REGISTRY_NAME") or "prole-registry")
        reg = _k3d(["registry", "list"], env=env, timeout=30)
        return _to_bool(reg.returncode == 0 and name in reg.stdout)

    target_ns = _registry_namespace(namespace, env)
    dep = _kubectl(
        ["-n", target_ns, "get", "deployment", "registry", "-o", "jsonpath={.status.readyReplicas}"],
        env=env,
        timeout=20,
    )
    svc = _kubectl(["-n", target_ns, "get", "svc", "registry"], env=env, timeout=20)
    return _to_bool(
        dep.returncode == 0 and (dep.stdout or "0").strip() not in {"", "0"} and svc.returncode == 0
    )