# plugins/docker/ingest.py
"""Persist host-scoped Docker samples pushed by the host agent.

Published as the "docker.persist_host_samples" capability (see __init__.setup),
so the host_agent plugin can hand off a sample's `docker` array (and, on a swarm
manager, its `swarm` object) WITHOUT importing the docker models — the coupling
is opportunistic and degrades to a no-op when the docker plugin is disabled.
Runs inside the caller's open transaction (the ingest handler's SAVEPOINT);
never opens or commits its own.
"""
from __future__ import annotations

import json
from datetime import datetime

from sqlalchemy import delete, select


# Stopped/terminal container states (anything not "running"/"paused"/"restarting").
_STOPPED_STATES = {"exited", "dead", "stopped"}


def _parse_started_at(value) -> datetime | None:
    """Parse the agent's ISO-8601 started_at string back to a datetime."""
    if not isinstance(value, str) or not value:
        return None
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return None


def _derive_events(old_state: dict, new_containers: list) -> list:
    """Diff a fresh snapshot against stored per-container state → lifecycle events.

    Pure function (no DB) so it's unit-testable. `old_state` maps name → the
    previously stored {status, health, oom_killed, exit_code, restart_count}; `new_containers`
    is the newest snapshot's list of container dicts. Returns a list of
    (container_name, event, detail) tuples:

      start         — a container transitions into running (or a genuinely new
                      container appears already running)
      die           — also for a new container first seen already exited
                      non-zero: a task that crashed between two samples (#5049)
      restart       — Docker restarted it (its restart count went up); said
                      instead of a start, since that is what happened
      stop          — running → not-running, or a running container is removed
      die           — same as stop but with a non-zero exit code (abnormal)
      oom           — OOMKilled flips False→True
      health_change — HEALTHCHECK status string changes

    The CALLER skips this entirely on a host's first-ever snapshot (empty
    old_state) so the baseline doesn't emit a spurious "start" per existing
    container — a later-appearing container still gets one because by then
    old_state is populated and that container's prior entry is simply absent.
    """
    events: list = []
    new_by_name = {c["name"]: c for c in new_containers if c.get("name")}

    for name, c in new_by_name.items():
        new_status = (c.get("status") or "").lower()
        new_running = new_status == "running"
        new_oom = bool(c.get("oom_killed"))
        new_health = c.get("health")
        new_exit = c.get("exit_code")
        old = old_state.get(name)

        if old is None:
            # Newly observed container: a start, or, when it is already down
            # with a non-zero exit, the death that happened between samples.
            if new_running:
                events.append((name, "start", c.get("image") or None))
            elif new_status in _STOPPED_STATES and new_exit not in (None, 0):
                events.append((name, "die", f"exit {new_exit}"))
            continue

        old_running = (old.get("status") or "").lower() == "running"
        new_restarts = c.get("restart_count") or 0
        old_restarts = old.get("restart_count") or 0
        if new_restarts > old_restarts:
            gained = new_restarts - old_restarts
            events.append((name, "restart", f"restart #{new_restarts}" if gained == 1
                           else f"{gained} restarts, now #{new_restarts}"))
        if new_running and not old_running and new_restarts <= old_restarts:
            events.append((name, "start", c.get("image") or None))
        elif old_running and not new_running:
            if new_exit not in (None, 0):
                events.append((name, "die", f"exit {new_exit}"))
            else:
                events.append((name, "stop", None))

        if new_oom and not bool(old.get("oom_killed")):
            events.append((name, "oom", None))

        if new_health != old.get("health") and (new_health or old.get("health")):
            events.append((name, "health_change",
                           f"{old.get('health') or '?'}→{new_health or '?'}"))

    # A running container that vanished from the listing entirely (removed).
    for name, old in old_state.items():
        if name not in new_by_name and (old.get("status") or "").lower() == "running":
            events.append((name, "stop", "removed"))

    return events


def _log_rows(log_batches, host_id: str):
    """Pure: flatten the agent's log batches → docker_logs row kwargs (no DB).

    `log_batches` is a list of (recorded_at, records) where each record is
    {"container", "stream", "ts", "line"}. A line's own Docker timestamp wins;
    recorded_at is the fallback when the agent couldn't parse one. Drops the
    agent's advisory truncation marker (container "_steward" — it signals a
    deferral, not a real container), malformed records, and lineless records;
    normalises an unknown stream to stdout. Unit-testable in isolation.
    """
    for recorded_at, records in log_batches:
        if not isinstance(records, list):
            continue
        for rec in records:
            if not isinstance(rec, dict):
                continue
            name = rec.get("container")
            line = rec.get("line")
            if not name or name == "_steward" or line is None:
                continue
            ts = _parse_started_at(rec.get("ts")) or recorded_at
            stream = rec.get("stream")
            if stream not in ("stdout", "stderr"):
                stream = "stdout"
            yield {"host_id": host_id, "container_name": str(name)[:255],
                   "ts": ts, "stream": stream, "line": str(line)}


async def _persist_logs(session, host, log_batches) -> None:
    """Append pushed container log lines (one row per line) for this host.

    Time-series / append-only — the per-container size+age ring (retention)
    bounds growth, so a chatty container just keeps a shorter window. The global
    toggle + exclude list (Settings, no restart) are enforced here: this is the
    authoritative drop point, since the push model has no channel to tell an
    agent to stop collecting.
    """
    from steward.core.settings import get_setting
    from .models import DockerLog

    if not await get_setting(session, "docker.logs.enabled"):
        return
    exclude = set(await get_setting(session, "docker.logs.exclude") or ())
    rows = [r for r in _log_rows(log_batches, host.id) if r["container_name"] not in exclude]
    stored = await _stored_lines(session, host.id, rows)
    for row in rows:
        if (row["container_name"], row["ts"], row["stream"], row["line"]) in stored:
            continue
        session.add(DockerLog(**row))


async def _stored_lines(session, host_id: str, rows: list) -> set:
    """The lines of `rows` this host already has, as (container, ts, stream,
    line). A restarted agent has lost its cursors and re-sends each
    container's last lines (#5809); a push of new lines overlaps nothing, so
    the one query, bounded per container to the batch's own span, finds none."""
    from sqlalchemy import and_, or_, select
    from .models import DockerLog

    spans: dict[str, list] = {}
    for r in rows:
        span = spans.setdefault(r["container_name"], [r["ts"], r["ts"]])
        span[0], span[1] = min(span[0], r["ts"]), max(span[1], r["ts"])
    if not spans:
        return set()
    found = await session.execute(
        select(DockerLog.container_name, DockerLog.ts, DockerLog.stream, DockerLog.line)
        .where(DockerLog.host_id == host_id, or_(*(
            and_(DockerLog.container_name == name, DockerLog.ts.between(lo, hi))
            for name, (lo, hi) in spans.items()))))
    return {tuple(r) for r in found}


async def _persist_swarm(session, host, swarm: dict) -> None:
    """Upsert this manager's swarm topology; drop rows no longer reported.

    Current-state tables (not time-series): a manager re-reports its full
    services/nodes set every sample, so we upsert what's present and delete what
    isn't (scoped to this host so two managers don't clobber each other).
    """
    from datetime import timezone
    from .models import DockerSwarmNode, DockerSwarmService

    now = datetime.now(timezone.utc)

    services = swarm.get("services") or []
    seen_services: set[str] = set()
    for s in services:
        name = s.get("service_name")
        if not name:
            continue
        seen_services.add(name)
        row = await session.get(DockerSwarmService, (host.id, name))
        if row is None:
            row = DockerSwarmService(host_id=host.id, service_name=name)
            session.add(row)
        row.mode = s.get("mode") or "replicated"
        row.desired = int(s.get("desired") or 0)
        row.running = int(s.get("running") or 0)
        row.image = s.get("image")
        row.placement_json = json.dumps(s.get("placement") or [])
        row.tasks_json = json.dumps(s.get("tasks") or [])
        row.updated_at = now
    stale_services = delete(DockerSwarmService).where(
        DockerSwarmService.host_id == host.id)
    if seen_services:
        stale_services = stale_services.where(
            DockerSwarmService.service_name.notin_(seen_services))
    await session.execute(stale_services)

    nodes = swarm.get("nodes") or []
    seen_nodes: set[str] = set()
    for n in nodes:
        nid = n.get("node_id")
        if not nid:
            continue
        seen_nodes.add(nid)
        row = await session.get(DockerSwarmNode, (host.id, nid))
        if row is None:
            row = DockerSwarmNode(host_id=host.id, node_id=nid)
            session.add(row)
        row.hostname = n.get("hostname") or ""
        row.role = n.get("role") or "worker"
        row.availability = n.get("availability") or "active"
        row.status = n.get("status") or "unknown"
        row.leader = bool(n.get("leader", False))
        row.updated_at = now
    stale_nodes = delete(DockerSwarmNode).where(DockerSwarmNode.host_id == host.id)
    if seen_nodes:
        stale_nodes = stale_nodes.where(DockerSwarmNode.node_id.notin_(seen_nodes))
    await session.execute(stale_nodes)


async def _persist_disk(session, host, disk: dict) -> None:
    """Upsert this host's /system/df summary + replace its image storage rows.

    Current-state (not time-series): the agent re-reports the full summary each
    interval, so we overwrite the single summary row and replace the image set
    (delete rows no longer present, upsert the rest), scoped to this host.
    """
    from datetime import timezone
    from .models import DockerDiskUsage, DockerImage

    now = datetime.now(timezone.utc)

    summary = await session.get(DockerDiskUsage, host.id)
    if summary is None:
        summary = DockerDiskUsage(host_id=host.id)
        session.add(summary)
    summary.layers_size = int(disk.get("layers_size") or 0)
    summary.images_size = int(disk.get("images_size") or 0)
    summary.images_reclaimable = int(disk.get("images_reclaimable") or 0)
    summary.containers_size = int(disk.get("containers_size") or 0)
    summary.volumes_size = int(disk.get("volumes_size") or 0)
    summary.build_cache_size = int(disk.get("build_cache_size") or 0)
    summary.images_total = int(disk.get("images_total") or 0)
    summary.images_active = int(disk.get("images_active") or 0)
    summary.containers_count = int(disk.get("containers_count") or 0)
    summary.volumes_count = int(disk.get("volumes_count") or 0)
    summary.updated_at = now
    # The totals as metrics too (#5304): a history of each host's growth, and
    # something an alert rule can watch. Resource: the host's name.
    from steward.core.alerts import MetricBatch
    from .disk_view import disk_metrics
    batch = MetricBatch()
    for metric, value in disk_metrics(disk).items():
        batch.add(source_module="docker", resource_name=host.name, metric_name=metric, value=value)
    await batch.record(session)

    images = disk.get("images") or []
    seen: set[str] = set()
    for im in images:
        iid = im.get("image_id")
        if not iid or iid in seen:
            continue
        seen.add(iid)
        row = await session.get(DockerImage, (host.id, iid))
        if row is None:
            row = DockerImage(host_id=host.id, image_id=iid)
            session.add(row)
        row.repo_tag = (im.get("repo_tag") or "<none>")[:512]
        row.size = int(im.get("size") or 0)
        row.shared_size = int(im.get("shared_size") or 0)
        row.containers = int(im.get("containers") or 0)
        row.updated_at = now
    stale = delete(DockerImage).where(DockerImage.host_id == host.id)
    if seen:
        stale = stale.where(DockerImage.image_id.notin_(seen))
    await session.execute(stale)


async def _persist_status(session, host, status: dict) -> None:
    """Record whether Docker answered on this host, and why not.

    A host reporting "unavailable" with no row yet never had Docker: nothing to
    record. Once a host has had Docker, losing it is recorded, so its last
    containers read as stale rather than as running."""
    from .models import DockerHostStatus

    row = await session.get(DockerHostStatus, host.id)
    if row is None:
        if status["status"] == "unavailable":
            return
        row = DockerHostStatus(host_id=host.id)
        session.add(row)
    if row.status != status["status"]:
        row.since = status["at"]
    row.status = status["status"]
    row.detail = (status.get("detail") or "")[:255] or None
    row.reported_at = status["at"]


async def _record_swarm_health(session, swarm: dict) -> None:
    """The swarm-wide figures alert rules watch (#5788), at most once every
    RECORD_EVERY whichever manager reported."""
    from datetime import timezone
    from steward.core.alerts import MetricBatch
    from . import swarm_health
    if not swarm_health.due(datetime.now(timezone.utc)):
        return
    batch = MetricBatch()
    for metric, value in swarm_health.figures(swarm).items():
        batch.add(source_module="docker", resource_name=swarm_health.RESOURCE, metric_name=metric, value=value)
    await batch.record(session)


async def _expire_swarm(session, host) -> None:
    """Drop the swarm rows a host reported while it was a manager."""
    from .models import DockerSwarmNode, DockerSwarmService

    await session.execute(delete(DockerSwarmService).where(DockerSwarmService.host_id == host.id))
    await session.execute(delete(DockerSwarmNode).where(DockerSwarmNode.host_id == host.id))


async def persist_host_docker(session, host, snapshots, swarm=None, disk=None,
                              logs=None, status=None, phases=None) -> None:
    """Upsert containers + time-series + lifecycle events + swarm for one host.

    `snapshots` is a list of (recorded_at: datetime, containers: list[dict]) —
    one entry per ingested sample carrying docker data (usually one; more when
    the agent flushes a backlog). Every snapshot's running containers record
    cpu_pct, mem_pct and mem_usage_bytes into the one metric store, at the
    snapshot's time; the newest snapshot drives current container state and
    lifecycle-event derivation. `swarm` is the newest
    sample's swarm object (or None off managers) — persisted when present.
    `disk` is the newest sample's /system/df summary (or None on Docker-less
    hosts) — persisted when present. `logs` is a list of (recorded_at, records)
    log batches (m79) — appended one row per line when present, independent of
    whether this sample also carried container metrics. `status` is the newest
    sample's {"status", "detail", "at", "swarm_manager"} (agent ≥ 1.9.0): an
    "ok" with an empty snapshot clears the host's containers, and a node that
    says it is no longer a manager loses the swarm rows it reported as one.

    `phases`, when given, is marked after each part (docker.status, .swarm,
    .disk, .logs, .metrics, .events, .containers, .alerts, .reap), each
    flushed before its mark, so a slow ingest names the part (#5181).
    """
    from steward.core.alerts import MetricBatch
    from .models import DockerContainer, DockerEvent

    async def mark(name: str) -> None:
        if phases is not None:
            await session.flush()
            phases.mark(f"docker.{name}")

    if status is not None:
        await _persist_status(session, host, status)
        if status.get("swarm_manager") is False:
            await _expire_swarm(session, host)
        await mark("status")
    if swarm is not None:
        await _persist_swarm(session, host, swarm)
        await _record_swarm_health(session, swarm)
        await mark("swarm")
    if disk is not None:
        await _persist_disk(session, host, disk)
        await mark("disk")
    if logs:
        await _persist_logs(session, host, logs)
        await mark("logs")

    if not snapshots:
        return
    ordered = sorted(snapshots, key=lambda s: s[0])
    latest_at, latest_containers = ordered[-1]

    # Time series for every snapshot (running containers with a CPU read), in
    # the one metric store (#5752): resource "<host>/<container>", so
    # containers of the same name on different hosts don't collide. Recorded
    # with the alert batch below, at each snapshot's own time.
    batch = MetricBatch()
    for recorded_at, containers in ordered:
        for c in containers:
            if c.get("status") != "running" or c.get("cpu_pct") is None or not c.get("name"):
                continue
            resource = f"{host.name}/{c['name']}"
            for metric in ("cpu_pct", "mem_pct", "mem_usage_bytes"):
                if c.get(metric) is not None:
                    batch.add(source_module="docker", resource_name=resource,
                              metric_name=metric, value=c[metric], recorded_at=recorded_at)
    await mark("metrics")

    # Snapshot of stored per-container state BEFORE the upsert overwrites it —
    # used to diff lifecycle events. Skip event derivation on the host's first
    # snapshot (empty state) so we don't emit a start per pre-existing container.
    old_rows = (await session.execute(
        select(DockerContainer).where(DockerContainer.host_id == host.id)
    )).scalars().all()
    old_state = {
        r.name: {"status": r.status, "health": r.health,
                 "oom_killed": r.oom_killed, "exit_code": r.exit_code,
                 "restart_count": r.restart_count}
        for r in old_rows
    }
    if old_state:
        # A stop is dated when Docker says it happened (agent ≥ 1.11), not when
        # the sample that noticed it was taken.
        finished = {c.get("name"): _parse_started_at(c.get("finished_at"))
                    for c in latest_containers}
        for name, event, detail in _derive_events(old_state, latest_containers):
            at = finished.get(name) if event in ("die", "stop") else None
            session.add(DockerEvent(
                host_id=host.id, container_name=name,
                event=event, detail=detail, at=at or latest_at,
            ))
    await mark("events")

    # Current state, and the restart and health alerts, from the newest snapshot only.
    for c in latest_containers:
        name = c.get("name")
        if not name:
            continue
        existing = await session.get(DockerContainer, (host.id, name))
        if existing is None:
            existing = DockerContainer(host_id=host.id, name=name)
            session.add(existing)
        existing.container_id = c.get("container_id", "") or ""
        existing.image = c.get("image", "") or ""
        existing.status = c.get("status", "unknown") or "unknown"
        existing.cpu_pct = c.get("cpu_pct")
        existing.mem_usage_bytes = c.get("mem_usage_bytes")
        existing.mem_limit_bytes = c.get("mem_limit_bytes")
        existing.mem_pct = c.get("mem_pct")
        existing.ports_json = json.dumps(c.get("ports") or [])
        existing.started_at = _parse_started_at(c.get("started_at"))
        existing.scraped_at = latest_at
        # Enrichment (agent ≥ 1.4.0; .get keeps older agents working — fields
        # stay None/0 when absent from the payload).
        existing.restart_count = c.get("restart_count", 0) or 0
        existing.health = c.get("health")
        existing.exit_code = c.get("exit_code")
        existing.oom_killed = bool(c.get("oom_killed", False))
        existing.compose_project = c.get("compose_project")
        existing.service_name = c.get("service_name")
        existing.task_id = c.get("task_id")
        existing.node_id = c.get("node_id")
        existing.net_rx_bytes = c.get("net_rx_bytes")
        existing.net_tx_bytes = c.get("net_tx_bytes")
        existing.blk_read_bytes = c.get("blk_read_bytes")
        existing.blk_write_bytes = c.get("blk_write_bytes")

        resource = f"{host.name}/{name}"
        # Restart count is alertable regardless of state (crash-looping matters
        # most while the container is down/restarting).
        batch.add(
            source_module="docker",
            resource_name=resource, metric_name="restart_count",
            value=float(c.get("restart_count", 0) or 0),
        )
        # Health → 1.0/0.0 only for containers that actually define a HEALTHCHECK
        # (health is None otherwise — recording 0 would false-alarm every plain
        # container).
        health = c.get("health")
        if health in ("healthy", "unhealthy", "starting"):
            batch.add(
                source_module="docker",
                resource_name=resource, metric_name="is_healthy",
                value=1.0 if health == "healthy" else 0.0,
            )

    await mark("containers")
    await batch.record(session)
    await mark("alerts")

    # Reap containers that vanished from the host's latest listing. Without this,
    # removed one-shot containers (CI job runners, buildkit builders, codex jobs)
    # accumulate forever as permanent "stopped" rows and inflate every count.
    # The listing is authoritative (event derivation already treats absence as
    # removal), so anything not in the newest snapshot no longer exists on the
    # host. Mirrors the image/swarm/disk persisters. An empty snapshot legitimately
    # means the host has no containers, so the unfiltered delete is correct.
    seen_names = {c["name"] for c in latest_containers if c.get("name")}
    reap = delete(DockerContainer).where(DockerContainer.host_id == host.id)
    if seen_names:
        reap = reap.where(DockerContainer.name.notin_(seen_names))
    await session.execute(reap)
    await mark("reap")