import json
import time
import random
import paho.mqtt.client as mqtt
import re
import os

# MQTT broker settings
broker_address = "localhost"
topic_base = "apu_tb/"  # Base topic, each node will append its ID
config_base = "apu_tb_config/"  # Base topic for configuration

# Define nodes to be part of the script
node_numbers = [0, 1, 2, 3, 4,
                5, 6, 7, 8, 9,
                10, 11, 12, 13, 14,
                15, 16, 17, 18, 19,
                20, 21, 22, 23, 24]  # Nodes 0 to 24 defined separately

# Read and parse node_addresses.sh
script_dir = os.path.dirname(__file__)
node_file = os.path.join(script_dir, "node_addresses.sh")

eth0MAC = {}
phy0MAC = {}
phy1MAC = {}
neighbors = {}

with open(node_file, 'r') as f:
    content = f.read()

# Function to extract associative arrays
def extract_array(name):
    pattern = re.compile(rf"declare -A {name}=\((.*?)\)", re.DOTALL)
    match = pattern.search(content)
    if not match:
        return {}
    entries = match.group(1)
    lines = re.findall(r'\["(.*?)"\]=\"(.*?)\"', entries)
    return dict(lines)

eth0MAC = extract_array("eth0MAC")
phy0MAC = extract_array("phy0MAC")
phy1MAC = extract_array("phy1MAC")

# Special handling for neighbors
neighbors_pattern = re.compile(r'declare -A neighbors=\((.*?)\)', re.DOTALL)
neighbors_match = neighbors_pattern.search(content)
if neighbors_match:
    entries = neighbors_match.group(1)
    lines = re.findall(r'\["(.*?)"\]=\"(.*?)\"', entries)
    for node, neighbor_list in lines:
        neighbors[node] = neighbor_list.split()

# Create a new MQTT client instance
client = mqtt.Client()

# Prepare node states
nodes = {}

# MQTT message callback
def on_message(client, userdata, msg):
    topic = msg.topic
    payload = msg.payload.decode()
    for node_id in nodes:
        if topic == config_base + node_id:
            print(f"[Config Received by {node_id}] {payload}")

client.on_message = on_message

# Connect and subscribe to configuration topics
client.connect(broker_address)

for num in node_numbers:
    node_id = f"apu{str(num).zfill(2)}"
    client.subscribe(config_base + node_id)


for num in node_numbers:
    node_id = f"apu{str(num).zfill(2)}"
    x_pos = num % 5
    y_pos = num // 5
    mesh0 = []
    mesh1 = []
    for neighbor in neighbors.get(node_id, []):
        if neighbor in phy0MAC:
            mesh0.append(phy0MAC[neighbor])
        if neighbor in phy1MAC:
            mesh1.append(phy1MAC[neighbor])
    nodes[node_id] = {
        "uptime": 0,
        "XPosition": x_pos,
        "YPosition": y_pos,
        "mesh0_neighbors": mesh0,
        "mesh1_neighbors": mesh1,
        "next_publish_time": time.time() + random.uniform(5, 10)
    }

# Start a background thread for MQTT network loop
client.loop_start()

# Publish continuously
try:
    while True:
        current_time = time.time()
        for node_id, state in nodes.items():
            if current_time >= state["next_publish_time"]:
                state["uptime"] += int(current_time - (state.get("last_publish_time") or current_time))
                payload = [
                    {
                        "uptime": state["uptime"],
                        "cpu_temp": random.randint(45, 65),
                        "mem_util": random.randint(10, 25),
                        "cpu_util": random.randint(1, 23),
                        "XPosition": state["XPosition"],
                        "YPosition": state["YPosition"],
                        "mesh0_neighbors": state["mesh0_neighbors"],
                        "mesh1_neighbors": state["mesh1_neighbors"]
                    },
                    {
                        "id": node_id
                    }
                ]
                payload_json = json.dumps(payload)
                topic = topic_base + node_id
                client.publish(topic, payload_json)
                # print(f"Published to {topic}: {payload_json}")
                state["last_publish_time"] = current_time
                state["next_publish_time"] = current_time + random.uniform(2, 5)
        time.sleep(0.1)
except KeyboardInterrupt:
    print("Stopping publishing.")
    client.loop_stop()
    client.disconnect()