# plugins/unifi/scheduler.py
from __future__ import annotations
import json
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING

from steward.core.scheduler import ScheduledTask

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

_client = None  # UnifiClient instance, initialised on first scrape


async def _drop_client() -> None:
    """Close and discard the cached client.

    The client holds a long-lived ``httpx.AsyncClient`` (a connection pool over
    real sockets). Nulling the reference without ``aclose()`` orphans that pool
    until GC, so every failure tick that re-auths would leak file descriptors —
    the same class of bug that took the app down via SNMP (Errno 24, "Too many
    open files"). Close first, then drop. Best-effort: a failed close must not
    break the poll loop.
    """
    global _client
    if _client is not None:
        try:
            await _client.close()
        except Exception:
            pass
        _client = None


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

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

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


async def _do_poll(app: "Quart") -> None:
    global _client

    from .client import UnifiClient
    from .models import UnifiClientSnapshot, UnifiDevice, UnifiWanStat
    from steward.core.alerts import record_metric

    cfg = app.config["PLUGINS"]["unifi"]

    if _client is None:
        _client = UnifiClient(
            host=cfg["host"],
            username=cfg["username"],
            password=cfg["password"],
            site=cfg.get("site", "default"),
            verify_ssl=cfg.get("verify_ssl", False),
        )
        try:
            await _client.login()
        except Exception as exc:
            logger.warning("UniFi initial login failed: %s", exc)
            await _drop_client()
            return

    try:
        health = await _client.get_health()
        clients = await _client.get_active_clients()
        devices = await _client.get_devices()
    except Exception as exc:
        logger.warning("UniFi poll failed — will retry next tick: %s", exc)
        await _drop_client()  # close + force re-auth on next tick
        return

    scraped_at = datetime.now(timezone.utc)
    top_n = int(cfg.get("top_clients", 10))

    # ── WAN stat ──────────────────────────────────────────────────────────────
    wan_data = next((h for h in health if h.get("subsystem") == "wan"), None)
    wan_stat = UnifiWanStat(scraped_at=scraped_at)
    if wan_data:
        wan_stat.is_up = wan_data.get("status") == "ok"
        wan_stat.wan_ip = wan_data.get("wan_ip")
        wan_stat.latency_ms = wan_data.get("latency")
        wan_stat.rx_bytes_rate = wan_data.get("rx_bytes-r")
        wan_stat.tx_bytes_rate = wan_data.get("tx_bytes-r")
        wan_stat.uptime_seconds = wan_data.get("uptime")

    # ── Client snapshot ───────────────────────────────────────────────────────
    wireless = [c for c in clients if not c.get("is_wired", True)]
    wired = [c for c in clients if c.get("is_wired", True)]
    guests = [c for c in clients if c.get("is_guest", False)]

    top_clients = sorted(
        clients,
        key=lambda c: (c.get("tx_bytes-r", 0) or 0) + (c.get("rx_bytes-r", 0) or 0),
        reverse=True,
    )[:top_n]

    top_clients_data = [
        {
            "mac": c.get("mac", ""),
            "hostname": c.get("hostname") or c.get("name") or c.get("mac", ""),
            "ip": c.get("ip", ""),
            "type": "wireless" if not c.get("is_wired", True) else "wired",
            "tx_rate": c.get("tx_bytes-r", 0) or 0,
            "rx_rate": c.get("rx_bytes-r", 0) or 0,
            "signal": c.get("signal"),
            "essid": c.get("essid"),
        }
        for c in top_clients
    ]

    client_snapshot = UnifiClientSnapshot(
        scraped_at=scraped_at,
        total_clients=len(clients),
        wireless_clients=len(wireless),
        wired_clients=len(wired),
        guest_clients=len(guests),
        top_clients_json=json.dumps(top_clients_data),
    )

    async with app.db_sessionmaker() as session:
        async with session.begin():
            session.add(wan_stat)
            session.add(client_snapshot)

            # ── Devices (upsert by MAC) ───────────────────────────────────────
            for d in devices:
                mac = d.get("mac", "")
                if not mac:
                    continue
                last_seen_ts = d.get("last_seen")
                last_seen_dt = (
                    datetime.fromtimestamp(last_seen_ts, tz=timezone.utc)
                    if last_seen_ts else None
                )
                existing = await session.get(UnifiDevice, mac)
                if existing:
                    existing.name = d.get("name")
                    existing.model = d.get("model")
                    existing.device_type = d.get("type")
                    existing.ip = d.get("ip")
                    existing.state = d.get("state", 0)
                    existing.uptime_seconds = d.get("uptime")
                    existing.last_seen = last_seen_dt
                    existing.version = d.get("version")
                    existing.scraped_at = scraped_at
                else:
                    session.add(UnifiDevice(
                        mac=mac,
                        name=d.get("name"),
                        model=d.get("model"),
                        device_type=d.get("type"),
                        ip=d.get("ip"),
                        state=d.get("state", 0),
                        uptime_seconds=d.get("uptime"),
                        last_seen=last_seen_dt,
                        version=d.get("version"),
                        scraped_at=scraped_at,
                    ))

            # ── Alert metrics ─────────────────────────────────────────────────
            if wan_stat.latency_ms is not None:
                await record_metric(
                    session=session,
                    source_module="unifi",
                    resource_name="wan",
                    metric_name="latency_ms",
                    value=wan_stat.latency_ms,
                )
            await record_metric(
                session=session,
                source_module="unifi",
                resource_name="wan",
                metric_name="is_up",
                value=1.0 if wan_stat.is_up else 0.0,
            )
            await record_metric(
                session=session,
                source_module="unifi",
                resource_name="clients",
                metric_name="total_clients",
                value=float(len(clients)),
            )

    logger.debug(
        "UniFi poll: wan=%s, clients=%d, devices=%d",
        "up" if wan_stat.is_up else "down",
        len(clients),
        len(devices),
    )

    # ── Expanded data (best-effort, failures don't abort core poll) ───────────
    try:
        await _do_expanded(app, scraped_at)
    except Exception as exc:
        logger.warning("UniFi expanded poll failed: %s", exc)


async def _do_expanded(app: "Quart", scraped_at: datetime) -> None:
    """Fetch supplementary UniFi data: WLAN stats, events, alarms, DPI,
    known clients, speedtest, networks, port forwards, firewall rules."""
    from .models import (
        UnifiAlarm, UnifiDpiSnapshot, UnifiEvent, UnifiFwRule,
        UnifiKnownClient, UnifiNetwork, UnifiPortForward, UnifiSpeedtestResult,
        UnifiWlanStat,
    )

    # ── Fetch from controller ─────────────────────────────────────────────────
    try:
        wlan_configs = await _client.get_wlan_configs()
    except Exception:
        wlan_configs = []

    try:
        events = await _client.get_events(limit=200)
    except Exception:
        events = []

    try:
        alarms = await _client.get_alarms(archived=False)
    except Exception:
        alarms = []

    try:
        dpi_cats = await _client.get_dpi_site_stats()
    except Exception:
        dpi_cats = []

    try:
        dpi_clients = await _client.get_dpi_stats()
    except Exception:
        dpi_clients = []

    try:
        known_clients = await _client.get_known_clients()
    except Exception:
        known_clients = []

    try:
        speedtest = await _client.get_speedtest_status()
    except Exception:
        speedtest = None

    try:
        networks = await _client.get_networks()
    except Exception:
        networks = []

    try:
        port_forwards = await _client.get_port_forwards()
    except Exception:
        port_forwards = []

    try:
        fw_rules = await _client.get_firewall_rules()
    except Exception:
        fw_rules = []

    # ── Derive WLAN stats from device VAP table ───────────────────────────────
    # WLAN usage stats live on the device objects' vap_table, not wlan configs.
    # We get per-radio SSID data from the active clients grouped by essid.
    # Store one row per SSID using wlan_configs for channel/satisfaction.
    wlan_rows: list[UnifiWlanStat] = []
    for wc in wlan_configs:
        ssid = wc.get("name") or wc.get("x_passphrase", "")
        if not ssid:
            continue
        wlan_rows.append(UnifiWlanStat(
            scraped_at=scraped_at,
            ssid=ssid,
            radio=None,
            num_sta=wc.get("num_sta", 0) or 0,
            tx_bytes=wc.get("tx_bytes"),
            rx_bytes=wc.get("rx_bytes"),
            channel=None,
            satisfaction=wc.get("satisfaction"),
        ))

    # ── Build DPI category summary ────────────────────────────────────────────
    cat_summary = []
    for entry in dpi_cats:
        cat_id = entry.get("cat") or entry.get("cat_id")
        cat_name = entry.get("cat_name") or str(cat_id)
        rx = entry.get("rx_bytes", 0) or 0
        tx = entry.get("tx_bytes", 0) or 0
        if rx or tx:
            cat_summary.append({"cat_id": cat_id, "cat_name": cat_name, "rx_bytes": rx, "tx_bytes": tx})
    cat_summary.sort(key=lambda x: x["rx_bytes"] + x["tx_bytes"], reverse=True)

    # Top 10 clients by DPI total
    client_dpi: dict[str, dict] = {}
    for entry in dpi_clients:
        mac = entry.get("mac", "")
        if not mac:
            continue
        if mac not in client_dpi:
            client_dpi[mac] = {"mac": mac, "hostname": entry.get("hostname", mac), "bytes": 0, "top_cats": []}
        client_dpi[mac]["bytes"] += (entry.get("rx_bytes", 0) or 0) + (entry.get("tx_bytes", 0) or 0)
    top_dpi_clients = sorted(client_dpi.values(), key=lambda x: x["bytes"], reverse=True)[:10]

    # ── Write to DB ───────────────────────────────────────────────────────────
    async with app.db_sessionmaker() as session:
        async with session.begin():
            # WLAN stats
            for row in wlan_rows:
                session.add(row)

            # DPI snapshot (only if we got data)
            if cat_summary or top_dpi_clients:
                session.add(UnifiDpiSnapshot(
                    scraped_at=scraped_at,
                    by_category_json=json.dumps(cat_summary),
                    by_client_json=json.dumps(top_dpi_clients),
                ))

            # Events (upsert by _id)
            _category_map = {
                "wlan": "wlan", "wireless": "wlan",
                "wired": "wired",
                "ugw": "wan", "wan": "wan",
                "ids": "ids",
                "system": "system",
                "ap": "ap",
            }
            for ev in events:
                uid = ev.get("_id", "")
                if not uid:
                    continue
                ts_raw = ev.get("datetime") or ev.get("time")
                try:
                    if isinstance(ts_raw, (int, float)):
                        event_time = datetime.fromtimestamp(ts_raw, tz=timezone.utc)
                    else:
                        event_time = datetime.fromisoformat(str(ts_raw).replace("Z", "+00:00"))
                except Exception:
                    event_time = scraped_at
                subsystem = ev.get("subsystem", "system")
                category = _category_map.get(subsystem, subsystem[:16])
                details = {k: v for k, v in ev.items()
                           if k not in ("_id", "datetime", "time", "key", "subsystem", "msg", "user", "src_ip")}
                existing = await session.get(UnifiEvent, uid)
                if not existing:
                    session.add(UnifiEvent(
                        unifi_id=uid,
                        event_time=event_time,
                        event_key=ev.get("key", ""),
                        category=category,
                        message=ev.get("msg", ""),
                        source_mac=ev.get("user") or ev.get("src_mac"),
                        source_ip=ev.get("src_ip"),
                        details_json=json.dumps(details),
                    ))

            # Alarms (upsert)
            for al in alarms:
                uid = al.get("_id", "")
                if not uid:
                    continue
                ts_raw = al.get("datetime") or al.get("time")
                try:
                    if isinstance(ts_raw, (int, float)):
                        alarm_time = datetime.fromtimestamp(ts_raw, tz=timezone.utc)
                    else:
                        alarm_time = datetime.fromisoformat(str(ts_raw).replace("Z", "+00:00"))
                except Exception:
                    alarm_time = scraped_at
                existing = await session.get(UnifiAlarm, uid)
                if existing:
                    existing.archived = al.get("archived", False)
                    existing.scraped_at = scraped_at
                else:
                    session.add(UnifiAlarm(
                        unifi_id=uid,
                        alarm_time=alarm_time,
                        alarm_key=al.get("key", ""),
                        message=al.get("msg", ""),
                        archived=al.get("archived", False),
                        scraped_at=scraped_at,
                    ))

            # Known clients (upsert by MAC)
            for kc in known_clients:
                mac = kc.get("mac", "")
                if not mac:
                    continue
                def _ts(val):
                    if not val:
                        return None
                    try:
                        if isinstance(val, (int, float)):
                            return datetime.fromtimestamp(val, tz=timezone.utc)
                        return datetime.fromisoformat(str(val).replace("Z", "+00:00"))
                    except Exception:
                        return None
                existing = await session.get(UnifiKnownClient, mac)
                if existing:
                    existing.hostname = kc.get("hostname")
                    existing.name = kc.get("name")
                    existing.ip = kc.get("ip")
                    existing.mac_vendor = kc.get("oui")
                    existing.note = kc.get("note")
                    existing.is_blocked = kc.get("blocked", False)
                    existing.first_seen = _ts(kc.get("first_seen"))
                    existing.last_seen = _ts(kc.get("last_seen"))
                    existing.scraped_at = scraped_at
                else:
                    session.add(UnifiKnownClient(
                        mac=mac,
                        hostname=kc.get("hostname"),
                        name=kc.get("name"),
                        ip=kc.get("ip"),
                        mac_vendor=kc.get("oui"),
                        note=kc.get("note"),
                        is_blocked=kc.get("blocked", False),
                        first_seen=_ts(kc.get("first_seen")),
                        last_seen=_ts(kc.get("last_seen")),
                        scraped_at=scraped_at,
                    ))

            # Speedtest (upsert by run_at)
            if speedtest:
                ts_raw = speedtest.get("rundate") or speedtest.get("time")
                try:
                    if isinstance(ts_raw, (int, float)):
                        run_at = datetime.fromtimestamp(ts_raw, tz=timezone.utc)
                    else:
                        run_at = datetime.fromisoformat(str(ts_raw).replace("Z", "+00:00"))
                except Exception:
                    run_at = scraped_at
                existing = await session.get(UnifiSpeedtestResult, run_at)
                if not existing:
                    session.add(UnifiSpeedtestResult(
                        run_at=run_at,
                        download_mbps=speedtest.get("xput_download"),
                        upload_mbps=speedtest.get("xput_upload"),
                        latency_ms=speedtest.get("latency"),
                        server_name=speedtest.get("server_name"),
                    ))

            # Networks (upsert by UniFi ID)
            for net in networks:
                uid = net.get("_id", "")
                if not uid:
                    continue
                existing = await session.get(UnifiNetwork, uid)
                if existing:
                    existing.name = net.get("name", "")
                    existing.purpose = net.get("purpose")
                    existing.vlan = net.get("vlan")
                    existing.ip_subnet = net.get("ip_subnet")
                    existing.dhcp_enabled = net.get("dhcpd_enabled", False)
                    existing.is_guest = net.get("is_guest", False)
                    existing.scraped_at = scraped_at
                else:
                    session.add(UnifiNetwork(
                        unifi_id=uid,
                        name=net.get("name", ""),
                        purpose=net.get("purpose"),
                        vlan=net.get("vlan"),
                        ip_subnet=net.get("ip_subnet"),
                        dhcp_enabled=net.get("dhcpd_enabled", False),
                        is_guest=net.get("is_guest", False),
                        scraped_at=scraped_at,
                    ))

            # Port forwards (upsert by UniFi ID)
            for pf in port_forwards:
                uid = pf.get("_id", "")
                if not uid:
                    continue
                existing = await session.get(UnifiPortForward, uid)
                if existing:
                    existing.name = pf.get("name")
                    existing.proto = pf.get("proto", "tcp")
                    existing.dst_port = pf.get("dst_port", "")
                    existing.fwd_ip = pf.get("fwd", "")
                    existing.fwd_port = pf.get("fwd_port", "")
                    existing.src_filter = pf.get("src")
                    existing.enabled = pf.get("enabled", True)
                    existing.scraped_at = scraped_at
                else:
                    session.add(UnifiPortForward(
                        unifi_id=uid,
                        name=pf.get("name"),
                        proto=pf.get("proto", "tcp"),
                        dst_port=pf.get("dst_port", ""),
                        fwd_ip=pf.get("fwd", ""),
                        fwd_port=pf.get("fwd_port", ""),
                        src_filter=pf.get("src"),
                        enabled=pf.get("enabled", True),
                        scraped_at=scraped_at,
                    ))

            # Firewall rules (upsert by UniFi ID)
            for rule in fw_rules:
                uid = rule.get("_id", "")
                if not uid:
                    continue
                existing = await session.get(UnifiFwRule, uid)
                if existing:
                    existing.name = rule.get("name", "")
                    existing.ruleset = rule.get("ruleset", "")
                    existing.rule_index = rule.get("rule_index")
                    existing.action = rule.get("action", "")
                    existing.protocol = rule.get("protocol")
                    existing.enabled = rule.get("enabled", True)
                    existing.src_address = rule.get("src_address")
                    existing.dst_address = rule.get("dst_address")
                    existing.scraped_at = scraped_at
                else:
                    session.add(UnifiFwRule(
                        unifi_id=uid,
                        name=rule.get("name", ""),
                        ruleset=rule.get("ruleset", ""),
                        rule_index=rule.get("rule_index"),
                        action=rule.get("action", ""),
                        protocol=rule.get("protocol"),
                        enabled=rule.get("enabled", True),
                        src_address=rule.get("src_address"),
                        dst_address=rule.get("dst_address"),
                        scraped_at=scraped_at,
                    ))

    logger.debug(
        "UniFi expanded: wlans=%d, events=%d, alarms=%d, known_clients=%d, "
        "networks=%d, port_fwds=%d, fw_rules=%d",
        len(wlan_rows), len(events), len(alarms), len(known_clients),
        len(networks), len(port_forwards), len(fw_rules),
    )