#!/usr/bin/env python3
"""Auto-rejoin helper: open permit_join, watch for device_joined events,
and rename any rejoining device back to its original friendly name.

Why: after the zigbee2mqtt database is wiped, devices need to rejoin. They
already know the network (network_key + pan_id + channel preserved on
RPi), so they will rejoin automatically when woken. This script:

1. Publishes permit_join for `time` seconds (default 120).
2. Subscribes to bridge/event on the local MQTT broker.
3. For each device_joined event with a known IEEE, sends a
   bridge/request/device/rename call to restore the friendly name.

Run from anywhere with network access to the RPi's mosquitto broker:

    uv run --with paho-mqtt scripts/z2m_repair.py --time 120

Mapping is hardcoded from the backup `state.json` / `configuration.yaml`
preserved from the previous working stack on the RPi. Override with
--mapping to add or change entries.
"""
import argparse
import json
import sys
import time

import paho.mqtt.client as mqtt

# IEEE -> friendly name, from the last known good state.
DEFAULT_MAPPING = {
    "0xffffb40e0604fbd0": "Squiggle",
    "0xffffb40e06065536": "Door Button",
    "0xffffb40e0603c9e4": "Sideboard Lamp",
    "0xffffb40e06050af6": "Bamboo Lights",
    "0xffffb40e0603ef11": "DJ Booth",
}


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="192.168.81.147")
    p.add_argument("--port", type=int, default=1883)
    p.add_argument("--time", type=int, default=120,
                   help="permit_join window in seconds (default 120)")
    p.add_argument("--mapping", help="JSON file overriding default mapping")
    args = p.parse_args()

    mapping = dict(DEFAULT_MAPPING)
    if args.mapping:
        with open(args.mapping) as f:
            override = json.load(f)
        mapping.update(override)

    print(f"[repair] mapping {len(mapping)} known devices:")
    for ieee, name in mapping.items():
        print(f"  {ieee}  ->  {name}")

    joined: set[str] = set()
    renamed: set[str] = set()

    def on_connect(client, userdata, flags, reason, properties=None):
        print(f"[repair] connected to mqtt://{args.host}:{args.port}")
        client.subscribe("zigbee2mqtt/bridge/event")
        client.publish(
            "zigbee2mqtt/bridge/request/permit_join",
            json.dumps({"time": args.time, "transaction": "repair-1"}),
        )
        print(f"[repair] permit_join opened for {args.time}s")
        print(f"[repair] waiting for device_joined events...")
        print(f"[repair] press Ctrl-C to abort")

    def on_message(client, userdata, msg):
        try:
            payload = json.loads(msg.payload)
        except (ValueError, TypeError):
            return
        if payload.get("type") != "device_joined":
            return
        ieee = payload.get("data", {}).get("ieee_address")
        if not ieee or ieee in joined:
            return
        joined.add(ieee)
        name = mapping.get(ieee, f"<unknown-{ieee[-6:]}>")
        if name.startswith("<unknown-"):
            print(f"[repair] joined (NEW): {ieee} -> {name}")
            print(f"[repair]           (will not auto-rename; use: just z2m-rename {ieee} <name>)")
            return
        if ieee in renamed:
            return
        # Ask zigbee2mqtt to rename
        client.publish(
            "zigbee2mqtt/bridge/request/device/rename",
            json.dumps({"from": ieee, "to": name, "transaction": f"rename-{ieee}"}),
        )
        renamed.add(ieee)
        print(f"[repair] joined + renamed: {ieee} -> {name}")
        if len(renamed) >= len(mapping):
            print(f"[repair] all {len(mapping)} known devices rejoined and renamed")

    client = mqtt.Client(callback_api_version=mqtt.CallbackAPIVersion.VERSION2)
    client.on_connect = on_connect
    client.on_message = on_message

    try:
        client.connect(args.host, args.port, keepalive=60)
        client.loop_start()
        time.sleep(args.time + 5)
        print(f"[repair] permit_join window elapsed ({args.time}s)")
        print(f"[repair] joined: {len(joined)}, renamed: {len(renamed)} of {len(mapping)}")
        missing = set(mapping) - joined
        if missing:
            print(f"[repair] missing (need re-pair):")
            for ieee in missing:
                print(f"  {ieee}  ({mapping[ieee]})")
        client.loop_stop()
        client.disconnect()
    except KeyboardInterrupt:
        print(f"\n[repair] aborted")
        client.loop_stop()
        client.disconnect()
        return 130

    return 0 if len(renamed) == len(mapping) else 1


if __name__ == "__main__":
    sys.exit(main())