# plugins/snmp/scheduler.py
from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from sqlalchemy import select

from steward.core.scheduler import ScheduledTask

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# Counters' previous samples, for their rates (service.RateKeeper).
RATES = None


def make_poll_task(app: "Quart") -> ScheduledTask:
    interval = app.config["PLUGINS"]["snmp"].get("poll_interval_seconds", 60)

    async def poll() -> None:
        await _do_poll(app)

    return ScheduledTask(
        name="snmp_poll",
        coro_factory=poll,
        interval_seconds=int(interval),
        run_on_startup=True,
    )


async def _do_poll(app: "Quart", poll=None) -> None:
    """Poll every device at once, then write each one's outcome to its row
    and its answers to the metric tables. `poll` stands in for
    poller.poll_device in tests."""
    global RATES
    from steward.core.alerts import MetricBatch

    from .models import SnmpDevice
    from .poller import poll_device
    from .service import RateKeeper, credentials, readings_for

    poll = poll or poll_device
    if RATES is None:
        RATES = RateKeeper()

    async with app.db_sessionmaker() as session:
        devices = (await session.execute(select(SnmpDevice))).scalars().all()
    if not devices:
        return

    asked = {d.id: readings_for(d) for d in devices}
    results = await asyncio.gather(*(
        poll(d.address, d.port, credentials(d), [(r.oid, r.key) for r in asked[d.id]])
        for d in devices), return_exceptions=True)

    now = datetime.now(timezone.utc)
    async with app.db_sessionmaker() as session:
        async with session.begin():
            batch = MetricBatch()
            for device, result in zip(devices, results):
                row = await session.get(SnmpDevice, device.id)
                if row is None:          # deleted while it was being polled
                    continue
                row.last_polled_at = now
                if isinstance(result, BaseException):
                    logger.warning("SNMP poll of %s failed: %s", device.name, result)
                    row.last_error = f"SNMP poll failed: {result}"
                    continue
                row.last_error = result.error
                row.missing = sorted(set(result.missing))
                if result.readings:
                    row.last_ok_at = now
                for reading in asked[device.id]:
                    value = result.readings.get(reading.key)
                    if value is None:
                        continue
                    if reading.counter:
                        value = RATES.rate(device.id, reading, value, now)
                        if value is None:
                            continue
                    else:
                        value = value * reading.scale + reading.offset
                    batch.add(source_module="snmp", resource_name=device.name,
                              metric_name=reading.key, value=value, recorded_at=now)
            await batch.record(session)