prole/knoe/core/cnpg_placement.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Iterable


CNPG_PLACEMENT_SCHEMA_VERSION = "v1"


def _coerce_positive_int(raw: object, default: int) -> int:
    try:
        val = int(str(raw).strip())
        if val < 1:
            return default
        return val
    except Exception:
        return default


def _normalize_nodes(nodes: Iterable[str] | None) -> list[str]:
    normalized = {
        str(node or "").strip()
        for node in (nodes or [])
        if str(node or "").strip()
    }
    return sorted(normalized)


def _normalize_assignments(raw: object) -> dict[int, str]:
    if not isinstance(raw, dict):
        return {}
    out: dict[int, str] = {}
    for k, v in raw.items():
        try:
            ordinal = int(str(k).strip())
        except Exception:
            continue
        if ordinal < 0:
            continue
        node = str(v or "").strip()
        if not node:
            continue
        out[ordinal] = node
    return out


def _assign_round_robin(
    *, desired_instances: int, eligible_nodes: list[str], base: dict[int, str] | None = None
) -> dict[int, str]:
    assignments = dict(base or {})
    if desired_instances < 1 or not eligible_nodes:
        return {}

    for ordinal in range(desired_instances):
        if ordinal in assignments:
            continue
        assignments[ordinal] = eligible_nodes[ordinal % len(eligible_nodes)]

    return {k: assignments[k] for k in sorted(assignments.keys()) if k < desired_instances}


def _plan_hash(payload: dict) -> str:
    canonical = json.dumps(payload, separators=(",", ":"), sort_keys=True)
    digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
    return digest[:16]


def _normalize_plan_for_hash(
    *,
    cluster_name: str,
    desired_instances: int,
    eligible_nodes: list[str],
    assignments: dict[int, str],
) -> dict:
    return {
        "schema_version": CNPG_PLACEMENT_SCHEMA_VERSION,
        "cluster_name": cluster_name,
        "desired_instances": desired_instances,
        "eligible_nodes": eligible_nodes,
        "assignments": {str(k): assignments[k] for k in sorted(assignments.keys())},
    }


def _prior_plan_matches_cluster(prior_plan: dict | None, cluster_name: str) -> bool:
    if not isinstance(prior_plan, dict):
        return False
    return str(prior_plan.get("cluster_name") or "").strip() == cluster_name


def _prior_plan_is_eligible(prior_assignments: dict[int, str], eligible_nodes: list[str]) -> bool:
    if not prior_assignments:
        return False
    if not eligible_nodes:
        return False
    eligible_set = set(eligible_nodes)
    return all(node in eligible_set for node in prior_assignments.values())


def plan_cnpg_placement(
    *,
    cluster_name: str,
    desired_instances: int,
    candidate_nodes: Iterable[str] | None,
    prior_plan: dict | None = None,
    rebalance: bool = False,
) -> dict:
    cluster = str(cluster_name or "").strip() or "knoe-db"
    desired = _coerce_positive_int(desired_instances, 1)
    eligible_nodes = _normalize_nodes(candidate_nodes)

    prior_assignments = _normalize_assignments(
        (prior_plan or {}).get("assignments") if isinstance(prior_plan, dict) else {}
    )
    prior_matches = _prior_plan_matches_cluster(prior_plan, cluster)
    prior_eligible = _prior_plan_is_eligible(prior_assignments, eligible_nodes)

    reused = False
    regenerated = True
    reason = "new_cluster"

    if rebalance:
        assignments = _assign_round_robin(
            desired_instances=desired,
            eligible_nodes=eligible_nodes,
        )
        reason = "rebalance_requested"
    elif prior_matches and prior_eligible:
        # Reuse existing ordinal bindings and only extend/trim as needed.
        assignments = _assign_round_robin(
            desired_instances=desired,
            eligible_nodes=eligible_nodes,
            base=prior_assignments,
        )
        reused = True
        regenerated = False
        reason = "reused"
    elif prior_matches and prior_assignments and not prior_eligible:
        assignments = _assign_round_robin(
            desired_instances=desired,
            eligible_nodes=eligible_nodes,
        )
        reason = "assigned_node_no_longer_eligible"
    else:
        assignments = _assign_round_robin(
            desired_instances=desired,
            eligible_nodes=eligible_nodes,
        )

    payload = _normalize_plan_for_hash(
        cluster_name=cluster,
        desired_instances=desired,
        eligible_nodes=eligible_nodes,
        assignments=assignments,
    )
    plan_hash = _plan_hash(payload)

    return {
        **payload,
        "plan_id": f"cnpg-placement-{plan_hash}",
        "plan_hash": plan_hash,
        "metadata": {
            "reused": reused,
            "regenerated": regenerated,
            "reason": reason,
            "prior_plan_present": isinstance(prior_plan, dict),
        },
    }


def load_cnpg_placement_plan(path: Path) -> dict | None:
    try:
        if not path.exists():
            return None
        payload = json.loads(path.read_text())
    except Exception:
        return None
    if not isinstance(payload, dict):
        return None
    return payload


def save_cnpg_placement_plan(path: Path, plan: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(plan, indent=2, sort_keys=True) + "\n")