kings-guard/src/kings_guard/stream.py

from __future__ import annotations

from collections.abc import Sequence
from datetime import UTC, datetime

from kings_guard.cadence import EmissionCadence, parse_timestamp
from kings_guard.contracts import (
    STREAM_FINDING_PREFIX,
    ImmuneObservation,
    ReconciliationView,
    StreamAssessment,
    StreamCompleteness,
    StreamHeartbeat,
)


def evaluate_stream(
    observations: Sequence[ImmuneObservation],
    cadence: EmissionCadence,
    *,
    now: str | datetime,
    watching_since: str | datetime | None = None,
    heartbeats: Sequence[StreamHeartbeat] = (),
    reconciliation: ReconciliationView | None = None,
) -> StreamAssessment:
    """Evaluate the stream, not its contents.

    Findings are prefixed with `stream:` so they cannot be mistaken for
    content findings such as `credential_exfil_probe`.
    """
    now_dt = _as_datetime(now)
    findings: list[str] = []
    reasons: list[str] = []
    observations = tuple(
        item for item in observations
        if item.source_system == cadence.source_system and parse_timestamp(item.timestamp) <= now_dt
    )
    heartbeats = tuple(
        item for item in heartbeats
        if item.source_system == cadence.source_system and parse_timestamp(item.timestamp) <= now_dt
    )
    observed_counts = _count_observations(observations)

    watch_start = _as_datetime(watching_since) if watching_since is not None else None
    if watch_start is None:
        timestamps = [parse_timestamp(item.timestamp) for item in observations]
        timestamps.extend(parse_timestamp(item.timestamp) for item in heartbeats)
        watch_start = min(timestamps) if timestamps else now_dt

    for rate in cadence.rates:
        window_start = now_dt - rate.window
        count = sum(
            1
            for item in observations
            if item.event_class == rate.event_class
            and parse_timestamp(item.timestamp) >= window_start
        )
        observed_counts[rate.event_class] = count
        watched_long_enough = (now_dt - watch_start) >= rate.window
        if not watched_long_enough:
            reasons.append(f"rate window for {rate.event_class} has not been fully observed")
        if watched_long_enough and count < rate.expected_min:
            findings.append(f"{STREAM_FINDING_PREFIX}cadence_unmet:{rate.event_class}")
            reasons.append(
                f"declared rate for {rate.event_class} is at least {rate.expected_min} "
                f"per {int(rate.window.total_seconds())}s; observed {count}"
            )

    for spec in cadence.heartbeats:
        due = (now_dt - watch_start) >= spec.interval
        latest = _latest_heartbeat(heartbeats, spec.event_class, spec.covered_event_class)
        missing = False
        if latest is None:
            missing = due
        else:
            missing = (now_dt - parse_timestamp(latest.timestamp)) > spec.interval
        if missing:
            findings.append(f"{STREAM_FINDING_PREFIX}heartbeat_missing:{spec.covered_event_class}")
            reasons.append(
                f"no {spec.assertion} heartbeat for {spec.covered_event_class} "
                f"within {int(spec.interval.total_seconds())}s"
            )
        elif not due and latest is None:
            reasons.append(
                f"heartbeat for {spec.covered_event_class} is not yet due; "
                "the stream cannot be vouched for until a positive claim arrives"
            )

    if cadence.reconciliations:
        if reconciliation is None:
            reasons.append("reconciliation view was not supplied; divergence cannot be ruled out")
        else:
            for spec in cadence.reconciliations:
                source_count = reconciliation.source_counts.get(spec.covered_event_class)
                evidence_count = reconciliation.evidence_counts.get(spec.covered_event_class)
                if any(type(value) is not int or value < 0
                       for value in (source_count, evidence_count)):
                    reasons.append(
                        "valid reconciliation counts were not supplied for "
                        f"{spec.covered_event_class}"
                    )
                    continue
                if evidence_count != source_count:
                    findings.append(
                        f"{STREAM_FINDING_PREFIX}reconciliation_divergence:{spec.covered_event_class}"
                    )
                    reasons.append(
                        f"{spec.covered_event_class} source transitions={source_count} "
                        f"but evidence count={evidence_count}"
                    )

    completeness = _completeness(findings, reasons, cadence)
    reason = (
        "; ".join(reasons) if reasons else "declared cadence is met and no stream finding is open"
    )
    return StreamAssessment(
        completeness=completeness,
        reason=reason,
        findings=tuple(findings),
        observed_counts=observed_counts,
        window_end=now_dt.isoformat().replace("+00:00", "Z"),
    )


def _completeness(
    findings: Sequence[str],
    reasons: Sequence[str],
    cadence: EmissionCadence,
) -> StreamCompleteness:
    if findings:
        return StreamCompleteness.DEGRADED
    if reasons:
        return StreamCompleteness.UNKNOWN
    if not cadence.heartbeats and not cadence.rates and not cadence.reconciliations:
        return StreamCompleteness.UNKNOWN
    return StreamCompleteness.COMPLETE


def _count_observations(observations: Sequence[ImmuneObservation]) -> dict[str, int]:
    counts: dict[str, int] = {}
    for item in observations:
        counts[item.event_class] = counts.get(item.event_class, 0) + 1
    return counts


def _latest_heartbeat(
    heartbeats: Sequence[StreamHeartbeat],
    event_class: str,
    covered_event_class: str,
) -> StreamHeartbeat | None:
    matching = [
        item for item in heartbeats
        if item.event_class == event_class
        and covered_event_class in item.covered_event_classes
    ]
    if not matching:
        return None
    return max(matching, key=lambda item: parse_timestamp(item.timestamp))


def _as_datetime(value: str | datetime) -> datetime:
    if isinstance(value, datetime):
        if value.tzinfo is None:
            return value.replace(tzinfo=UTC)
        return value.astimezone(UTC)
    return parse_timestamp(value)