"""HA Device Cleanup — remove devices and their entities from .storage/ JSON files.

Run inside the HA container:
    docker cp /tmp/ha_cleanup_devices.py homeassistant:/tmp/
    docker exec homeassistant python3 /tmp/ha_cleanup_devices.py <device_id> [device_id ...]

Or edit the DEVICE_IDS constant below and run without args.
"""

import json
import sys

DEVICE_IDS = []

if len(sys.argv) > 1:
    DEVICE_IDS = sys.argv[1:]
elif not DEVICE_IDS:
    print("Usage: python3 ha_cleanup_devices.py <device_id> [device_id ...]")
    print("Or edit DEVICE_IDS in the script and run without args.")
    sys.exit(1)

print("Removing devices:", DEVICE_IDS)

# 1. Remove from device_registry
with open('/config/.storage/core.device_registry', 'r') as f:
    dev_data = json.load(f)

before = len(dev_data['data']['devices'])
dev_data['data']['devices'] = [
    d for d in dev_data['data']['devices']
    if d['id'] not in DEVICE_IDS
]
after = len(dev_data['data']['devices'])

with open('/config/.storage/core.device_registry', 'w') as f:
    json.dump(dev_data, f, indent=2)

print('Devices: {} -> {} (removed {})'.format(before, after, before - after))

# 2. Remove their entities from entity_registry
with open('/config/.storage/core.entity_registry', 'r') as f:
    ent_data = json.load(f)

before = len(ent_data['data']['entities'])
ent_data['data']['entities'] = [
    e for e in ent_data['data']['entities']
    if e.get('device_id') not in DEVICE_IDS
]
after = len(ent_data['data']['entities'])

with open('/config/.storage/core.entity_registry', 'w') as f:
    json.dump(ent_data, f, indent=2)

print('Entities: {} -> {} (removed {})'.format(before, after, before - after))
print('Done. No HA restart needed.')