Aligning asynchronous sensor timestamps in Python

import logging
from datetime import datetime, timezone
from typing import Optional, Union

from pydantic import BaseModel, Field, ValidationError, field_validator

logger = logging.getLogger(__name__)


class SensorTelemetry(BaseModel):
    sensor_id: str
    temperature_c: float = Field(ge=-80.0, le=80.0)
    raw_timestamp: str
    zone_id: Optional[str] = None
    gateway_id: str

    @field_validator("raw_timestamp", mode="before")
    @classmethod
    def normalize_timestamp(cls, v: Union[str, int, float]) -> str:
        """Standardize incoming timestamp strings to UTC-aware ISO 8601.

        Heuristic for numeric inputs: epochs above 1e10 (≈ year 2286 in
        seconds) are treated as milliseconds. Microsecond epochs would be
        ≥ 1e13 and must be detected by digit length, not value comparison.
        """
        if isinstance(v, (int, float)):
            epoch = float(v)
            if epoch >= 1e16:        # nanoseconds since epoch
                epoch /= 1_000_000_000.0
            elif epoch >= 1e13:      # microseconds
                epoch /= 1_000_000.0
            elif epoch >= 1e10:      # milliseconds
                epoch /= 1_000.0
            return datetime.fromtimestamp(epoch, tz=timezone.utc).isoformat()
        if isinstance(v, str) and v.endswith("Z"):
            return v.replace("Z", "+00:00")
        return v

    def get_utc_datetime(self) -> datetime:
        """Parse and return timezone-aware UTC datetime using standard library."""
        try:
            dt = datetime.fromisoformat(self.raw_timestamp)
        except ValueError as e:
            raise ValueError(f"Invalid timestamp format: {self.raw_timestamp}") from e
        if dt.tzinfo is None:
            return dt.replace(tzinfo=timezone.utc)
        return dt.astimezone(timezone.utc)

from typing import List

import pandas as pd


def normalize_and_sort_telemetry(records: List[SensorTelemetry]) -> pd.DataFrame:
    """Convert validated records to a UTC-indexed DataFrame, deterministically sorted."""
    if not records:
        return pd.DataFrame(
            columns=["sensor_id", "zone_id", "temperature_c", "raw_ts"]
        ).set_index(pd.DatetimeIndex([], tz="UTC", name="timestamp"))

    df = pd.DataFrame([
        {
            "timestamp": rec.get_utc_datetime(),
            "sensor_id": rec.sensor_id,
            "zone_id": rec.zone_id,
            "temperature_c": rec.temperature_c,
            "raw_ts": rec.raw_timestamp,
        }
        for rec in records
    ])
    df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True)
    # Single stable sort breaks timestamp ties by sensor_id, then sets the index.
    df = df.sort_values(by=["timestamp", "sensor_id"], kind="stable").set_index("timestamp")
    return df

def align_to_compliance_grid(df: pd.DataFrame, freq: str = "1min") -> pd.DataFrame:
    """Resample telemetry to a strict compliance grid without speculative interpolation.

    Forward-fill is capped at the configured tolerance; any cell still NaN
    after the bounded ffill is left NaN and surfaces as an audit-visible gap.
    """
    aligned = (
        df.groupby(["sensor_id", "zone_id"])["temperature_c"]
        .resample(freq)
        .mean()
    )
    aligned_df = aligned.reset_index().rename(columns={"temperature_c": "avg_temp_c"})
    aligned_df = aligned_df.sort_values(["sensor_id", "timestamp"])

    # Capture which cells were NaN before any imputation so we can flag the
    # filled rows (and leave still-NaN rows as visible gaps).
    pre_fill_na = aligned_df["avg_temp_c"].isna()
    aligned_df["avg_temp_c"] = (
        aligned_df.groupby("sensor_id")["avg_temp_c"].ffill(limit=2)
    )
    aligned_df["is_interpolated"] = pre_fill_na & aligned_df["avg_temp_c"].notna()

    return aligned_df

import asyncio
from collections import deque

BATCH_SIZE = 5000
QUARANTINE_QUEUE = deque()

async def process_telemetry_stream(raw_stream: asyncio.Queue, output_sink: asyncio.Queue):
    """Async worker that validates, aligns, and batches telemetry."""
    batch = []
    while True:
        try:
            payload = await asyncio.wait_for(raw_stream.get(), timeout=1.0)
            try:
                validated = SensorTelemetry(**payload)
                batch.append(validated)
            except ValidationError as e:
                QUARANTINE_QUEUE.append({"payload": payload, "error": str(e)})
                continue

            if len(batch) >= BATCH_SIZE:
                df = normalize_and_sort_telemetry(batch)
                aligned = align_to_compliance_grid(df)
                await output_sink.put(aligned)
                batch.clear()

        except asyncio.TimeoutError:
            if batch:
                df = normalize_and_sort_telemetry(batch)
                aligned = align_to_compliance_grid(df)
                await output_sink.put(aligned)
                batch.clear()