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

import asyncio
import logging
from datetime import datetime, timedelta, 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

# How long a device's discovered table rows stand before they are walked
# again: a disk swapped or a volume added shows within the hour, and a NAS is
# walked 24 times a day, not 1,440.
TABLES_REFRESH = timedelta(hours=1)


def tables_due(device, now: datetime) -> bool:
    """Whether this poll walks the device's vendor tables: it reads some, it
    answered its last poll (a device that isn't answering is only asked its
    GET), and its rows are unknown, reset by a save, or an hour old."""
    from .readings import tables_for
    return (bool(tables_for(device.presets or [])) and not device.last_error
            and (device.tables_at is None or now - device.tables_at >= TABLES_REFRESH))


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, walk=None) -> None:
    """Poll every device at once, then write each one's outcome to its row
    and its answers to the metric tables. A device whose vendor tables are
    due is walked first, so a new disk is read in the same poll. `poll` and
    `walk` stand in for poller.poll_device and poller.walk_tables in tests."""
    global RATES
    from steward.core.alerts import MetricBatch

    from .models import SnmpDevice
    from .poller import poll_device, walk_tables
    from .readings import state_of, tables_for
    from .service import RateKeeper, credentials, readings_for

    poll = poll or poll_device
    walk = walk or walk_tables
    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

    started = datetime.now(timezone.utc)
    scans, asked = {}, {}

    async def one(d):
        creds = credentials(d)
        if tables_due(d, started):
            scans[d.id] = await walk(d.address, d.port, creds, tables_for(d.presets or []))
        found = scans.get(d.id)
        asked[d.id] = readings_for(d, found.rows if found is not None and found.error is None else None)
        return await poll(d.address, d.port, creds, [(r.oid, r.key) for r in asked[d.id]])

    results = await asyncio.gather(*(one(d) 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
                scan = scans.get(device.id)
                if scan is not None and scan.error is None:
                    row.tables, row.tables_at = scan.rows, now
                row.last_error = result.error
                missing = set(result.missing)
                if result.readings or result.texts:
                    row.last_ok_at = now
                texts = dict(row.texts or {}) if result.error else {}
                for reading in asked[device.id]:
                    value = result.readings.get(reading.key)
                    if reading.unit == "state":
                        # Its word is shown as answered; its level is stored,
                        # unless Steward doesn't know the word (level None).
                        raw = value if value is not None else result.texts.get(reading.key)
                        if raw is None:
                            continue
                        texts[reading.key], value = state_of(reading, raw)
                    elif value is None:
                        if reading.key in result.texts:      # a word where a number was asked for
                            missing.add(reading.key)
                        continue
                    elif reading.counter:
                        value = RATES.rate(device.id, reading, value, now)
                    else:
                        value = value * reading.scale + reading.offset
                    if value is None:
                        continue
                    batch.add(source_module="snmp", resource_name=device.name,
                              metric_name=reading.key, value=value, recorded_at=now)
                row.missing = sorted(missing)
                row.texts = texts
            await batch.record(session)