import network, time, hashlib, ujson as json
from umqtt.simple import MQTTClient

WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""
MQTT_BROKER = "test.mosquitto.org"
MQTT_PORT = 1883
NODE_ID = "NODE_A"
ROUTE = "NODE_A->NODE_B->NODE_C->GATEWAY"

TOPIC_ANNOUNCE = b"iot/cncp/routing/announce"
TOPIC_STATUS = b"iot/cncp/nodes/status"

def sha256b(data):
    return hashlib.sha256(data).digest()

def to_hex(b):
    return "".join("{:02x}".format(x) for x in b)

def filtered_body(d, exact_exclude=(), prefix_exclude=()):
    out = {}
    for k in d:
        if k in exact_exclude:
            continue
        skip = False
        for p in prefix_exclude:
            if k.startswith(p):
                skip = True
                break
        if not skip:
            out[k] = d[k]
    return out

def canonical_json(d):
    parts = []
    for k in sorted(d.keys()):
        parts.append('"%s":%s' % (k, json.dumps(d[k])))
    return "{" + ",".join(parts) + "}"

class Crypto:
    def __init__(self, node_id):
        self.priv = sha256b((node_id + "_PRIVATE_SEED_BLOCKCHAIN").encode())
        x = sha256b(self.priv)
        self.pub = bytes([0x04]) + x + sha256b(self.priv + x)

    def sign(self, d, exact_exclude=(), prefix_exclude=()):
        body = canonical_json(filtered_body(d, exact_exclude, prefix_exclude)).encode()
        return to_hex(sha256b(sha256b(self.pub[1:33]) + body))

    def pub_hex(self):
        return to_hex(self.pub)

def connect_wifi():
    w = network.WLAN(network.STA_IF)
    w.active(True)
    w.connect(WIFI_SSID, WIFI_PASSWORD)
    print("[NODE_A] Connecting to WiFi...")
    for _ in range(20):
        if w.isconnected():
            print("[NODE_A] WiFi connected:", w.ifconfig()[0])
            return True
        time.sleep(0.5)
    print("[NODE_A] WiFi failed")
    return False

def connect_mqtt(client):
    while True:
        try:
            client.connect()
            print("[NODE_A] MQTT connected")
            return
        except Exception as e:
            print("[NODE_A] MQTT retry:", e)
            time.sleep(2)

crypto = Crypto(NODE_ID)
print("[NODE_A] Keypair:", crypto.pub_hex()[:16] + "...")

if not connect_wifi():
    raise Exception("WiFi failed")

client = MQTTClient(NODE_ID + "_client", MQTT_BROKER, MQTT_PORT, keepalive=60)
connect_mqtt(client)

msg_count = 0

while True:
    msg_count += 1

    temp = 22.0 + (msg_count % 10) * 0.5
    humidity = 55.0 + (msg_count % 8)

    msg = {
        "node_id": NODE_ID,
        "route": ROUTE,
        "hops": 3,
        "hop_index": 0,
        "timestamp": time.ticks_ms(),
        "route_id": "{}_{}_{}".format(NODE_ID, time.ticks_ms(), msg_count),
        "temp": temp,
        "humidity": humidity,
        "version": 1,
    }

    msg["signature"] = crypto.sign(msg, exact_exclude=("signature", "public_key"))
    msg["public_key"] = crypto.pub_hex()

    try:
        client.publish(TOPIC_ANNOUNCE, json.dumps(msg))
        print("[NODE_A] Sent:", msg["route_id"], "temp=", temp)
    except Exception as e:
        print("[NODE_A] Publish error:", e)
        connect_mqtt(client)

    status = json.dumps({
        "node_id": NODE_ID,
        "status": "active",
        "timestamp": time.ticks_ms()
    })

    try:
        client.publish(TOPIC_STATUS, status)
    except Exception:
        connect_mqtt(client)

    time.sleep(15)