#!/usr/bin/env python3
"""
ha-nut-setup.py — Add the Home Assistant NUT integration via REST config-flow API.

Stdlib-only (no pip). Idempotent: skips if a NUT config entry already exists.

Env:
  HA_HOST   (default pfv-bms.knel.net)
  HA_PORT   (default 8123)
  HA_TOKEN  (long-lived access token)
  NUT_HOST  (default 100.121.189.98)
  NUT_PORT  (default 3493)
  NUT_USER  (default homeassistant)
  NUT_PASS  (required)
  NUT_UPS   (default apc-smartups-c1500)
"""
import os, json, sys, time, urllib.request, urllib.error

HA_HOST  = os.environ.get("HA_HOST", "pfv-bms.knel.net")
HA_PORT  = int(os.environ.get("HA_PORT", "8123"))
TOKEN    = os.environ["HA_TOKEN"]
NUT_HOST = os.environ.get("NUT_HOST", "192.168.3.11")
NUT_PORT = int(os.environ.get("NUT_PORT", "3493"))
NUT_USER = os.environ.get("NUT_USER", "homeassistant")
NUT_PASS = os.environ["NUT_PASS"]
NUT_UPS  = os.environ.get("NUT_UPS", "apc-smartups-c1500")
BASE     = f"http://{HA_HOST}:{HA_PORT}"

def api(method, path, data=None):
    body = json.dumps(data).encode() if data else None
    req = urllib.request.Request(
        f"{BASE}/api{path}", data=body, method=method,
        headers={"Authorization": f"Bearer {TOKEN}",
                 "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=20) as r:
            return json.loads(r.read())
    except urllib.error.HTTPError as e:
        raw = e.read().decode()
        try:
            return json.loads(raw)
        except Exception:
            return {"_http_error": e.code, "_raw": raw[:300]}
    except Exception as e:
        return {"_error": str(e)}

# ── verify token ──
cfg = api("GET", "/config")
if "_http_error" in cfg or "_error" in cfg:
    print(f"Cannot reach HA or token invalid: {cfg}"); sys.exit(1)
print(f"HA {cfg.get('version')} — token valid")

# ── check existing entries (idempotent) ──
entries = api("GET", "/config/config_entries/entry")
existing = [e for e in entries if e.get("domain") == "nut"]
if existing:
    for e in existing:
        print(f"NUT already configured: {e.get('title')} "
              f"(data={json.dumps(e.get('data', {}))})")
    print("Skipping — delete it in HA UI first if you want to re-run.")
    sys.exit(0)
print("No existing NUT entry. Starting config flow.")

# ── initiate flow ──
flow = api("POST", "/config/config_entries/flow", {"handler": "nut"})
if "flow_id" not in flow:
    print(f"Flow init failed: {json.dumps(flow)}"); sys.exit(1)
fid = flow["flow_id"]
print(f"Flow started: step={flow.get('step_id')} "
      f"fields={[f.get('name') for f in flow.get('data_schema', [])]}")

# ── submit connection details ──
creds = {"host": NUT_HOST, "port": NUT_PORT,
         "username": NUT_USER, "password": NUT_PASS}
flow = api("POST", f"/config/config_entries/flow/{fid}", creds)
if flow.get("errors"):
    print(f"Validation errors: {flow['errors']}"); sys.exit(1)
print(f"After submit: type={flow.get('type')} step={flow.get('step_id')}")

# ── handle follow-up steps (UPS selection etc.) ──
while flow.get("type") == "form":
    step = flow.get("step_id", "?")
    schema = flow.get("data_schema", [])
    print(f"Step '{step}': fields={[f.get('name') for f in schema]}")
    for f in schema:
        opts = f.get("options") or f.get("values")
        if opts:
            print(f"  {f.get('name')} options: {opts}")
    submission = {}
    for f in schema:
        nm = f.get("name")
        ftype = f.get("type", "")
        if ftype == "multi_select":
            opts = f.get("options", [])
            vals = [o[0] if isinstance(o, list) else o for o in opts]
            submission[nm] = [NUT_UPS] if NUT_UPS in vals else vals[:1]
        elif nm in creds:
            submission[nm] = creds[nm]
        elif "default" in f:
            submission[nm] = f["default"]
        elif ftype == "select":
            opts = f.get("options", [])
            vals = [o[0] if isinstance(o, list) else o for o in opts]
            submission[nm] = NUT_UPS if NUT_UPS in vals else (vals[0] if vals else "")
    fid = flow.get("flow_id", fid)
    flow = api("POST", f"/config/config_entries/flow/{fid}", submission)
    if flow.get("errors"):
        print(f"Validation errors: {flow['errors']}"); sys.exit(1)
    print(f"  -> type={flow.get('type')} step={flow.get('step_id')}")

# ── result ──
if flow.get("type") == "create_entry":
    print(f"\nNUT integration created: {flow.get('title')}")
elif flow.get("type") == "abort":
    print(f"\nFlow aborted: {flow.get('reason')}"); sys.exit(1)
else:
    print(f"\nFinal state: {flow.get('type')} — {json.dumps(flow)[:200]}")

# ── verify sensors ──
print("\nWaiting 10s for entities ...")
time.sleep(10)
states = api("GET", "/states")
ups = [s for s in states
       if "apc_smartups" in s["entity_id"].lower()
       or "sensor.ups_" in s["entity_id"].lower()]
if ups:
    print(f"Found {len(ups)} UPS sensors:")
    for e in sorted(ups, key=lambda x: x["entity_id"]):
        st = e.get("state", "?")
        unit = e.get("attributes", {}).get("unit_of_measurement", "")
        name = e.get("attributes", {}).get("friendly_name", "")
        print(f"  {e['entity_id']:55s} {st:>8} {unit:4s} {name}")
else:
    print("No UPS sensors yet (may still be initialising — check HA UI).")