# plugins/host_agent/maintenance.py
"""A host's maintenance job, followed to its end through any restart (#5767).

Steward plans and the host's agent acts: a job is one host's update and, only
when the update says one is needed, its reboot, with the host's swarm node
drained first and made active again after. Each action is one step the agent
hands to its root helper (maint.py, #5766). Everything lives on the job's row,
so Steward restarting, or the host Steward runs on rebooting, loses nothing:
the next post from the agent carries the step's result, and the `started`
event a host sends after booting is what completes a reboot.

The transitions here are model-free (a job is anything with the attributes),
so they're unit-tested without a database; routes.ingest and the scheduler's
tick apply them.

States: queued → updating → [draining →] rebooting → [undraining →] done,
or failed / timed_out. A job with no reboot needed goes updating → done.
"""
from __future__ import annotations

from datetime import datetime, timedelta, timezone

ACTIVE = ("queued", "updating", "draining", "rebooting", "undraining")
FINISHED = ("done", "failed", "timed_out")
# The helper action each handed state asks for.
STEP_ACTION = {"updating": "update", "draining": "drain", "rebooting": "reboot",
               "undraining": "undrain"}
# How long a state may last before the job is given up. Queued counts from the
# job's creation (the host may be silent); the rest from their hand-off.
DEADLINES = {
    "queued": timedelta(hours=1),
    "updating": timedelta(minutes=90),     # the helper itself stops at 60
    "draining": timedelta(minutes=10),
    "rebooting": timedelta(minutes=30),
    "undraining": timedelta(minutes=10),
}
# A host's clock and Steward's differ by seconds; a boot this far before the
# reboot was handed out is still taken as that reboot.
BOOT_SLACK = timedelta(seconds=60)
DETAIL_MAX = 500


def step_id(job, state: str) -> str:
    """The helper request id for a job's step: unique per job and step, and in
    the helper's id alphabet (letters, digits, dashes)."""
    return f"{job.id}-{STEP_ACTION[state]}"


def _enter(job, state: str, now: datetime) -> None:
    job.state, job.state_at = state, now
    job.step_id = step_id(job, state) if state in STEP_ACTION else None
    job.step_seen = False               # a new step is handed out until its result arrives
    if state in FINISHED:
        job.finished_at = now


def _finish(job, state: str, now: datetime, detail: str | None = None) -> None:
    if detail:
        job.detail = detail[:DETAIL_MAX]
    _enter(job, state, now)


def start(job, now: datetime) -> None:
    """The host's agent is about to be handed the job: update first."""
    if job.state == "queued":
        _enter(job, "updating", now)


def executor(job) -> str:
    """The host whose agent carries out the current step: the job's own host,
    except a drain or undrain, which a swarm manager does."""
    if job.state in ("draining", "undraining") and job.drain_host_id:
        return job.drain_host_id
    return job.host_id


def handoff(job) -> dict | None:
    """What to hand the executing agent now, or None. Handed on every post
    until a result for the step arrives; the agent and the helper both ignore
    an id they've already taken, so a repeat is harmless."""
    if job.state not in STEP_ACTION or job.step_seen:
        return None
    out = {"id": job.step_id, "action": STEP_ACTION[job.state]}
    if job.state in ("draining", "undraining") and job.node_id:
        out["node"] = job.node_id
    return out


REBOOT_POLICIES = ("if_required", "always", "never")


def wants_reboot(job) -> bool:
    """Reboot only when the update says one is needed (the operator's
    default), always (a plan may ask), or never."""
    if job.reboot_policy == "always":
        return True
    return bool(job.reboot_needed) and job.reboot_policy == "if_required"


def _after_update(job, now: datetime) -> None:
    if not wants_reboot(job):
        _enter(job, "done", now)
    elif job.node_id and job.drain_host_id:
        _enter(job, "draining", now)
    else:
        _enter(job, "rebooting", now)


def apply_result(job, event: dict, now: datetime) -> bool:
    """Apply a `maintenance` event for this job's current step; True when the
    job changed. A result for a step the job has moved past is ignored, so a
    repeated or late report does nothing."""
    if job.state not in STEP_ACTION or event.get("id") != job.step_id:
        return False
    state = event.get("state")
    if state in ("running", "rebooting"):
        if job.step_seen:
            return False
        job.step_seen = True
        return True
    if state not in ("done", "failed"):
        return False
    if state == "failed" or event.get("ok") is False:
        what = STEP_ACTION[job.state]
        detail = f"{what} failed: {event.get('detail') or 'no detail'}"
        if job.state == "undraining":
            detail += " (the node is still drained)"
        _finish(job, "failed", now, detail)
        return True
    if job.state == "updating":
        if isinstance(event.get("packages"), int):
            job.packages = event["packages"]
        job.reboot_needed = event.get("reboot_required") is True
        _after_update(job, now)
    elif job.state == "draining":
        _enter(job, "rebooting", now)
    elif job.state == "undraining":
        _enter(job, "done", now)
    else:                               # rebooting: the boot completes it, not this
        return False
    return True


def apply_boot(job, boot_at: datetime | None, now: datetime) -> bool:
    """A `started` event from the job's host. A boot after the reboot was
    handed out is that reboot: the host is back."""
    if job.state != "rebooting" or boot_at is None or job.state_at is None:
        return False
    if boot_at < job.state_at - BOOT_SLACK:
        return False                    # the agent restarted; the host didn't
    job.back_at = now
    if job.node_id and job.drain_host_id:
        _enter(job, "undraining", now)
    else:
        _enter(job, "done", now)
    return True


def check_deadline(job, now: datetime) -> bool:
    """Give up on a job whose state has lasted past its deadline."""
    limit = DEADLINES.get(job.state)
    since = job.state_at or job.created_at
    if limit is None or since is None or now - since <= limit:
        return False
    waiting = {"queued": "the host's agent didn't check in",
               "updating": "the update didn't finish",
               "draining": "the node wasn't drained",
               "rebooting": "the host didn't come back",
               "undraining": "the node wasn't made active again"}[job.state]
    _finish(job, "timed_out", now, f"{waiting} within {int(limit.total_seconds() // 60)} min")
    return True


def window_end(job, now: datetime) -> datetime:
    """When the job's maintenance window should end: the longest the rest of
    the job could take, or now once it has finished."""
    if job.state in FINISHED:
        return now
    left = sum((DEADLINES[s] for s in ACTIVE[ACTIVE.index(job.state):]), timedelta())
    return now + left


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


# ── applied to the database: ingest, the tick, and starting a job ────────────

async def _window(session, job, host_name: str, now: datetime) -> None:
    """Keep the job's maintenance window over the host until the job ends: one
    scoped to this host's agent metrics, so an alert the job causes waits."""
    from steward.models.alerts import MaintenanceWindow
    if job.window_id:
        w = await session.get(MaintenanceWindow, job.window_id)
        if w is not None:
            w.end_at = max(w.start_at, window_end(job, now))
        return
    if job.created_by is None or job.state in FINISHED:
        return                          # a window needs an owner; none for a finished job
    w = MaintenanceWindow(name=f"Maintenance: {host_name}"[:128], start_at=now,
                          end_at=window_end(job, now), scope_source="host_agent",
                          scope_resource=host_name, created_by=job.created_by)
    session.add(w)
    await session.flush()
    job.window_id = w.id


async def _host_name(session, host_id: str) -> str:
    from steward.models.hosts import Host
    host = await session.get(Host, host_id)
    return host.name if host is not None else host_id


async def apply_events(session, host_id: str, events, now: datetime) -> int:
    """Apply one post's `maintenance` results and `started` boots to the jobs
    they belong to; returns how many jobs changed. A result is matched by its
    step id, so a drain reported by a manager reaches the job of the node it
    drained."""
    from sqlalchemy import select
    from .downtime import _at
    from .models import HostMaintenanceJob
    if not isinstance(events, list):
        return 0
    changed = set()
    for event in events:
        if not isinstance(event, dict):
            continue
        if event.get("kind") == "maintenance" and isinstance(event.get("id"), str):
            job = (await session.execute(select(HostMaintenanceJob).where(
                HostMaintenanceJob.step_id == event["id"],
                HostMaintenanceJob.state.in_(ACTIVE)))).scalar_one_or_none()
            if job is not None and apply_result(job, event, now):
                changed.add(job)
        elif event.get("kind") == "started":
            boot_at = _at(event.get("boot_at"))
            for job in (await session.execute(select(HostMaintenanceJob).where(
                    HostMaintenanceJob.host_id == host_id,
                    HostMaintenanceJob.state == "rebooting"))).scalars():
                if apply_boot(job, boot_at, now):
                    changed.add(job)
    for job in changed:
        await _window(session, job, await _host_name(session, job.host_id), now)
    return len(changed)


async def handoffs(session, host_id: str, host_name: str, now: datetime) -> list[dict]:
    """The steps this host's agent should carry out now. A queued job for the
    host starts here, when its agent is known to be listening."""
    from sqlalchemy import or_, select
    from .models import HostMaintenanceJob
    jobs = (await session.execute(select(HostMaintenanceJob).where(
        HostMaintenanceJob.state.in_(ACTIVE),
        or_(HostMaintenanceJob.host_id == host_id,
            HostMaintenanceJob.drain_host_id == host_id)
    ).order_by(HostMaintenanceJob.created_at))).scalars().all()
    out = []
    for job in jobs:
        if job.state == "queued" and job.host_id == host_id:
            start(job, now)
            await _window(session, job, host_name, now)
        if executor(job) == host_id:
            step = handoff(job)
            if step is not None:
                out.append(step)
    return out


async def tick(session, now: datetime) -> int:
    """Give up on jobs past their deadline; returns how many."""
    from sqlalchemy import select
    from .models import HostMaintenanceJob
    gave_up = 0
    for job in (await session.execute(select(HostMaintenanceJob).where(
            HostMaintenanceJob.state.in_(ACTIVE)))).scalars():
        if check_deadline(job, now):
            gave_up += 1
            await _window(session, job, await _host_name(session, job.host_id), now)
    return gave_up


class CannotStart(Exception):
    """Why a host can't take a maintenance job now, in words for the page."""


def pick_drainer(node_hostname_to_id: dict, managers: dict, host_name: str,
                 host_id: str, helpers: set, availability: str = "active") -> tuple:
    """(node_id, drain_host_id) for a host in a swarm, (None, None) for one
    that isn't. `managers` maps a manager's host id to its swarm node id;
    `helpers` is the hosts whose agent has the maintenance helper. The host
    drains itself when it is a manager, else another manager does. A node the
    operator already set to drain or pause (a manager kept free of work) is
    left as it is: no drain before, and no making it active after."""
    node_id = node_hostname_to_id.get(host_name.lower())
    if node_id is None:
        return None, None
    if availability != "active":
        return node_id, None
    if host_id in managers and host_id in helpers:
        return node_id, host_id
    for mgr in sorted(managers):
        if mgr != host_id and mgr in helpers:
            return node_id, mgr
    raise CannotStart("it is a swarm node and no swarm manager has the maintenance "
                      "helper to drain it: run Update agent on a manager")


async def plan_job(session, host, reg, reboot_policy: str = "if_required") -> tuple:
    """(node_id, drain_host_id) for a job on `host`, or CannotStart saying why
    it can't have one now. The panel asks this before offering the button."""
    from sqlalchemy import select
    from .models import HostAgentRegistration, HostMaintenanceJob
    if reg is None or not reg.maint_helper:
        raise CannotStart("its agent has no maintenance helper: run Update agent first")
    if reboot_policy not in REBOOT_POLICIES:
        raise CannotStart(f"unknown reboot policy {reboot_policy!r}")
    busy = (await session.execute(select(HostMaintenanceJob.id).where(
        HostMaintenanceJob.host_id == host.id,
        HostMaintenanceJob.state.in_(ACTIVE)))).first()
    if busy is not None:
        raise CannotStart("a maintenance job is already running on it")
    if reboot_policy == "never":
        return None, None
    nodes, managers, availability = await _swarm(session)
    helpers = set((await session.execute(select(HostAgentRegistration.host_id).where(
        HostAgentRegistration.maint_helper.is_(True)))).scalars())
    node_id = nodes.get(host.name.lower())
    return pick_drainer(nodes, managers, host.name, host.id, helpers,
                        availability.get(node_id, "active"))


async def create_job(session, host, reg, user_id: str | None,
                     reboot_policy: str = "if_required", now: datetime | None = None,
                     run_id: str | None = None):
    """Queue a maintenance job for `host`; it starts with the agent's next post."""
    from .models import HostMaintenanceJob
    now = now or utcnow()
    node_id, drain_host_id = await plan_job(session, host, reg, reboot_policy)
    job = HostMaintenanceJob(host_id=host.id, state="queued", state_at=now,
                             reboot_policy=reboot_policy, node_id=node_id,
                             drain_host_id=drain_host_id, created_by=user_id, created_at=now,
                             run_id=run_id)
    session.add(job)
    await session.flush()
    return job


async def latest_job(session, host_id: str, now: datetime, days: int = 366):
    """The host's running job, else its last one (within a year), for the panel."""
    from sqlalchemy import select
    from .models import HostMaintenanceJob
    return (await session.execute(select(HostMaintenanceJob).where(
        HostMaintenanceJob.host_id == host_id,
        HostMaintenanceJob.created_at >= now - timedelta(days=days),
    ).order_by(HostMaintenanceJob.state.in_(ACTIVE).desc(),
               HostMaintenanceJob.created_at.desc()).limit(1))).scalar_one_or_none()


LABELS = {
    "queued": "Waiting for the agent",
    "updating": "Updating",
    "draining": "Draining its swarm node",
    "rebooting": "Rebooting",
    "undraining": "Making its swarm node active again",
    "done": "Done",
    "failed": "Failed",
    "timed_out": "Timed out",
}


def summary(job) -> str:
    """One line on what the job did, for the panel."""
    if job.state not in FINISHED:
        return LABELS[job.state]
    parts = []
    if job.packages is not None:
        parts.append(f"{job.packages} package{'s' if job.packages != 1 else ''} updated")
    if job.back_at is not None:
        parts.append("rebooted")
    elif job.reboot_needed and job.reboot_policy == "never":
        parts.append("reboot needed, not allowed by this job")
    elif job.state == "done" and job.reboot_needed is False:
        parts.append("no reboot needed")
    head = LABELS[job.state]
    return f"{head}: {', '.join(parts)}" if parts else head


def docker_models():
    """The loaded Docker plugin's models module, or None when it isn't loaded."""
    from steward.core.plugin_manager import loaded_plugin_submodule
    return loaded_plugin_submodule("docker", "models")


async def _swarm(session) -> tuple:
    """({node hostname: node id}, {manager host id: its node id},
    {node id: availability}) from what the Docker plugin last stored; empty
    when it isn't installed. A manager is a host whose agent reported the
    swarm's nodes."""
    models = docker_models()
    if models is None:
        return {}, {}, {}
    DockerSwarmNode = models.DockerSwarmNode
    from sqlalchemy import select
    from steward.models.hosts import Host
    rows = (await session.execute(select(DockerSwarmNode))).scalars().all()
    nodes = {(n.hostname or "").lower(): n.node_id for n in rows if n.hostname}
    reporters = {n.host_id for n in rows}
    names = dict((await session.execute(select(Host.id, Host.name).where(
        Host.id.in_(reporters)))).all()) if reporters else {}
    managers = {hid: nodes.get(name.lower()) for hid, name in names.items()}
    availability = {n.node_id: (n.availability or "active")
                    for n in sorted(rows, key=lambda n: n.updated_at)}
    return nodes, managers, availability