# plugins/host_agent/rebalance.py
"""A maintenance run spreads swarm services back out after its last host.

A run drains each swarm node in turn, and the swarm never moves a task back
on its own, so every run used to end with the services on whichever node was
drained first. Once the run's last host is back in service, Steward picks
services to move from the fullest node to the emptiest and hands the list to
a swarm manager's helper (maint.py `rebalance`), which force-updates each in
turn: spec unchanged, no image pulled, and the swarm starts it again on its
emptiest node.

A service moves only when the emptier node already holds the exact image the
service is pinned to (its agent reports every image digest it holds), so a
pruned image or a digest the registry has since deleted never strands one.
A service the plan keeps in place, one below its replicas, and one a
constraint places are left where they are, and the run says why.

The choice is model-free (`choose`); the database side is below it, driven by
cadence.advance, the agent's posts (handoffs, results) and the minute tick.
"""
from __future__ import annotations

import json
from datetime import datetime, timedelta

from . import maintenance

# The first agent whose helper can rebalance, and that reports image digests.
REBALANCE_AGENT = (1, 26, 0)
# How many services one rebalance moves at most; the helper refuses more.
SERVICES_MAX = 48
# A node's list of images older than this isn't trusted (agents report it
# every 10 minutes).
DIGESTS_FRESH = timedelta(minutes=30)
# The helper starts no move after 40 min and gives each 5 to settle.
DEADLINE = timedelta(minutes=60)
STATES = ("handed", "done", "failed", "timed_out", "nothing")
WINDOW_NAME = "Maintenance: spreading swarm services back out"


def keep_list(text: str | None) -> list[str]:
    """A plan's keep_in_place as names, in the order written."""
    out = []
    for name in (text or "").replace("\n", ",").split(","):
        name = name.strip()
        if name and name not in out:
            out.append(name)
    return out


def _repo(image: str) -> str:
    """An image reference's repository, without tag or digest, in the short
    form Docker reports (no docker.io/ or library/)."""
    ref = image.split("@", 1)[0]
    slash = ref.rfind("/")
    colon = ref.rfind(":")
    if colon > slash:
        ref = ref[:colon]
    for prefix in ("docker.io/", "index.docker.io/"):
        if ref.startswith(prefix):
            ref = ref[len(prefix):]
    if ref.startswith("library/"):
        ref = ref[len("library/"):]
    return ref


def pinned(image: str | None, digest: str | None) -> str | None:
    """"repo@sha256:…" for a service's image, or None when it isn't pinned."""
    if not image or not digest or not digest.startswith("sha256:"):
        return None
    return f"{_repo(image)}@{digest}"


def held(digests) -> set:
    """A node's reported digests, in the same form as pinned()."""
    return {f"{_repo(d)}@{d.split('@', 1)[1]}" for d in digests or []
            if isinstance(d, str) and "@sha256:" in d}


def choose(services: list, nodes: dict, images: dict, keep=(),
           limit: int = SERVICES_MAX) -> tuple[list, list]:
    """(services to move, [{"service", "reason"}] left where they are).

    services: {"name", "mode", "desired", "running", "image" ("repo@sha256:…"
    or None), "placement": [{"node_id", "running"}]}; nodes: {node id: name}
    for the nodes to balance across (ready and active); images: {node id: set
    of "repo@sha256:…"}, a node missing when its list is unknown.

    While the fullest node has two or more tasks than the emptiest, a service
    whose tasks all run on the fullest moves; one with several replicas there
    first, as one node holding them all defeats the replicas. A move counts
    as its share of the replicas going to the emptiest. What is reported left
    is what kept the nodes from evening out: the services that would have
    moved next but can't."""
    if len(nodes) < 2:
        return [], []
    counts = {n: 0 for n in nodes}
    for s in services:
        for p in s.get("placement") or []:
            if p.get("node_id") in counts:
                counts[p["node_id"]] += int(p.get("running") or 0)
    keep = set(keep)
    order = sorted(services, key=lambda s: (-int(s.get("running") or 0), s["name"]))
    moves: list = []
    left: dict = {}
    while len(moves) < limit:
        src = max(counts, key=lambda n: (counts[n], n))
        dst = min(counts, key=lambda n: (counts[n], n))
        gap = counts[src] - counts[dst]
        if gap < 2:
            break
        pick, blocked = None, {}
        for s in order:
            name, running = s["name"], int(s.get("running") or 0)
            on = {p.get("node_id"): int(p.get("running") or 0) for p in s.get("placement") or []
                  if int(p.get("running") or 0) > 0}
            if name in moves or s.get("mode") != "replicated" or set(on) != {src}:
                continue
            share = max(1, running // len(nodes))
            if 2 * share > gap:
                continue                    # it would only tip the balance the other way
            if name in keep:
                blocked[name] = "kept in place by the plan"
            elif running < int(s.get("desired") or 0):
                blocked[name] = "below its replicas"
            elif not s.get("image"):
                blocked[name] = "its image isn't pinned to a digest"
            elif dst not in images:
                blocked[name] = f"no image list from {nodes[dst]} (its agent needs 1.26.0)"
            elif s["image"] not in images[dst]:
                blocked[name] = f"its image isn't on {nodes[dst]}"
            else:
                pick = s
                break
        if pick is None:
            left = blocked                  # what kept the nodes from evening out
            break
        share = max(1, int(pick.get("running") or 0) // len(nodes))
        moves.append(pick["name"])
        counts[src] -= share
        counts[dst] += share
    return moves, [{"service": n, "reason": r} for n, r in sorted(left.items())]


def can_rebalance(agent_version) -> bool:
    try:
        return tuple(int(x) for x in str(agent_version).split(".")[:3]) >= REBALANCE_AGENT
    except ValueError:
        return False


def step_id(run) -> str:
    return f"{run.id}-rebalance"


def _set(run, **changes) -> None:
    """Change the run's rebalance record (a new dict, so the JSON column
    sees the change)."""
    run.rebalance = {**(run.rebalance or {}), **changes}


def handoff(run) -> dict | None:
    """The step to hand the manager's agent, until its result arrives."""
    rb = run.rebalance or {}
    if rb.get("state") != "handed" or rb.get("seen"):
        return None
    return {"id": rb["step_id"], "action": "rebalance", "services": rb.get("planned") or []}


def apply_result(run, event: dict, now: datetime) -> bool:
    """A `maintenance` event for the run's rebalance; True when it changed."""
    rb = run.rebalance or {}
    if rb.get("state") != "handed" or event.get("id") != rb.get("step_id"):
        return False
    state = event.get("state")
    if state == "running":
        if rb.get("seen"):
            return False
        _set(run, seen=True)
        return True
    if state not in ("done", "failed"):
        return False
    moved = [str(m)[:64] for m in event.get("moved") or [] if isinstance(m, str)][:SERVICES_MAX]
    gone = [{"service": str(x.get("service"))[:64], "reason": str(x.get("reason"))[:160]}
            for x in event.get("left") or [] if isinstance(x, dict)][:SERVICES_MAX]
    ok = state == "done" and event.get("ok") is not False
    _set(run, state="done" if ok else "failed", moved=moved, left=(rb.get("left") or []) + gone,
         detail=str(event.get("detail") or "")[:maintenance.DETAIL_MAX] or None,
         finished_at=now.isoformat())
    return True


def check_deadline(run, now: datetime) -> bool:
    rb = run.rebalance or {}
    if rb.get("state") != "handed" or now - datetime.fromisoformat(rb["at"]) <= DEADLINE:
        return False
    _set(run, state="timed_out", finished_at=now.isoformat(),
         detail=f"the manager didn't finish within {int(DEADLINE.total_seconds() // 60)} min")
    return True


def outcome(run) -> tuple[str, str | None]:
    """How a run whose hosts are all done ends: ("done" | "failed", detail)."""
    rb = run.rebalance or {}
    state = rb.get("state")
    left = rb.get("left") or []
    tail = f"; {len(left)} left where they are" if left else ""
    if state == "done":
        n = len(rb.get("moved") or [])
        return "done", f"spread {n} swarm service{'s' if n != 1 else ''} back out{tail}"
    if state in ("failed", "timed_out"):
        return "failed", f"spreading swarm services back out: {rb.get('detail') or state}"
    if state == "nothing" and rb.get("detail"):
        return "done", f"swarm services not spread back out: {rb['detail']}"
    if state == "nothing" and left:
        return "done", f"no swarm service could be spread back out{tail}"
    return "done", None


def label(run) -> tuple[str, bool, bool] | None:
    """(text, bad, busy) for the run's swarm row on the page, or None when
    the run has no rebalance."""
    rb = run.rebalance or {}
    state = rb.get("state")
    if state is None:
        return None
    if state == "handed":
        n = len(rb.get("planned") or [])
        return f"Spreading {n} service{'s' if n != 1 else ''} back out", False, True
    text = outcome(run)[1] or "Nothing to spread back out"
    return text[:1].upper() + text[1:], state in ("failed", "timed_out"), False


# ── applied to the database ──────────────────────────────────────────────────

async def _freshest(session, model, key: str) -> list:
    """The newest row per `key` (several managers may report the swarm)."""
    from sqlalchemy import select
    rows: dict = {}
    for r in (await session.execute(select(model))).scalars():
        k = getattr(r, key)
        if k not in rows or r.updated_at > rows[k].updated_at:
            rows[k] = r
    return list(rows.values())


async def _swarm_view(session, now: datetime) -> tuple[list, dict, dict] | None:
    """(services, nodes, images) as choose() takes them, from what the Docker
    plugin last stored; None without it."""
    from sqlalchemy import select
    from steward.models.hosts import Host
    models = maintenance.docker_models()
    if models is None:
        return None
    services = [{"name": s.service_name, "mode": s.mode, "desired": s.desired,
                 "running": s.running, "image": pinned(s.image, s.image_digest),
                 "placement": json.loads(s.placement_json or "[]")}
                for s in await _freshest(session, models.DockerSwarmService, "service_name")]
    nodes = {n.node_id: n.hostname or n.node_id
             for n in await _freshest(session, models.DockerSwarmNode, "node_id")
             if n.status == "ready" and n.availability == "active"}
    hosts = dict((await session.execute(select(Host.name, Host.id))).all())
    by_name = {name.lower(): hid for name, hid in hosts.items()}
    images = {}
    for node_id, name in nodes.items():
        host_id = by_name.get(name.lower())
        disk = await session.get(models.DockerDiskUsage, host_id) if host_id else None
        if disk is not None and now - disk.updated_at <= DIGESTS_FRESH:
            got = held(json.loads(disk.digests_json or "[]"))
            if got:
                images[node_id] = got
    return services, nodes, images


async def _manager(session, jobs: list) -> tuple[str | None, str | None]:
    """(host id, None) of a swarm manager whose helper can rebalance, the one
    that drained the run's nodes first; (None, why) when there is none."""
    from sqlalchemy import select
    from .models import HostAgentRegistration
    _, managers, _ = await maintenance._swarm(session)
    regs = {r.host_id: r for r in (await session.execute(select(HostAgentRegistration).where(
        HostAgentRegistration.host_id.in_(list(managers)),
        HostAgentRegistration.maint_helper.is_(True)))).scalars()} if managers else {}
    drained_by = [j.drain_host_id for j in jobs if j.drain_host_id]
    for host_id in drained_by + sorted(regs):
        reg = regs.get(host_id)
        if reg is not None and can_rebalance(reg.agent_version):
            return host_id, None
    if regs:
        return None, "no swarm manager's agent can do it yet: run Update agent on a manager"
    return None, "no swarm manager has the maintenance helper"


async def begin(session, run, plan, now: datetime) -> bool:
    """The run's last host is back in service: hand a manager the services to
    move, or record why there are none. True while a rebalance is handed out."""
    from sqlalchemy import select
    from .models import HostMaintenanceJob
    if not plan.spread or run.rebalance is not None:
        return False
    jobs = (await session.execute(select(HostMaintenanceJob).where(
        HostMaintenanceJob.run_id == run.id))).scalars().all()
    if not any(maintenance.drains(j) for j in jobs):
        return False                        # nothing was drained: nothing to put back
    view = await _swarm_view(session, now)
    if view is None:
        return False
    services, nodes, images = view
    moves, left = choose(services, nodes, images, keep_list(plan.keep_in_place))
    if not moves:
        run.rebalance = {"state": "nothing", "left": left, "at": now.isoformat()}
        return False
    host_id, why = await _manager(session, jobs)
    if host_id is None:
        run.rebalance = {"state": "nothing", "left": left, "detail": why, "at": now.isoformat()}
        return False
    run.rebalance = {"state": "handed", "step_id": step_id(run), "host_id": host_id,
                     "at": now.isoformat(), "planned": moves, "left": left, "seen": False,
                     "window_id": await _window(session, plan, now)}
    return True


async def _window(session, plan, now: datetime) -> str | None:
    """A window over the Docker figures while services restart, so their
    restarts raise no alert; ended when the rebalance does."""
    from steward.models.alerts import MaintenanceWindow
    if plan.created_by is None:
        return None
    w = MaintenanceWindow(name=WINDOW_NAME, start_at=now, end_at=now + DEADLINE,
                          scope_source="docker", scope_resource=None, created_by=plan.created_by)
    session.add(w)
    await session.flush()
    return w.id


async def finish(session, run, now: datetime) -> None:
    """The rebalance ended: its window ends, and it goes on the timeline."""
    from steward.models.alerts import MaintenanceWindow
    from . import timeline
    rb = run.rebalance or {}
    if rb.get("window_id"):
        w = await session.get(MaintenanceWindow, rb["window_id"])
        if w is not None:
            w.end_at = max(w.start_at, now)
    if rb.get("state") in ("done", "failed", "timed_out"):
        await timeline.write(session, f"maintenance run {run.id}", timeline.run_entry(session, run))


async def running_runs(session) -> list:
    from sqlalchemy import select
    from .models import HostMaintenanceRun
    return (await session.execute(select(HostMaintenanceRun).where(
        HostMaintenanceRun.state == "running"))).scalars().all()


async def handoffs(session, host_id: str) -> list[dict]:
    """The rebalance steps this host's agent should carry out now."""
    out = []
    for run in await running_runs(session):
        if (run.rebalance or {}).get("host_id") == host_id:
            step = handoff(run)
            if step is not None:
                out.append(step)
    return out


async def apply_event(session, event: dict, now: datetime) -> bool:
    """Apply a `maintenance` event to the running run whose rebalance it is."""
    for run in await running_runs(session):
        if apply_result(run, event, now):
            if (run.rebalance or {}).get("state") != "handed":
                await finish(session, run, now)
            return True
    return False