#!/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 eco     <ip...> [--go]    turn device.eco_mode OFF (--on to re-enable)
  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_eco(specs, go, want=False):
    """Set device.eco_mode across the fleet (TP17.3, 2026-08-23).

    eco_mode throttles the CPU clock and TX power. It was applied in TP17 by an
    RSSI rule and turned out to cause real misbehaviour: the s23 phase-cut dimmer
    flickered, and s26's `momentary` input swallowed the first edge. Both were
    fixed by switching it off, so the fleet decision is: OFF everywhere.

    No reboot needed — the device answers `restart_required: false` to this write.
    Devices that don't support it report eco_mode=None and are skipped (s24).
    """
    # Weak links (s20 sits at -85 dBm) time out intermittently. A single timeout
    # must not abort the sweep, so every call is retried and failures are
    # collected and reported at the end instead of raising.
    def _try(ip, method, params=None, tries=3):
        for n in range(tries):
            try:
                return _rpc(ip, method, params, timeout=10.0)
            except Exception as e:
                last = e
        raise last

    failed = []
    for ip in _resolve(specs):
        ident = discover.probe(ip)
        if not ident:
            continue
        name = _name_for(ident["mac"]) or ident.get("name_dev") or "?"
        try:
            cur = _try(ip, "Sys.GetConfig").get("result", {}).get("device", {}).get("eco_mode")
            if cur is None:
                print(f"  {ip:15} {name:26} eco not supported — skipped")
                continue
            if cur == want:
                print(f"  {ip:15} {name:26} already {want}")
                continue
            if not go:
                print(f"  dry: {ip:11} {name:26} eco {cur} -> {want}")
                continue
            res = _try(ip, "Sys.SetConfig", {"config": {"device": {"eco_mode": want}}})
            new = _try(ip, "Sys.GetConfig").get("result", {}).get("device", {}).get("eco_mode")
            rr = (res.get("result") or {}).get("restart_required")
            ok = "OK " if new == want else "MISMATCH"
            print(f"  {ip:15} {name:26} eco {cur} -> {new}  {ok}"
                  + ("  (restart_required was already set before this write)" if rr else ""))
            if new != want:
                failed.append((ip, name, f"read back {new}"))
        except Exception as e:
            print(f"  {ip:15} {name:26} FAILED: {type(e).__name__}")
            failed.append((ip, name, type(e).__name__))
    if failed:
        print(f"\n{len(failed)} device(s) need a retry:")
        for ip, name, why in failed:
            print(f"  {ip:15} {name:26} {why}")


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 == "eco":
        cmd_eco(specs, go, want=("--on" in rest))
    elif cmd == "table":
        cmd_table()
    else:
        print(f"unknown command: {cmd}\n{__doc__}")


if __name__ == "__main__":
    main()