#!/usr/bin/env python3
# rf-analyze.py -- run the unifi-wifi airtime-reduction analyses against the UniFi cloud
# CONNECTOR (api.ui.com) for NON-UOS consoles. Same logic as the Mongo-fed model-rank.sh /
# optimize-radios.sh, but the data comes from the console's LOCAL Network API proxied through
# the cloud connector (see references/site-manager-api.md):
#   - per-AP/band airtime history: POST /stat/report/hourly.ap (SAME schema as ace_stat.stat_hourly)
#   - identity/zone:               GET  /stat/device  (type=uap, name -> floor/zone)
#   - roam graph (optional):       POST /stat/event   (often empty on small/stationary sites ->
#                                  degrades to airtime-pressure ranking / power-down-only planning)
#
# Usage (normally invoked by model-rank.sh/optimize-radios.sh with --console):
#   UNIFI_SM_KEY=<key> rf-analyze.py rank|optimize --console "<name|hostId>" [--days N] [--band ng|na|6e] [--site default]
#   env (optimize): ROAM_MIN(4) CAP(85) ZONE_DISABLE_PCT(40) REDUN_NG(2) REDUN_OTHER(1)
import json, os, sys, math, re, urllib.request, urllib.error

KEY  = os.environ.get('UNIFI_SM_KEY', '')
BASE = 'https://api.ui.com'

def api(path, method='GET', body=None, quiet=False):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(BASE + path, data=data, method=method,
          headers={'X-API-KEY': KEY, 'Accept': 'application/json', 'Content-Type': 'application/json'})
    try:
        return json.loads(urllib.request.urlopen(req, timeout=90).read())
    except urllib.error.HTTPError as e:
        if not quiet: sys.stderr.write('[rf-analyze] %s %s -> HTTP %s\n' % (method, path, e.code))
        return None
    except Exception as e:
        if not quiet: sys.stderr.write('[rf-analyze] %s %s -> %s\n' % (method, path, e))
        return None

def die(m): sys.stderr.write('[ERROR] %s\n' % m); sys.exit(1)

# ---- args ----
if len(sys.argv) < 2 or sys.argv[1] not in ('rank', 'optimize'):
    die('usage: rf-analyze.py rank|optimize --console <name> [--days N] [--band ng|na|6e] [--site default]')
MODE = sys.argv[1]
opt = {'--console': None, '--days': '7' if MODE == 'rank' else '14', '--band': 'ng', '--site': 'default'}
a = sys.argv[2:]; i = 0
while i < len(a):
    if a[i] in opt: opt[a[i]] = a[i+1]; i += 2
    else: i += 1
CONSOLE = opt['--console'] or die('--console required')
DAYS = int(opt['--days']); BAND = opt['--band']; SITE = opt['--site']
if BAND not in ('ng', 'na', '6e'): die('band must be ng|na|6e')
if not KEY: die('no UNIFI_SM_KEY (the .sh wrapper resolves it from vault services/unifi-site-manager)')

# ---- resolve console -> hostId ----
hosts = (api('/v1/hosts') or {}).get('data', [])
HID = CNAME = None
for h in hosts:
    nm = ((h.get('reportedState') or {}).get('name')) or ''
    if h.get('id', '').lower() == CONSOLE.lower() or CONSOLE.lower() in nm.lower():
        HID, CNAME = h['id'], nm; break
if not HID: die("no console matching '%s' in the Site Manager fleet" % CONSOLE)
CB = '/v1/connector/consoles/%s/proxy/network/api/s/%s' % (HID, SITE)

# ---- identity + zone from /stat/device ----
dev = (api(CB + '/stat/device') or {}).get('data', [])
name, zone = {}, {}
for d in dev:
    if d.get('type') != 'uap': continue
    n = d.get('name') or d.get('mac'); name[d['mac']] = n
    fz = re.search(r'(\d)(?:st|nd|rd|th)\s*floor', str(n), re.I)
    rm = re.search(r'\b(\d)\d{2}\b', str(n))
    zone[d['mac']] = ('Floor ' + fz.group(1)) if fz else (('Floor ' + rm.group(1)) if rm else 'misc')
if not name: die('no APs (type=uap) on console %s site %s' % (CNAME, SITE))

# ---- airtime history from /stat/report/hourly.ap (same schema as ace_stat.stat_hourly) ----
import time as _t
end = int(_t.time() * 1000); start = end - DAYS * 86400000
attrs = ['time', 'ap'] + ['%s-%s' % (BAND, f) for f in
         ('cu_total', 'cu_self_rx', 'cu_self_tx', 'cu_interf', 'num_sta', 'tx_retries', 'wifi_tx_attempts', 'satisfaction')]
# IMPORTANT: pass macs[] explicitly -- the UniFi report endpoint returns only a small default
# subset of APs otherwise (Cascades returned 10 of 77 until the MAC list was supplied).
rep = api(CB + '/stat/report/hourly.ap', 'POST', {'attrs': attrs, 'start': start, 'end': end, 'macs': list(name.keys())})
rows = (rep or {}).get('data', [])
prof = {}
for d in rows:
    ap = d.get('ap')
    if not ap or ap not in name: continue
    cu = d.get(BAND + '-cu_total');
    if cu is None: continue
    self_ = (d.get(BAND + '-cu_self_rx') or 0) + (d.get(BAND + '-cu_self_tx') or 0)
    intf = d.get(BAND + '-cu_interf')
    if intf is None: intf = max(0.0, (cu or 0) - self_)        # derive when not reported
    sta = d.get(BAND + '-num_sta') or 0
    retr = d.get(BAND + '-tx_retries') or 0; att = d.get(BAND + '-wifi_tx_attempts') or 0
    sat = d.get(BAND + '-satisfaction')
    p = prof.setdefault(ap, {'cu': 0, 'intf': 0, 'self': 0, 'sta': 0, 'staPk': 0, 'retr': 0, 'att': 0, 'sat': 0, 'satN': 0, 'n': 0})
    p['cu'] += cu or 0; p['intf'] += intf; p['self'] += self_; p['sta'] += sta; p['staPk'] = max(p['staPk'], sta)
    p['retr'] += retr; p['att'] += att
    if isinstance(sat, (int, float)): p['sat'] += sat; p['satN'] += 1
    p['n'] += 1
for k, p in prof.items():
    for f in ('cu', 'intf', 'self', 'sta'): p[f] /= p['n']
    p['retrPct'] = min(100.0, 100.0 * p['retr'] / p['att']) if p['att'] > 0 else 0.0
    p['sat'] = p['sat'] / p['satN'] if p['satN'] else 100.0

# ---- roam graph (optional; often empty on small/stationary sites) ----
roam = {}; dir_ = {}; rss = {}
ev = api(CB + '/stat/event', 'POST', {'_limit': 3000, 'within': DAYS * 24}, quiet=True)
for e in ((ev or {}).get('data', []) if ev else []):
    if 'Roam' not in str(e.get('key', '')): continue
    A = e.get('ap_from') or e.get('ap'); B = e.get('ap_to') or e.get('ap')
    if A: roam[A] = roam.get(A, 0) + 1
    if B: roam[B] = roam.get(B, 0) + 1
    if A and B and A != B: dir_[A + '>' + B] = dir_.get(A + '>' + B, 0) + 1
HAVE_ROAM = bool(roam)

bn = {'ng': '2.4', 'na': '5', '6e': '6'}[BAND]
print('[INFO] console=%s site=%s band=%s(%sGHz) window=%dd  APs profiled=%d  roam-graph=%s  (source: cloud connector)'
      % (CNAME, SITE, BAND, bn, DAYS, len(prof), 'yes' if HAVE_ROAM else 'EMPTY -> airtime-only'))
if not prof: die('no airtime history for band %s (try another band, or the APs may be idle)' % BAND)

# ============================== RANK (model-rank port) ==============================
if MODE == 'rank':
    rows_out = []
    for ap, p in prof.items():
        rm = roam.get(ap, 0)
        # with roam: (cu+intf)*log(1+roam)/(1+sta); without: airtime pressure / load (roam-neutral)
        score = (p['cu'] + p['intf']) * (math.log(1 + rm) if HAVE_ROAM else 1.0) / (1 + p['sta'])
        rows_out.append((ap, name.get(ap, ap), p['cu'], p['intf'], p['sta'], rm, score))
    rows_out.sort(key=lambda r: -r[6])
    print('\n==== band=%s  top airtime-reduction candidates (disable/power-down) ====' % BAND)
    print('  rank  AP                       cu%   interf%  ~clients  roams   score   hint')
    for i, r in enumerate(rows_out[:15]):
        ap, nm, cu, intf, sta, rm, sc = r
        hint = ('DISABLE (redundant, low load)' if (rm > 50 and sta < 3)
                else 'POWER-DOWN (busy)' if (cu + intf > 40) else 'review')
        print('  %d\t%-24s %.0f\t%.0f\t%.1f\t%d\t%.1f\t%s' % (i+1, nm[:24], cu, intf, sta, rm, sc, hint))
    print('  (APs profiled on %s: %d)' % (BAND, len(rows_out)))
    note = 'score = (cu_total + cu_interf) * log(1+roams) / (clients+1)'
    if not HAVE_ROAM:
        note = ('score = (cu_total + cu_interf) / (clients+1)  [ROAM GRAPH EMPTY on this console -- '
                'ranking by airtime pressure only; no redundancy weighting]')
    print('\n[note] v1 heuristic: %s. High = busy+interfered (and clients have somewhere to roam, if roam known) '
          '= safe to shrink/disable. Validate per-zone before applying.' % note)
    sys.exit(0)

# ============================== OPTIMIZE (optimize-radios port) ==============================
ROAM_MIN = int(os.environ.get('ROAM_MIN', '4')); CAP = float(os.environ.get('CAP', '85'))
ZPCT = float(os.environ.get('ZONE_DISABLE_PCT', '40'))
REDUN = int(os.environ.get('REDUN_NG', '2')) if BAND == 'ng' else int(os.environ.get('REDUN_OTHER', '1'))
RSSI_OK = {'ng': -68, 'na': -72, '6e': -75}[BAND]

# strong bidirectional neighbors (from roam dir edges). No roam + no neighbor matrix -> empty -> no disables.
strong = {a: {} for a in prof}
for A in prof:
    for B in prof:
        if A == B: continue
        if min(dir_.get(A + '>' + B, 0), dir_.get(B + '>' + A, 0)) >= ROAM_MIN:
            strong[A][B] = 1
# optional measured SNR matrix (neighbor-collect) via env NEIGHBOR_JSON
nj = os.environ.get('NEIGHBOR_JSON', '')
if nj and os.path.isfile(nj):
    nbr = json.load(open(nj)); MIN = int(os.environ.get('NBR_SNR_MIN', '20'))
    def ms(x, y): o = nbr.get(x, {}); return max([o.get(b, {}).get(y, -1) for b in ('2.4', '5', '6')] + [-1])
    nm2mac = {v: k for k, v in name.items()}
    strong = {a: {} for a in prof}
    aps_n = list(nbr)
    for x in aps_n:
        for y in aps_n:
            if x == y: continue
            if ms(x, y) >= MIN and ms(y, x) >= MIN:
                A, B = nm2mac.get(x), nm2mac.get(y)
                if A in prof and B in prof: strong[A][B] = 1; strong[B][A] = 1

aps = list(prof); active = {a: True for a in aps}; projCu = {a: prof[a]['cu'] for a in aps}
zoneTot = {}; zoneDis = {}
for a in aps: zoneTot[zone[a]] = zoneTot.get(zone[a], 0) + 1
def activeStrong(a): return [n for n in strong.get(a, {}) if active.get(n)]
def absorbers(a): return [n for n in activeStrong(a) if projCu[n] + prof[a]['self'] <= CAP]
def soleNeighborOf(a):
    for x in aps:
        if not active.get(x) or x == a or not strong.get(x, {}).get(a): continue
        if sum(1 for n in strong.get(x, {}) if n != a and active.get(n)) == 0: return True
    return False
disabled = []
go = True
while go:
    go = False; best = None; bestBen = -1
    for a in aps:
        if not active[a]: continue
        p = prof[a]
        if len(activeStrong(a)) < REDUN: continue
        if len(absorbers(a)) < 1: continue
        if soleNeighborOf(a): continue
        if (zoneDis.get(zone[a], 0) + 1) > math.floor(zoneTot[zone[a]] * ZPCT / 100): continue
        ben = p['intf'] + 0.5 * p['retrPct']
        if ben > bestBen: bestBen = ben; best = a
    if best:
        ab = absorbers(best); active[best] = False; zoneDis[zone[best]] = zoneDis.get(zone[best], 0) + 1
        share = prof[best]['self'] / len(ab)
        for n in ab: projCu[n] += share
        disabled.append({'ap': best, 'cover': [name[n] for n in ab]}); go = True

powerdown = []; keep = []
for a in aps:
    if not active[a]: continue
    p = prof[a]
    if p['n'] < 3 or (p['cu'] < 5 and p['intf'] < 5): keep.append((a, 'idle/off')); continue
    if p['cu'] >= 40 or p['intf'] >= 35 or p['retrPct'] >= 15 or p['sat'] < 80: powerdown.append(a)
    else: keep.append((a, 'low-util'))

def fmt(a):
    p = prof[a]
    return '%s  cu=%.0f%% interf=%.0f%% self=%.0f%% clients(avg/pk)=%.1f/%.0f retr=%.0f%% sat=%.0f' % (
        name[a], p['cu'], p['intf'], p['self'], p['sta'], p['staPk'], p['retrPct'], p['sat'])
def byZone(items, kf):
    z = {}
    for x in items: z.setdefault(zone[kf(x)], []).append(x)
    return z

print('[INFO] band=%s rssi>=%d roam>=%d cap=%.0f%% need_neighbors=%d zone_cap=%.0f%%  adjacency=%s'
      % (BAND, RSSI_OK, ROAM_MIN, CAP, REDUN, ZPCT,
         ('neighbor SNR matrix' if nj else ('roam graph' if HAVE_ROAM else 'NONE -> power-down only'))))
print('==== PLAN band=%s  (APs=%d)  POWER-DOWN(first)=%d  DISABLE(after)=%d  KEEP=%d ===='
      % (BAND, len(aps), len(powerdown), len(disabled), len(keep)))
print('Phase A: POWER-DOWN the busy/thrashing radios to ~Low (smaller cells cut mutual cu_interf zone-wide). Do FIRST.')
pz = byZone(powerdown, lambda a: a)
for z in sorted(pz):
    print(' [%s] %d' % (z, len(pz[z])))
    for a in pz[z][:6]: print('   ' + fmt(a))
    if len(pz[z]) > 6: print('   ...(+%d)' % (len(pz[z]) - 6))
print('\nPhase B: re-measure after power-down settles (cu_interf should drop, headroom appears).')
save = sum(prof[d['ap']]['intf'] for d in disabled)
print('\nPhase C: DISABLE only-when-redundant radios (each has >=%d bidirectional good neighbors WITH headroom). '
      'Est. interference airtime removed: %.0f.' % (REDUN, save))
if not disabled:
    why = 'roam graph empty + no neighbor SNR matrix -> no proven overlap' if not HAVE_ROAM and not nj else 'none clear the capacity+coverage bar yet'
    print('   (none -- %s. Power-down is still coverage-safe; for disables, supply NEIGHBOR_JSON from neighbor-collect.)' % why)
dz = byZone(disabled, lambda d: d['ap'])
for z in sorted(dz):
    print(' [%s]' % z)
    for d in dz[z]: print('   ' + fmt(d['ap']) + '  -> covered by: ' + ', '.join(d['cover'][:3]))
print('\nKEEP: %d radios (isolated-essential or already efficient). Apply per ZONE; never >%.0f%% disabled/zone; validate before+after.'
      % (len(keep), ZPCT))