#!/usr/bin/env python3
"""TP17 migration runner — anchored VLAN1→VLAN30 cutover for Shellys.

Per-device sequence (Gen2+ RPC), ordered so recovery is ALWAYS possible:

  1. probe old IP        — must be reachable; record current sta for the log
  2. ANCHOR: ap.enable=true   — device now also broadcasts its own AP
                                (Shelly*-<mac>, open) as an emergency anchor.
                                AP+STA run concurrently, no reboot, sta untouched.
  3. set sta -> dirtyvader-iot / VLAN30 static 10.40.30.<octet>
  4. Shelly.Reboot
  5. poll new IP (10.40.30.<octet>) until it answers, or time out
       -> success: AP stays on until end-of-run `deanchor`
       -> fail:    device is still reachable via its own AP (step 2) — no reset

Only AFTER step 2 do we overwrite sta, so the working fallback that vanishes on
reboot (old sta) is replaced by a fallback that does NOT depend on any WLAN infra
(the device's own AP). That's the whole point of the anchor.

Reuses discover.py (probe/expand/load_inventory) and the same .env as server.py.

Commands:
  python3 migrate.py show    <ip>              read WiFi.GetConfig (read-only)
  python3 migrate.py migrate <ip...> [--go]    anchor+sta+reboot+verify (dry-run w/o --go)
  python3 migrate.py deanchor <ip...> [--go]   turn the AP anchor back off (.30.x IPs)
  python3 migrate.py auth    <ip...> [--go]    set uniform admin password (last step)
  python3 migrate.py table                     final tabular overview (old/new/model/fw)

<ip...> accepts single IPs, ranges (10.40.41.205-233) or 'all' (the default sweep).
Batches run sequentially and STOP on the first failed verify (safety).
"""
import hashlib
import json
import os
import sys
import time
import urllib.error
import urllib.request

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import discover  # noqa: E402

# ---- shared .env loader (same as server.py) ----
_envfile = os.path.join(os.path.dirname(os.path.abspath(__file__)), ".env")
if os.path.exists(_envfile):
    for _line in open(_envfile):
        _line = _line.strip()
        if _line and not _line.startswith("#") and "=" in _line:
            _k, _v = _line.split("=", 1)
            os.environ.setdefault(_k.strip(), _v.strip().strip('"').strip("'"))

IOT_SSID = os.environ.get("SHELLY_IOT_SSID", "dirtyvader-iot")
IOT_PSK = os.environ.get("SHELLY_IOT_PSK", "")
DEVICE_PW = os.environ.get("SHELLY_DEVICE_PW", "")
V30_NET, V30_MASK, V30_GW = "10.40.30", "255.255.255.0", "10.40.30.1"
LOG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "migration-log.jsonl")


def _rpc(ip, method, params=None, timeout=8.0):
    body = json.dumps({"id": 1, "method": method, "params": params or {}}).encode()

    def _mk():
        return urllib.request.Request(f"http://{ip}/rpc", data=body,
                                      headers={"Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(_mk(), timeout=timeout) as r:
            return json.load(r)
    except urllib.error.HTTPError as e:
        if e.code == 401 and DEVICE_PW:
            mgr = urllib.request.HTTPPasswordMgrWithDefaultRealm()
            mgr.add_password(None, f"http://{ip}/rpc", "admin", DEVICE_PW)
            opener = urllib.request.build_opener(urllib.request.HTTPDigestAuthHandler(mgr))
            with opener.open(_mk(), timeout=timeout) as r:
                return json.load(r)
        raise


def _octet(ip):
    return ip.rsplit(".", 1)[-1]


def _name_for(mac):
    m = discover.load_inventory().get((mac or "").lower())
    return m["name"] if m else None


def _log(rec):
    with open(LOG, "a") as f:
        f.write(json.dumps(rec) + "\n")


def _resolve(specs):
    """Expand CLI ip args into a concrete IP list. 'all' -> default VLAN1 sweep."""
    if not specs or specs == ["all"]:
        return discover.expand(os.environ.get("SHELLY_SWEEP", "10.40.41.180-250"))
    ips = []
    for s in specs:
        ips.extend(discover.expand(s))
    return ips


# ---------------- commands ----------------
def cmd_show(ip):
    print(json.dumps(_rpc(ip, "WiFi.GetConfig"), indent=2))


def _anchor(ip):
    """Enable the device's own AP as an emergency anchor. Live, no reboot.

    This firmware refuses an OPEN AP (`-103 'pass' should not be empty`), so the
    anchor is WPA2 secured with the IOT PSK — a value the operator already has,
    and it's removed again by `deanchor`. AP+STA run concurrently; sta untouched.
    """
    res = _rpc(ip, "WiFi.SetConfig",
               {"config": {"ap": {"enable": True, "is_open": False, "pass": IOT_PSK}}})
    if res.get("error"):
        return {"enabled": False, "error": res["error"]}
    ap = _rpc(ip, "WiFi.GetConfig").get("result", {}).get("ap", {})
    return {"enabled": bool(ap.get("enable")), "ap_ssid": ap.get("ssid"),
            "secured": not ap.get("is_open")}


def _migrate_one(ip, go, poll=90):
    ident = discover.probe(ip)
    if not ident:
        return {"ip": ip, "ok": False, "error": "not reachable on old IP"}
    if not IOT_PSK:
        return {"ip": ip, "ok": False, "error": "SHELLY_IOT_PSK not set"}
    cur = _rpc(ip, "WiFi.GetConfig").get("result", {}).get("sta", {})
    new_ip = f"{V30_NET}.{_octet(ip)}"
    name = _name_for(ident["mac"]) or ident.get("name_dev") or "?"
    sta = {"ssid": IOT_SSID, "pass": IOT_PSK, "enable": True, "ipv4mode": "static",
           "ip": new_ip, "netmask": V30_MASK, "gw": V30_GW, "nameserver": V30_GW}
    plan = {"ip": ip, "name": name, "mac": ident["mac"], "model": ident["model"],
            "fw": ident["fw"], "old_ssid": cur.get("ssid"), "new_ip": new_ip,
            "new_ssid": IOT_SSID}
    if not go:
        return {**plan, "ok": False, "dry_run": True,
                "would": "anchor AP -> set sta(static .30) -> reboot -> verify"}

    # 2. anchor
    anc = _anchor(ip)
    if not anc["enabled"]:
        return {**plan, "ok": False, "error": "AP anchor did not enable — aborting before sta write"}
    print(f"    anchor up: AP '{anc['ap_ssid']}' (WPA2, PSK=IOT-PSK)")
    # 3+4. set sta, reboot
    _rpc(ip, "WiFi.SetConfig", {"config": {"sta": sta}})
    _rpc(ip, "Shelly.Reboot")
    print(f"    sta set -> {new_ip}, rebooting; polling…")
    # 5. poll new IP
    t0 = time.time()
    while time.time() - t0 < poll:
        time.sleep(5)
        got = discover.probe(new_ip)
        if got and got["mac"].lower() == ident["mac"].lower():
            st = _rpc(new_ip, "WiFi.GetStatus").get("result", {})
            rec = {**plan, "ok": True, "verified_ssid": st.get("ssid"),
                   "rssi": st.get("rssi"), "secs": round(time.time() - t0)}
            _log(rec)
            return rec
    return {**plan, "ok": False,
            "error": f"did not appear on {new_ip} within {poll}s — "
                     f"reach it via its AP '{anc['ap_ssid']}' (no reset needed)"}


def cmd_migrate(specs, go):
    ips = _resolve(specs)
    print(f"{'GO' if go else 'DRY-RUN'} migrate {len(ips)} device(s) "
          f"-> {IOT_SSID}/VLAN30 (PSK {'set' if IOT_PSK else 'MISSING'})\n")
    ok = 0
    for ip in ips:
        if not discover.probe(ip):
            continue  # skip empty slots in a range
        print(f"[{ip}] …")
        r = _migrate_one(ip, go)
        tag = "OK " if r.get("ok") else ("dry" if r.get("dry_run") else "FAIL")
        print(f"  {tag}: {r.get('name','?')} {ip} -> {r.get('new_ip','?')} "
              f"{r.get('error') or r.get('would') or ('rssi '+str(r.get('rssi')))}")
        if r.get("ok"):
            ok += 1
        elif go and not r.get("dry_run"):
            print("  >> STOP: batch halted on failure (device recoverable via its AP).")
            break
    print(f"\n{ok} migrated OK.")


def cmd_deanchor(specs, go):
    ips = _resolve(specs) if specs != ["all"] else discover.expand("10.40.30.180-250")
    for ip in ips:
        if not discover.probe(ip):
            continue
        if not go:
            print(f"dry: would disable AP on {ip}")
            continue
        _rpc(ip, "WiFi.SetConfig", {"config": {"ap": {"enable": False}}})
        cfg = _rpc(ip, "WiFi.GetConfig").get("result", {}).get("ap", {})
        print(f"{ip}: AP enable -> {cfg.get('enable')}")


def cmd_auth(specs, go):
    if not DEVICE_PW:
        print("SHELLY_DEVICE_PW not set — aborting.")
        return
    for ip in _resolve(specs):
        if not discover.probe(ip):
            continue
        info = _rpc(ip, "Shelly.GetDeviceInfo").get("result", {})
        realm = info.get("id")
        ha1 = hashlib.sha256(f"admin:{realm}:{DEVICE_PW}".encode()).hexdigest()
        if not go:
            print(f"dry: {ip} realm={realm} ha1={ha1[:12]}…")
            continue
        _rpc(ip, "Shelly.SetAuth", {"user": "admin", "realm": realm, "ha1": ha1})
        print(f"{ip}: auth set (realm {realm})")


def cmd_table():
    inv = discover.load_inventory()
    ips = discover.expand("10.40.30.180-250")
    from concurrent.futures import ThreadPoolExecutor
    rows = []
    with ThreadPoolExecutor(max_workers=32) as ex:
        for r in ex.map(discover.probe, ips):
            if r:
                m = inv.get(r["mac"])
                rows.append({"id": (m["shelly_id"] if m else "?"),
                             "name": (m["name"] if m else r.get("name_dev") or "?"),
                             "old": f"10.40.41.{_octet(r['ip'])}", "new": r["ip"],
                             "gen": r["gen"], "model": r["model"], "fw": r["fw"]})
    rows.sort(key=lambda x: x["id"])
    print(f"\n{'id':5} {'name':26} {'old IP':15} {'new IP':15} {'gen':4} {'model':18} fw")
    print("-" * 108)
    for x in rows:
        print(f"{x['id']:5} {x['name']:26} {x['old']:15} {x['new']:15} "
              f"g{x['gen']:<3} {x['model']:18} {x['fw']}")
    print(f"\n{len(rows)} Shellys on VLAN30.")


def main():
    a = sys.argv[1:]
    if not a:
        print(__doc__)
        return
    cmd, rest = a[0], a[1:]
    go = "--go" in rest
    specs = [x for x in rest if not x.startswith("--")]
    if cmd == "show":
        cmd_show(specs[0])
    elif cmd == "migrate":
        cmd_migrate(specs, go)
    elif cmd == "deanchor":
        cmd_deanchor(specs or ["all"], go)
    elif cmd == "auth":
        cmd_auth(specs, go)
    elif cmd == "table":
        cmd_table()
    else:
        print(f"unknown command: {cmd}\n{__doc__}")


if __name__ == "__main__":
    main()