# plugins/host_agent/scheduler.py
from __future__ import annotations
import logging
from datetime import datetime, timezone
from typing import Iterable

from sqlalchemy import select

from steward.core.alerts import MetricBatch
from steward.core.scheduler import ScheduledTask
from steward.core.watch import watch
from steward.models.hosts import Host
from .models import HostAgentRegistration

logger = logging.getLogger(__name__)


def arrivals(app) -> dict[str, datetime]:
    """{host id: when its agent's newest post reached Steward}, in memory (#5514).

    last_seen_at is written in the ingest's own transaction, so when Steward
    can't store a post (its storage stalls and the request is dropped) the
    agent's report is lost with it, and an agent posting every few seconds read
    as offline. A post that arrived is the agent reaching Steward; storing it
    is Steward's part. Set before any write, so it holds however slow they are.
    """
    got = getattr(app, "_agent_arrivals", None)
    if got is None:
        got = app._agent_arrivals = {}
    return got


def seconds_since_seen(regs: Iterable, host_names: dict[str, str],
                       now: datetime, watching_since: datetime | None = None,
                       arrived: dict[str, datetime] | None = None
                       ) -> list[tuple[str, float]]:
    """(host name, seconds since its agent last reported) per registration.

    An agent that has never reported counts from when it was registered. A
    registration whose host is gone is skipped. Nothing counts from before
    `watching_since`, when Steward last started watching without a break: an
    agent can't report while Steward is down, so that time is not its silence
    (#5441). An agent that is really gone reads as offline once it has been
    quiet that long while Steward watched. `arrived` (arrivals(app)) counts a
    post Steward received but couldn't store as the agent seen (#5514).
    """
    out = []
    for r in regs:
        name = host_names.get(r.host_id)
        seen = max((t for t in (r.last_seen_at or r.created_at, (arrived or {}).get(r.host_id))
                    if t is not None), default=None)
        if name is None or seen is None:
            continue
        if watching_since is not None:
            seen = max(seen, watching_since)
        out.append((name, max(0.0, (now - seen).total_seconds())))
    return out


async def record_seconds_since_seen(app) -> None:
    """Record host_agent/<host>/seconds_since_seen for every registered agent,
    so "agent offline" is an ordinary alert rule (seconds_since_seen > 180)."""
    now = datetime.now(timezone.utc)
    async with app.db_sessionmaker() as session:
        async with session.begin():
            regs = (await session.execute(select(HostAgentRegistration))).scalars().all()
            watching_since = watch(app, "host_agent_seconds_since_seen").tick(now)
            if not regs:
                return
            host_names = dict((await session.execute(
                select(Host.id, Host.name).where(Host.id.in_([r.host_id for r in regs]))
            )).all())
            batch = MetricBatch()
            for name, seconds in seconds_since_seen(regs, host_names, now, watching_since,
                                                    arrivals(app)):
                batch.add(source_module="host_agent", resource_name=name,
                          metric_name="seconds_since_seen", value=seconds, recorded_at=now)
            await batch.record(session)


def make_task(app) -> ScheduledTask:
    async def _record():
        try:
            await record_seconds_since_seen(app)
        except Exception:
            logger.exception("host_agent: recording seconds_since_seen failed")

    return ScheduledTask(
        name="host_agent_seconds_since_seen",
        coro_factory=_record,
        interval_seconds=60,
        run_on_startup=False,
    )