#!/usr/bin/env python3
"""
WiFi Recon Dashboard
Reads from events.db (written by udp_ingest.py), serves a live web UI.
"""

import asyncio
import sqlite3
import datetime
import os
import json
from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse, FileResponse

DB_PATH      = os.path.join(os.path.dirname(__file__), "events.db")
HTML_PATH    = os.path.join(os.path.dirname(__file__), "templates", "dashboard.html")
CSS_PATH     = os.path.join(os.path.dirname(__file__), "static", "dashboard.css")
JS_PATH      = os.path.join(os.path.dirname(__file__), "static", "dashboard.js")
OUI_PATH     = os.path.join(os.path.dirname(__file__), "oui.txt")
ONLINE_SECS  = 30   # node considered online if heartbeat within this many seconds
EVENT_CAP    = 100  # max events returned / shown in dashboard

# ─── OUI lookup ──────────────────────────────────────────────────────────────

def _load_oui(path: str) -> dict[str, str]:
    db: dict[str, str] = {}
    if not os.path.exists(path):
        return db
    with open(path, encoding="utf-8", errors="ignore") as f:
        for line in f:
            if "(hex)" not in line:
                continue
            parts = line.split("(hex)")
            prefix = parts[0].strip().replace("-", ":").upper()   # "AA-BB-CC" → "AA:BB:CC"
            vendor = parts[1].strip()
            if prefix and vendor:
                db[prefix] = vendor
    return db

_OUI_DB: dict[str, str] = _load_oui(OUI_PATH)


def oui_lookup(mac: str) -> tuple[str, bool]:
    """
    Returns (vendor_name, is_randomized).
    Randomized MACs have the locally-administered bit set (bit 1 of first byte).
    OUI lookup is meaningless for randomized MACs.
    """
    if not mac or len(mac) < 8:
        return ("Unknown", False)
    try:
        first_byte = int(mac.replace(":", "").replace("-", "")[:2], 16)
        randomized = bool(first_byte & 0x02)
    except ValueError:
        return ("Unknown", False)
    if randomized:
        return ("Randomized MAC", True)
    prefix = mac.upper()[:8]
    return (_OUI_DB.get(prefix, "Unknown"), False)

# ─── DB helpers ──────────────────────────────────────────────────────────────

def get_conn() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def query(sql: str, params: tuple = ()) -> list[dict]:
    with get_conn() as conn:
        rows = conn.execute(sql, params).fetchall()
    return [dict(r) for r in rows]


def _node_status(last_seen: str | None, threshold: int = ONLINE_SECS) -> str:
    if not last_seen:
        return "offline"
    try:
        dt = datetime.datetime.fromisoformat(last_seen)
        delta = (datetime.datetime.utcnow() - dt).total_seconds()
        return "online" if delta <= threshold else "offline"
    except ValueError:
        return "offline"

# ─── API data builders ────────────────────────────────────────────────────────

def build_nodes() -> list[dict]:
    rows = query("""
        SELECT
            node_id,
            MAX(received_at)  AS last_beacon,
            COUNT(*)          AS total_events,
            MIN(received_at)  AS first_seen
        FROM beacon_events
        GROUP BY node_id
        ORDER BY last_beacon DESC
    """)
    # Latest heartbeat per node
    hb_rows = query("""
        SELECT node_id, received_at AS last_heartbeat,
               uptime_ms, free_heap, wifi_rssi
        FROM heartbeat_events
        WHERE id IN (SELECT MAX(id) FROM heartbeat_events GROUP BY node_id)
    """)
    hb_map = {r["node_id"]: r for r in hb_rows}
    for r in rows:
        hb = hb_map.get(r["node_id"], {})
        r["last_heartbeat"] = hb.get("last_heartbeat")
        r["uptime_ms"]      = hb.get("uptime_ms")
        r["free_heap"]      = hb.get("free_heap")
        r["wifi_rssi"]      = hb.get("wifi_rssi")
        r["last_seen"]      = r["last_heartbeat"] or r["last_beacon"]
        # Online if heartbeat within 30s; fall back to beacon within 60s if no heartbeat yet
        if r["last_heartbeat"]:
            r["status"] = _node_status(r["last_heartbeat"], threshold=30)
        else:
            r["status"] = _node_status(r["last_beacon"], threshold=60)
    return rows


def build_events(limit: int = EVENT_CAP, node_id: str | None = None) -> list[dict]:
    if node_id:
        return query("""
            SELECT * FROM beacon_events
            WHERE node_id = ?
            ORDER BY id DESC LIMIT ?
        """, (node_id, limit))
    return query("""
        SELECT * FROM beacon_events
        ORDER BY id DESC LIMIT ?
    """, (limit,))


def build_networks() -> list[dict]:
    rows = query("""
        SELECT
            ssid,
            bssid,
            MAX(rssi)                AS best_rssi,
            ROUND(AVG(rssi), 1)      AS avg_rssi,
            channel,
            encryption,
            COUNT(*)                 AS times_seen,
            COUNT(DISTINCT node_id)  AS node_count,
            MAX(received_at)         AS last_seen
        FROM beacon_events
        WHERE bssid IS NOT NULL
        GROUP BY bssid
        ORDER BY times_seen DESC
    """)
    return rows


def build_clients() -> list[dict]:
    """
    Per unique src_mac: aggregated stats and the SSIDs they were probing for.
    """
    rows = query("""
        SELECT
            src_mac,
            COUNT(*)                        AS times_seen,
            COUNT(DISTINCT node_id)         AS node_count,
            MAX(rssi)                        AS best_rssi,
            ROUND(AVG(rssi), 1)             AS avg_rssi,
            COUNT(DISTINCT NULLIF(ssid,'')) AS unique_ssids,
            MIN(received_at)                AS first_seen,
            MAX(received_at)                AS last_seen
        FROM probe_events
        WHERE src_mac IS NOT NULL
        GROUP BY src_mac
        ORDER BY times_seen DESC
    """)
    # Fetch all probed SSIDs in one query, group in Python — avoids N+1
    ssid_rows = query("""
        SELECT DISTINCT src_mac, ssid FROM probe_events
        WHERE src_mac IS NOT NULL AND ssid IS NOT NULL AND ssid != ''
        ORDER BY src_mac, ssid
    """)
    ssids_by_mac: dict = {}
    for s in ssid_rows:
        ssids_by_mac.setdefault(s["src_mac"], []).append(s["ssid"])

    for r in rows:
        r["probed_ssids"] = ssids_by_mac.get(r["src_mac"], [])
        vendor, randomized = oui_lookup(r["src_mac"])
        r["vendor"]     = vendor
        r["randomized"] = randomized
    return rows


def build_cross_node() -> dict:
    """
    For every BSSID, return per-node RSSI stats side by side.
    Result: { node_ids: [...], networks: [{bssid, ssid, channel, encryption,
              node_count, nodes: {node_id: {best_rssi, avg_rssi, times_seen}}}] }
    """
    node_rows = query("SELECT DISTINCT node_id FROM beacon_events ORDER BY node_id")
    node_ids  = [r["node_id"] for r in node_rows]

    rows = query("""
        SELECT
            bssid,
            MAX(ssid)            AS ssid,
            MAX(channel)         AS channel,
            MAX(encryption)      AS encryption,
            node_id,
            MAX(rssi)            AS best_rssi,
            ROUND(AVG(rssi), 1)  AS avg_rssi,
            COUNT(*)             AS times_seen
        FROM beacon_events
        WHERE bssid IS NOT NULL
        GROUP BY bssid, node_id
        ORDER BY bssid, node_id
    """)

    networks: dict = {}
    for r in rows:
        bssid = r["bssid"]
        if bssid not in networks:
            networks[bssid] = {
                "bssid":      bssid,
                "ssid":       r["ssid"],
                "channel":    r["channel"],
                "encryption": r["encryption"],
                "nodes":      {},
            }
        networks[bssid]["nodes"][r["node_id"]] = {
            "best_rssi":  r["best_rssi"],
            "avg_rssi":   r["avg_rssi"],
            "times_seen": r["times_seen"],
        }

    result = sorted(
        networks.values(),
        key=lambda x: len(x["nodes"]),
        reverse=True,
    )
    for net in result:
        net["node_count"] = len(net["nodes"])

    return {"node_ids": node_ids, "networks": result}


def build_presence() -> dict:
    now        = datetime.datetime.utcnow()
    cutoff_1h  = (now - datetime.timedelta(hours=1)).isoformat(timespec="seconds")
    cutoff_24h = (now - datetime.timedelta(hours=24)).isoformat(timespec="seconds")

    # ── Present now: real MACs seen 2+ times in the last hour ────────────────
    present_rows = query("""
        SELECT src_mac,
               COUNT(*)              AS times_seen,
               MAX(rssi)             AS best_rssi,
               ROUND(AVG(rssi), 1)   AS avg_rssi,
               MAX(received_at)      AS last_seen,
               COUNT(DISTINCT node_id) AS node_count
        FROM probe_events
        WHERE src_mac IS NOT NULL AND received_at > ?
        GROUP BY src_mac
        HAVING COUNT(*) > 1
        ORDER BY last_seen DESC
    """, (cutoff_1h,))
    present = []
    for r in present_rows:
        vendor, randomized = oui_lookup(r["src_mac"])
        if not randomized:
            r["vendor"] = vendor
            present.append(r)

    # ── New arrivals: first seen in last 24h ─────────────────────────────────
    new_mac_rows = query("""
        SELECT src_mac,
               MIN(received_at)      AS first_seen,
               MAX(rssi)             AS best_rssi,
               COUNT(*)              AS times_seen,
               COUNT(DISTINCT node_id) AS node_count
        FROM probe_events
        WHERE src_mac IS NOT NULL
        GROUP BY src_mac
        HAVING MIN(received_at) > ?
        ORDER BY first_seen DESC
        LIMIT 100
    """, (cutoff_24h,))
    new_macs = []
    for r in new_mac_rows:
        vendor, randomized = oui_lookup(r["src_mac"])
        r["vendor"]     = vendor
        r["randomized"] = randomized
        new_macs.append(r)

    new_networks = query("""
        SELECT bssid,
               MAX(ssid)             AS ssid,
               MIN(received_at)      AS first_seen,
               MAX(rssi)             AS best_rssi,
               MAX(channel)          AS channel,
               MAX(encryption)       AS encryption,
               COUNT(DISTINCT node_id) AS node_count
        FROM beacon_events
        WHERE bssid IS NOT NULL
        GROUP BY bssid
        HAVING MIN(received_at) > ?
        ORDER BY first_seen DESC
        LIMIT 100
    """, (cutoff_24h,))

    # ── Regulars: seen on 2+ distinct calendar days ───────────────────────────
    reg_mac_rows = query("""
        SELECT src_mac,
               COUNT(DISTINCT date(received_at)) AS days_seen,
               MIN(received_at)                  AS first_seen,
               MAX(received_at)                  AS last_seen,
               MAX(rssi)                         AS best_rssi,
               COUNT(*)                          AS times_seen,
               COUNT(DISTINCT node_id)           AS node_count
        FROM probe_events
        WHERE src_mac IS NOT NULL
        GROUP BY src_mac
        HAVING COUNT(DISTINCT date(received_at)) >= 2
        ORDER BY days_seen DESC, times_seen DESC
        LIMIT 50
    """)
    reg_macs = []
    for r in reg_mac_rows:
        vendor, randomized = oui_lookup(r["src_mac"])
        r["vendor"]     = vendor
        r["randomized"] = randomized
        reg_macs.append(r)

    reg_networks = query("""
        SELECT bssid,
               MAX(ssid)                         AS ssid,
               COUNT(DISTINCT date(received_at)) AS days_seen,
               MIN(received_at)                  AS first_seen,
               MAX(received_at)                  AS last_seen,
               MAX(rssi)                         AS best_rssi,
               MAX(channel)                      AS channel,
               MAX(encryption)                   AS encryption,
               COUNT(DISTINCT node_id)           AS node_count
        FROM beacon_events
        WHERE bssid IS NOT NULL
        GROUP BY bssid
        HAVING COUNT(DISTINCT date(received_at)) >= 2
        ORDER BY days_seen DESC, best_rssi DESC
        LIMIT 50
    """)

    return {
        "present":      present,
        "new_macs":     new_macs,
        "new_networks": new_networks,
        "reg_macs":     reg_macs,
        "reg_networks": reg_networks,
    }


def build_node_detail(node_id: str) -> dict | None:
    rows = query("""
        SELECT
            node_id,
            MIN(received_at) AS first_seen,
            MAX(received_at) AS last_seen,
            COUNT(*)         AS total_events,
            COUNT(DISTINCT ssid)  AS unique_ssids,
            COUNT(DISTINCT bssid) AS unique_bssids,
            AVG(rssi)        AS avg_rssi,
            MIN(rssi)        AS min_rssi,
            MAX(rssi)        AS max_rssi
        FROM beacon_events
        WHERE node_id = ?
        GROUP BY node_id
    """, (node_id,))
    if not rows:
        return None
    detail = rows[0]
    hb_rows = query("""
        SELECT received_at AS last_heartbeat, uptime_ms, free_heap, wifi_rssi
        FROM heartbeat_events
        WHERE node_id = ?
        ORDER BY id DESC LIMIT 1
    """, (node_id,))
    hb = hb_rows[0] if hb_rows else {}
    detail["last_heartbeat"] = hb.get("last_heartbeat")
    detail["uptime_ms"]      = hb.get("uptime_ms")
    detail["free_heap"]      = hb.get("free_heap")
    detail["wifi_rssi"]      = hb.get("wifi_rssi")
    if detail["last_heartbeat"]:
        detail["status"] = _node_status(detail["last_heartbeat"], threshold=30)
    else:
        detail["status"] = _node_status(detail["last_seen"], threshold=60)
    detail["avg_rssi"] = round(detail["avg_rssi"], 1) if detail["avg_rssi"] else None
    detail["events"]   = build_events(limit=EVENT_CAP, node_id=node_id)
    return detail

# ─── Routes ──────────────────────────────────────────────────────────────────

def ensure_schema():
    """Create tables if they don't exist. Runs once at startup."""
    with get_conn() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS beacon_events (
                id INTEGER PRIMARY KEY AUTOINCREMENT, received_at TEXT NOT NULL,
                node_id TEXT, node_ts INTEGER, ssid TEXT, bssid TEXT,
                rssi INTEGER, channel INTEGER, encryption TEXT, importance TEXT, confidence TEXT
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS probe_events (
                id INTEGER PRIMARY KEY AUTOINCREMENT, received_at TEXT NOT NULL,
                node_id TEXT, node_ts INTEGER, src_mac TEXT, ssid TEXT,
                rssi INTEGER, importance TEXT
            )
        """)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS heartbeat_events (
                id INTEGER PRIMARY KEY AUTOINCREMENT, received_at TEXT NOT NULL,
                node_id TEXT, node_ts INTEGER,
                uptime_ms INTEGER, free_heap INTEGER, wifi_rssi INTEGER
            )
        """)
        conn.commit()


app = FastAPI(title="ESP32 Recon Dashboard")
ensure_schema()


@app.get("/", response_class=HTMLResponse)
async def index():
    return FileResponse(HTML_PATH)


@app.get("/static/dashboard.css")
async def serve_css():
    return FileResponse(CSS_PATH, media_type="text/css")


@app.get("/static/dashboard.js")
async def serve_js():
    return FileResponse(JS_PATH, media_type="application/javascript")


@app.get("/api/nodes")
async def api_nodes():
    return build_nodes()


@app.get("/api/events")
async def api_events(limit: int = EVENT_CAP):
    limit = min(limit, EVENT_CAP)
    return build_events(limit=limit)


@app.get("/api/networks")
async def api_networks():
    return build_networks()


@app.get("/api/clients")
async def api_clients():
    return build_clients()


@app.get("/api/cross-node")
async def api_cross_node():
    return build_cross_node()


@app.get("/api/presence")
async def api_presence():
    return build_presence()


@app.get("/api/rssi-history")
async def api_rssi_history(bssid: str, hours: int = 2):
    hours = min(max(hours, 1), 48)
    cutoff = (datetime.datetime.utcnow() - datetime.timedelta(hours=hours)).isoformat(timespec="seconds")
    rows = query("""
        SELECT node_id, received_at, rssi
        FROM beacon_events
        WHERE bssid = ? AND received_at >= ?
        ORDER BY received_at ASC
    """, (bssid, cutoff))
    series: dict = {}
    for r in rows:
        nid = r["node_id"]
        if nid not in series:
            series[nid] = []
        series[nid].append({"t": r["received_at"], "rssi": r["rssi"]})
    return series


@app.get("/api/heartbeats")
async def api_heartbeats():
    """Latest heartbeat per node."""
    return query("""
        SELECT node_id, received_at, uptime_ms, free_heap, wifi_rssi
        FROM heartbeat_events
        WHERE id IN (SELECT MAX(id) FROM heartbeat_events GROUP BY node_id)
    """)


@app.get("/api/nodes/{node_id}")
async def api_node_detail(node_id: str):
    detail = build_node_detail(node_id)
    if detail is None:
        from fastapi import HTTPException
        raise HTTPException(status_code=404, detail="Node not found")
    return detail


@app.get("/api/stream")
async def api_stream():
    """SSE stream — polls DB every 2 s, pushes new events to the browser."""
    async def event_generator():
        try:
            last_id = 0
            # Start from the current latest event so we don't replay history on connect
            rows = query("SELECT MAX(id) AS max_id FROM beacon_events")
            if rows and rows[0]["max_id"]:
                last_id = rows[0]["max_id"]

            while True:
                await asyncio.sleep(2)
                new_rows = query("""
                    SELECT * FROM beacon_events
                    WHERE id > ?
                    ORDER BY id ASC
                """, (last_id,))
                for row in new_rows:
                    last_id = row["id"]
                    payload = json.dumps(row)
                    yield f"data: {payload}\n\n"
        except asyncio.CancelledError:
            pass  # client disconnected or server shutting down — expected, not an error

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


# ─── Entry point ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("dashboard:app", host="0.0.0.0", port=8080, reload=False)