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

WIFI_SSID = "Wokwi-GUEST"
WIFI_PASSWORD = ""

MQTT_BROKER = "broker.hivemq.com"
MQTT_PORT = 1883
MQTT_KEEPALIVE = 60
NODE_ID = "GATEWAY"

TOPIC_FORWARD = b"iot/cncp/routing/forward"
TOPIC_VERIFIED = b"iot/cncp/routing/verified"
TOPIC_ALERT = b"iot/cncp/security/alert"
TOPIC_STATUS = b"iot/cncp/nodes/status"

STATUS_INTERVAL_SEC = 30
PING_EVERY_LOOPS = 80

processed = set()
stats = {
    "routes_received": 0,
    "routes_verified": 0,
    "routes_rejected": 0,
    "attacks_detected": 0,
}

client = None
last_status_at = 0


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


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


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


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


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

    def verify(self, message, signature, pub_hex, exact_exclude=(), prefix_exclude=()):
        try:
            pub = bytes(int(pub_hex[i:i + 2], 16) for i in range(0, len(pub_hex), 2))
            verify_key = sha256b(pub[1:33])
            body = canonical_json(filtered_body(message, exact_exclude, prefix_exclude)).encode()
            return to_hex(sha256b(verify_key + body)) == signature
        except Exception:
            return False

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


crypto = Crypto(NODE_ID)


def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(WIFI_SSID, WIFI_PASSWORD)
    print("[GW] Connecting WiFi...")
    for _ in range(30):
        if wlan.isconnected():
            print("[GW] WiFi:", wlan.ifconfig()[0])
            return True
        time.sleep(0.5)
    return False


def build_client():
    client_id = "{}_{}".format(NODE_ID, int(time.time()))
    mqtt_client = MQTTClient(client_id, MQTT_BROKER, MQTT_PORT, keepalive=MQTT_KEEPALIVE)
    mqtt_client.set_callback(on_message)
    return mqtt_client


def connect_mqtt():
    while True:
        mqtt_client = build_client()
        try:
            mqtt_client.connect()
            mqtt_client.subscribe(TOPIC_FORWARD)
            print("[GW] MQTT connected")
            return mqtt_client
        except Exception as exc:
            print("[GW] MQTT retry:", exc)
            time.sleep(2)


def publish_status():
    global last_status_at
    now = int(time.time())
    if now - last_status_at < STATUS_INTERVAL_SEC:
        return
    payload = {
        "node_id": NODE_ID,
        "status": "active",
        "timestamp": now,
        "public_key": crypto.pub_hex(),
        "routes_received": stats["routes_received"],
        "routes_verified": stats["routes_verified"],
        "routes_rejected": stats["routes_rejected"],
        "attacks_detected": stats["attacks_detected"],
    }
    client.publish(TOPIC_STATUS, json.dumps(payload))
    last_status_at = now


def publish_verified(message, route_id):
    payload = {
        "route_id": route_id,
        "node_id": message.get("node_id", "NODE_A"),
        "route": message.get("route", "NODE_A->NODE_B->NODE_C->GATEWAY"),
        "nonce": message.get("nonce", route_id),
        "node_a_sig": message.get("signature", ""),
        "node_b_sig": message.get("relay_1_signature", ""),
        "node_c_sig": message.get("relay_2_signature", ""),
        "temp": message.get("temp", 0),
        "humidity": message.get("humidity", 0),
        "timestamp": time.time(),
        "status": "verified",
    }
    client.publish(TOPIC_VERIFIED, json.dumps(payload))
    stats["routes_verified"] += 1
    print("[GW] VERIFIED + Sent to MQTT:", route_id)


def publish_mitm_alert(route_id, offender):
    payload = {
        "attack_type": "MITM",
        "attacker_id": offender,
        "evidence": "signature_failed",
        "route_id": route_id,
        "timestamp": time.time(),
    }
    client.publish(TOPIC_ALERT, json.dumps(payload))
    stats["attacks_detected"] += 1


def on_message(topic, payload):
    try:
        msg = json.loads(payload)
        route_id = msg.get("route_id", "")

        if msg.get("relay_2_node") != "NODE_C":
            return
        if "relay_2_signature" not in msg:
            return

        stats["routes_received"] += 1
        if route_id in processed:
            return
        processed.add(route_id)

        print("\n[GW] Verifying", route_id)

        ok_a = crypto.verify(
            msg,
            msg.get("signature", ""),
            msg.get("public_key", ""),
            exact_exclude=("signature", "public_key"),
            prefix_exclude=("relay_",),
        )
        ok_b = crypto.verify(
            msg,
            msg.get("relay_1_signature", ""),
            msg.get("relay_1_public_key", ""),
            exact_exclude=("relay_1_signature", "relay_1_public_key"),
            prefix_exclude=("relay_2_",),
        )
        ok_c = crypto.verify(
            msg,
            msg.get("relay_2_signature", ""),
            msg.get("relay_2_public_key", ""),
            exact_exclude=("relay_2_signature", "relay_2_public_key"),
        )

        print("[GW] A={} B={} C={}".format(ok_a, ok_b, ok_c))

        if ok_a and ok_b and ok_c:
            print("[GW] ALL VALID")
            publish_verified(msg, route_id)
        else:
            stats["routes_rejected"] += 1
            print("[GW] SIGNATURE FAIL")
            if not ok_a:
                offender = msg.get("node_id", "UNKNOWN")
            elif not ok_b:
                offender = msg.get("relay_1_node", "UNKNOWN")
            else:
                offender = msg.get("relay_2_node", "UNKNOWN")
            publish_mitm_alert(route_id, offender)
    except Exception as exc:
        print("[GW] Error:", exc)


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

client = connect_mqtt()
print("[GW] Listening...")

ping_counter = 0
while True:
    try:
        client.check_msg()
        publish_status()
        ping_counter += 1
        if ping_counter >= PING_EVERY_LOOPS:
            client.ping()
            ping_counter = 0
    except Exception as exc:
        print("[GW] MQTT Error:", exc)
        client = connect_mqtt()
        ping_counter = 0
    time.sleep(0.5)