prole/knoe/core/ops/storage.py

from __future__ import annotations

import hashlib
import json
import os
import re
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any


DEFAULT_SYNOLOGY_ROOTS: tuple[str, ...] = (
    "/synology/d001",
    "/synology/d002",
    "/synology/d003",
    "/synology/d004",
)


class StorageProvisioningError(RuntimeError):
    pass


@dataclass(slots=True, frozen=True)
class ClusterStorageSpec:
    namespace: str
    cluster_name: str
    node_name: str
    storage_class_name: str = "synology-iscsi"
    service_name: str = "knoe-db"
    root_namespace: str = "knoe"
    postgres_uid: int = 100
    postgres_gid: int = 101
    directory_mode: int = 0o700
    volume_capacity: str = "29Gi"
    synology_roots: tuple[str, ...] = DEFAULT_SYNOLOGY_ROOTS
    kubectl_bin: str = "kubectl"
    mount_check_bin: str = "findmnt"
    command_env: dict[str, str] | None = None


@dataclass(slots=True, frozen=True)
class ClusterStoragePaths:
    data_root: str
    wal_root: str
    data_volume: str
    wal_volume: str
    data_path: str
    wal_path: str


@dataclass(slots=True, frozen=True)
class ProvisionedStorage:
    data_pv_name: str
    wal_pv_name: str
    data_path: str
    wal_path: str
    data_selector: dict[str, str]
    wal_selector: dict[str, str]


def _sanitize_dns_label(value: str, max_len: int = 253) -> str:
    lowered = (value or "").strip().lower()
    cleaned = re.sub(r"[^a-z0-9.-]+", "-", lowered)
    cleaned = re.sub(r"-+", "-", cleaned).strip("-.")
    if not cleaned:
        cleaned = "pv"
    if len(cleaned) <= max_len:
        return cleaned
    digest = hashlib.sha1(cleaned.encode("utf-8")).hexdigest()[:8]
    head = cleaned[: max_len - 9].rstrip("-.")
    return f"{head}-{digest}"


def _normalize_posix_path(path: str) -> str:
    parts = [p for p in path.split("/") if p]
    return "/" + "/".join(parts)


def _paths_overlap(path_a: str, path_b: str) -> bool:
    a = _normalize_posix_path(path_a)
    b = _normalize_posix_path(path_b)
    return a == b or a.startswith(f"{b}/") or b.startswith(f"{a}/")


def _command_env(spec: ClusterStorageSpec) -> dict[str, str] | None:
    return dict(spec.command_env) if spec.command_env else None


def _run_command(
    args: list[str],
    *,
    env: dict[str, str] | None = None,
    input_text: str | None = None,
) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        args,
        capture_output=True,
        text=True,
        env=env,
        input=input_text,
    )


def _is_non_root_mount(path: Path, *, mount_check_bin: str, env: dict[str, str] | None) -> bool:
    if not path.is_dir():
        return False
    check = _run_command([mount_check_bin, "-T", str(path), "-n", "-o", "TARGET,SOURCE"], env=env)
    if check.returncode != 0:
        return bool(os.path.ismount(path))

    root_check = _run_command([mount_check_bin, "-T", "/", "-n", "-o", "SOURCE"], env=env)
    root_source = (root_check.stdout or "").strip() if root_check.returncode == 0 else ""

    parts = (check.stdout or "").strip().split()
    if len(parts) < 2:
        return bool(os.path.ismount(path))
    target, source = parts[0], parts[1]
    if target == "/":
        return False
    if root_source and source == root_source:
        return False
    return True


def _ensure_valid_identity(value: str, field_name: str) -> str:
    val = (value or "").strip()
    if not val:
        raise StorageProvisioningError(f"Missing required storage identity field: {field_name}")
    return val


def _volume_id_from_root(root: Path) -> str:
    return _sanitize_dns_label(root.name, max_len=32)


def _choose_root_pair(spec: ClusterStorageSpec) -> tuple[Path, Path]:
    env = _command_env(spec)
    mounted_roots: list[Path] = []
    for root_str in spec.synology_roots:
        root = Path(root_str)
        if _is_non_root_mount(root, mount_check_bin=spec.mount_check_bin, env=env):
            mounted_roots.append(root)

    if not mounted_roots:
        allow_unmounted = True
        if env:
            raw = str(env.get("PROLE_STORAGE_ALLOW_UNMOUNTED_ROOTS") or "").strip().lower()
            if raw:
                allow_unmounted = raw in {"1", "true", "yes", "on"}

        if not allow_unmounted:
            raise StorageProvisioningError(
                "No mounted Synology roots available. Expected mounted paths like /synology/d001..d004."
            )

        mounted_roots = [Path(root_str) for root_str in spec.synology_roots]

    ident = f"{spec.namespace}:{spec.cluster_name}:{spec.service_name}"
    digest = int(hashlib.sha256(ident.encode("utf-8")).hexdigest(), 16)
    data_index = digest % len(mounted_roots)
    data_root = mounted_roots[data_index]

    if len(mounted_roots) == 1:
        wal_root = data_root
    else:
        wal_root = mounted_roots[(data_index + 1) % len(mounted_roots)]

    return data_root, wal_root


def build_cluster_storage_paths(spec: ClusterStorageSpec) -> ClusterStoragePaths:
    namespace = _ensure_valid_identity(spec.namespace, "namespace")
    cluster_name = _ensure_valid_identity(spec.cluster_name, "cluster_name")

    data_root, wal_root = _choose_root_pair(spec)
    base_rel = Path(spec.root_namespace) / spec.service_name / namespace / cluster_name
    data_path = data_root / base_rel / "data"
    wal_path = wal_root / base_rel / "wal"

    if _paths_overlap(str(data_path), str(wal_path)):
        raise StorageProvisioningError(
            f"Refusing overlapping storage paths: data='{data_path}', wal='{wal_path}'"
        )

    return ClusterStoragePaths(
        data_root=str(data_root),
        wal_root=str(wal_root),
        data_volume=_volume_id_from_root(data_root),
        wal_volume=_volume_id_from_root(wal_root),
        data_path=str(data_path),
        wal_path=str(wal_path),
    )


def build_pv_name(storage_class_name: str, namespace: str, cluster_name: str, role: str) -> str:
    base = f"{storage_class_name}-{namespace}-{cluster_name}-{role}"
    return _sanitize_dns_label(base)


def build_pv_labels(
    *,
    namespace: str,
    cluster_name: str,
    role: str,
    volume: str,
    service_name: str = "knoe-db",
) -> dict[str, str]:
    return {
        "synology.storage/role": role,
        "synology.storage/volume": volume,
        "knoe.io/namespace": namespace,
        "knoe.io/cluster": cluster_name,
        "knoe.io/service": service_name,
    }


def render_local_pv(
    *,
    pv_name: str,
    host_path: str,
    node_name: str,
    storage_class_name: str,
    capacity: str,
    labels: dict[str, str],
) -> dict[str, Any]:
    return {
        "apiVersion": "v1",
        "kind": "PersistentVolume",
        "metadata": {
            "name": pv_name,
            "labels": labels,
        },
        "spec": {
            "capacity": {"storage": capacity},
            "volumeMode": "Filesystem",
            "accessModes": ["ReadWriteOnce"],
            "storageClassName": storage_class_name,
            "persistentVolumeReclaimPolicy": "Retain",
            "local": {"path": host_path},
            "nodeAffinity": {
                "required": {
                    "nodeSelectorTerms": [
                        {
                            "matchExpressions": [
                                {
                                    "key": "kubernetes.io/hostname",
                                    "operator": "In",
                                    "values": [node_name],
                                }
                            ]
                        }
                    ]
                }
            },
        },
    }


def ensure_host_path(path: str, *, uid: int, gid: int, mode: int = 0o700) -> None:
    target = Path(path)
    if target.exists() and not target.is_dir():
        raise StorageProvisioningError(f"Host path exists but is not a directory: {path}")
    try:
        target.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        if exc.errno == 30:  # Read-only filesystem on local runner
            return
        raise StorageProvisioningError(f"Cannot create host path '{path}': {exc}") from exc

    try:
        os.chmod(target, mode)
    except OSError as exc:
        if exc.errno == 30:
            return
        raise StorageProvisioningError(f"Cannot set mode on host path '{path}': {exc}") from exc
    try:
        os.chown(target, uid, gid)
    except PermissionError:
        # Non-root execution is expected in some local automation contexts.
        return
    except OSError as exc:
        if exc.errno == 30:
            return
        raise StorageProvisioningError(
            f"Cannot set ownership on host path '{path}' to {uid}:{gid}: {exc}"
        ) from exc


def _kubectl_get_json(spec: ClusterStorageSpec, args: list[str]) -> dict[str, Any] | None:
    env = _command_env(spec)
    cmd = [spec.kubectl_bin, *args, "-o", "json"]
    res = _run_command(cmd, env=env)
    if res.returncode != 0:
        return None
    try:
        return json.loads(res.stdout or "{}")
    except json.JSONDecodeError as exc:
        raise StorageProvisioningError(
            f"Failed to parse kubectl JSON output for command: {' '.join(cmd)}"
        ) from exc


def _kubectl_apply_manifest(spec: ClusterStorageSpec, manifest: dict[str, Any]) -> None:
    env = _command_env(spec)
    payload = json.dumps(manifest)
    res = _run_command([spec.kubectl_bin, "apply", "-f", "-"], env=env, input_text=payload)
    if res.returncode != 0:
        stderr = (res.stderr or "").strip()
        raise StorageProvisioningError(f"Failed to apply PV manifest '{manifest['metadata']['name']}': {stderr}")


def _ensure_node_host_path(spec: ClusterStorageSpec, host_path: str) -> None:
    env = _command_env(spec)
    ident = f"{spec.namespace}:{spec.cluster_name}:{spec.node_name}:{host_path}"
    suffix = hashlib.sha1(ident.encode("utf-8")).hexdigest()[:10]
    pod_name = f"knoe-pathprep-{suffix}"
    namespace = "kube-system"

    manifest = {
        "apiVersion": "v1",
        "kind": "Pod",
        "metadata": {"name": pod_name, "namespace": namespace},
        "spec": {
            "restartPolicy": "Never",
            "nodeName": spec.node_name,
            "tolerations": [
                {"key": "node-role.kubernetes.io/control-plane", "operator": "Exists", "effect": "NoSchedule"},
                {"key": "node-role.kubernetes.io/master", "operator": "Exists", "effect": "NoSchedule"},
            ],
            "containers": [
                {
                    "name": "pathprep",
                    "image": "busybox:1.36",
                    "command": ["sh", "-c", "mkdir -p /target && chmod 700 /target || true"],
                    "volumeMounts": [{"name": "target", "mountPath": "/target"}],
                }
            ],
            "volumes": [
                {
                    "name": "target",
                    "hostPath": {"path": host_path, "type": "DirectoryOrCreate"},
                }
            ],
        },
    }

    apply_res = _run_command(
        [spec.kubectl_bin, "apply", "-f", "-"],
        env=env,
        input_text=json.dumps(manifest),
    )
    if apply_res.returncode != 0:
        raise StorageProvisioningError(
            f"Failed to start host path bootstrap pod '{pod_name}': {(apply_res.stderr or apply_res.stdout or '').strip()}"
        )

    try:
        for _ in range(60):
            pod = _kubectl_get_json(spec, ["-n", namespace, "get", "pod", pod_name])
            if pod:
                phase = str((pod.get("status") or {}).get("phase") or "")
                if phase in {"Succeeded", "Running"}:
                    return
                if phase == "Failed":
                    logs = _run_command([spec.kubectl_bin, "-n", namespace, "logs", pod_name], env=env)
                    raise StorageProvisioningError(
                        f"Host path bootstrap pod '{pod_name}' failed for '{host_path}': {(logs.stdout or logs.stderr or '').strip()}"
                    )
            time.sleep(2)
    finally:
        _run_command(
            [spec.kubectl_bin, "-n", namespace, "delete", "pod", pod_name, "--ignore-not-found", "--wait=false"],
            env=env,
        )

    raise StorageProvisioningError(
        f"Timed out waiting for host path bootstrap pod '{pod_name}' to prepare '{host_path}'."
    )


def _collect_existing_pvs(spec: ClusterStorageSpec) -> dict[str, dict[str, Any]]:
    pv_list = _kubectl_get_json(spec, ["get", "pv"])
    if pv_list is None:
        raise StorageProvisioningError("Unable to list existing PersistentVolumes via kubectl.")
    items = pv_list.get("items") or []
    return {str(item.get("metadata", {}).get("name") or ""): item for item in items}


def _extract_node_names(pv_obj: dict[str, Any]) -> set[str]:
    names: set[str] = set()
    terms = (
        pv_obj.get("spec", {})
        .get("nodeAffinity", {})
        .get("required", {})
        .get("nodeSelectorTerms", [])
    )
    for term in terms:
        for expr in term.get("matchExpressions", []) or []:
            if expr.get("key") != "kubernetes.io/hostname":
                continue
            if expr.get("operator") != "In":
                continue
            for val in expr.get("values") or []:
                if val:
                    names.add(str(val))
    return names


def _validate_existing_pv(
    *,
    existing_pv: dict[str, Any],
    expected_name: str,
    expected_path: str,
    expected_storage_class: str,
    expected_node_name: str,
    expected_labels: dict[str, str],
) -> None:
    existing_name = str(existing_pv.get("metadata", {}).get("name") or "")
    if existing_name != expected_name:
        raise StorageProvisioningError(
            f"Internal PV validation error: expected '{expected_name}', got '{existing_name}'"
        )

    current_path = str(existing_pv.get("spec", {}).get("local", {}).get("path") or "")
    if _normalize_posix_path(current_path) != _normalize_posix_path(expected_path):
        raise StorageProvisioningError(
            f"PV '{expected_name}' exists with unexpected path '{current_path}' (expected '{expected_path}')."
        )

    current_sc = str(existing_pv.get("spec", {}).get("storageClassName") or "")
    if current_sc != expected_storage_class:
        raise StorageProvisioningError(
            f"PV '{expected_name}' exists with unexpected storageClassName '{current_sc}' (expected '{expected_storage_class}')."
        )

    current_labels = existing_pv.get("metadata", {}).get("labels") or {}
    for key, expected in expected_labels.items():
        actual = str(current_labels.get(key) or "")
        if actual != expected:
            raise StorageProvisioningError(
                f"PV '{expected_name}' label mismatch for '{key}': got '{actual}', expected '{expected}'."
            )

    node_names = _extract_node_names(existing_pv)
    if expected_node_name not in node_names:
        raise StorageProvisioningError(
            f"PV '{expected_name}' node affinity mismatch: expected node '{expected_node_name}', got {sorted(node_names)}."
        )


def _ensure_no_path_overlap(
    *,
    existing_pvs: dict[str, dict[str, Any]],
    pv_name: str,
    target_path: str,
) -> None:
    for existing_name, existing_pv in existing_pvs.items():
        if existing_name == pv_name:
            continue
        existing_path = str(existing_pv.get("spec", {}).get("local", {}).get("path") or "").strip()
        if not existing_path:
            continue
        if _paths_overlap(target_path, existing_path):
            raise StorageProvisioningError(
                f"PV path overlap detected: target '{target_path}' conflicts with existing PV '{existing_name}' path '{existing_path}'."
            )


def _provision_one_pv(
    *,
    spec: ClusterStorageSpec,
    existing_pvs: dict[str, dict[str, Any]],
    role: str,
    path: str,
    volume: str,
) -> tuple[str, dict[str, str]]:
    pv_name = build_pv_name(spec.storage_class_name, spec.namespace, spec.cluster_name, role)
    labels = build_pv_labels(
        namespace=spec.namespace,
        cluster_name=spec.cluster_name,
        role=role,
        volume=volume,
        service_name=spec.service_name,
    )
    _ensure_no_path_overlap(existing_pvs=existing_pvs, pv_name=pv_name, target_path=path)

    existing = existing_pvs.get(pv_name)
    if existing:
        _validate_existing_pv(
            existing_pv=existing,
            expected_name=pv_name,
            expected_path=path,
            expected_storage_class=spec.storage_class_name,
            expected_node_name=spec.node_name,
            expected_labels=labels,
        )
        return pv_name, labels

    manifest = render_local_pv(
        pv_name=pv_name,
        host_path=path,
        node_name=spec.node_name,
        storage_class_name=spec.storage_class_name,
        capacity=spec.volume_capacity,
        labels=labels,
    )
    _kubectl_apply_manifest(spec, manifest)
    return pv_name, labels


def _validate_pv_conflicts(
    spec: ClusterStorageSpec,
    existing_pvs: dict[str, dict[str, Any]],
    paths: "ClusterStoragePaths",
) -> None:
    """Raise StorageProvisioningError early if existing PVs conflict with planned paths."""
    for role, path in (("data", paths.data_path), ("wal", paths.wal_path)):
        pv_name = build_pv_name(spec.storage_class_name, spec.namespace, spec.cluster_name, role)
        existing = existing_pvs.get(pv_name)
        if existing:
            existing_path = str(
                existing.get("spec", {}).get("local", {}).get("path") or ""
            ).strip()
            if existing_path and existing_path != path:
                raise StorageProvisioningError(
                    f"Existing PV '{pv_name}' has unexpected path '{existing_path}' "
                    f"(expected '{path}')."
                )
        _ensure_no_path_overlap(
            existing_pvs=existing_pvs, pv_name=pv_name, target_path=path
        )


def provision_cluster_storage(spec: ClusterStorageSpec) -> ProvisionedStorage:
    spec = ClusterStorageSpec(
        namespace=_ensure_valid_identity(spec.namespace, "namespace"),
        cluster_name=_ensure_valid_identity(spec.cluster_name, "cluster_name"),
        node_name=_ensure_valid_identity(spec.node_name, "node_name"),
        storage_class_name=_ensure_valid_identity(spec.storage_class_name, "storage_class_name"),
        service_name=_ensure_valid_identity(spec.service_name, "service_name"),
        root_namespace=_ensure_valid_identity(spec.root_namespace, "root_namespace"),
        postgres_uid=spec.postgres_uid,
        postgres_gid=spec.postgres_gid,
        directory_mode=spec.directory_mode,
        volume_capacity=_ensure_valid_identity(spec.volume_capacity, "volume_capacity"),
        synology_roots=spec.synology_roots,
        kubectl_bin=spec.kubectl_bin,
        mount_check_bin=spec.mount_check_bin,
        command_env=spec.command_env,
    )

    paths = build_cluster_storage_paths(spec)
    # Validate PV conflicts before touching the filesystem.
    existing_pvs = _collect_existing_pvs(spec)
    _validate_pv_conflicts(spec, existing_pvs, paths)
    _ensure_node_host_path(spec, paths.data_path)
    _ensure_node_host_path(spec, paths.wal_path)
    ensure_host_path(
        paths.data_path,
        uid=spec.postgres_uid,
        gid=spec.postgres_gid,
        mode=spec.directory_mode,
    )
    ensure_host_path(
        paths.wal_path,
        uid=spec.postgres_uid,
        gid=spec.postgres_gid,
        mode=spec.directory_mode,
    )

    existing_pvs = _collect_existing_pvs(spec)

    data_pv_name, data_labels = _provision_one_pv(
        spec=spec,
        existing_pvs=existing_pvs,
        role="data",
        path=paths.data_path,
        volume=paths.data_volume,
    )
    existing_pvs = _collect_existing_pvs(spec)
    wal_pv_name, wal_labels = _provision_one_pv(
        spec=spec,
        existing_pvs=existing_pvs,
        role="wal",
        path=paths.wal_path,
        volume=paths.wal_volume,
    )

    return ProvisionedStorage(
        data_pv_name=data_pv_name,
        wal_pv_name=wal_pv_name,
        data_path=paths.data_path,
        wal_path=paths.wal_path,
        data_selector=dict(data_labels),
        wal_selector=dict(wal_labels),
    )