import json, time, urllib.request, urllib.parse

KEY = "e274865463f2d4ca5676f076f86ac133"
GEO_URL = "https://restapi.amap.com/v3/geocode/geo"

def geocode(address):
    params = urllib.parse.urlencode({'key': KEY, 'address': address})
    url = f"{GEO_URL}?{params}"
    try:
        req = urllib.request.Request(url)
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read().decode())
            if data.get('status') == '1' and data.get('geocodes'):
                loc = data['geocodes'][0]['location']
                lng, lat = loc.split(',')
                return [float(lat), float(lng)]
    except Exception as e:
        print(f"ERR: {address[:30]} -> {e}")
    return None

# Load all devices
with open('/tmp/all_devices_raw.json') as f:
    data = json.load(f)

# Collect unique addresses
addr_map = {}
for sys_name, devices in data.items():
    for d in devices:
        addr = (d.get('address') or '').strip()
        addr = ' '.join(addr.split())
        if addr and addr not in addr_map:
            addr_map[addr] = None

print(f"Unique addresses to geocode: {len(addr_map)}")

# Load cache
try:
    with open('/tmp/geocode_cache.json') as f:
        cache = json.load(f)
except:
    cache = {}

new_count = 0
total = len(addr_map)
for i, addr in enumerate(addr_map):
    if addr in cache:
        addr_map[addr] = cache[addr]
        continue
    loc = geocode(addr)
    addr_map[addr] = loc
    cache[addr] = loc
    if loc:
        new_count += 1
        print(f"[{i+1}/{total}] OK  {addr[:50]} -> {loc}")
    else:
        print(f"[{i+1}/{total}] NF  {addr[:50]}")
    if (i + 1) % 30 == 0:
        time.sleep(0.5)
    else:
        time.sleep(0.1)

# Save cache
with open('/tmp/geocode_cache.json', 'w') as f:
    json.dump(cache, f, ensure_ascii=False)

# Annotate devices
for sys_name, devices in data.items():
    for d in devices:
        addr = (d.get('address') or '').strip()
        addr = ' '.join(addr.split())
        d['latlon'] = addr_map.get(addr)

with open('/tmp/all_devices_geo.json', 'w') as f:
    json.dump(data, f, ensure_ascii=False)

cached = total - new_count
failed = sum(1 for v in addr_map.values() if v is None)
print(f"\nDone! Total: {total}, Cached: {cached}, New: {new_count}, Failed: {failed}")