import streamlit as st
import paho.mqtt.client as mqtt
import json
import time


# Create a form to get the MQTT broker details
st.title("PioReactor MQTT Client")
with st.form("mqtt_form"):
    hivemq_host = st.text_input("HiveMQ Host", "9bb69fc9bf6b4092abc4c92d76e602cd.s1.eu.hivemq.cloud", type="password")
    hivemq_username = st.text_input("HiveMQ Username", "Website", type="password")
    hivemq_password = st.text_input("HiveMQ Password", "Abcabcabc1", type="password")
    hivemq_port = st.number_input("HiveMQ Port", min_value=1, max_value=65535, step=1, value=8883)
    reactor = st.text_input("PioReactor", "pio1")

    submit_button = st.form_submit_button("Connect")

experiment = None
running = []
jobs = {
    "temperature_automation": False, 
    "growth_rate_calculating": False, 
    "stirring": False, 
    "od_reading": False, 
    "led_automation": False
}

# Create the MQTT client
@st.cache_resource
def create_mqtt_client(hivemq_host, hivemq_port, hivemq_username, hivemq_password):
    client = mqtt.Client()

    # Set authentication and TLS settings
    client.username_pw_set(hivemq_username, hivemq_password)
    client.tls_set(tls_version=mqtt.ssl.PROTOCOL_TLS_CLIENT)

    # Define callback functions
    def on_connect(client, userdata, flags, rc):
        if rc == 0:
            st.success("Connected to HiveMQ successfully!")
        else:
            st.error(f"Failed to connect, return code {rc}")

    client.on_connect = on_connect
    
    return client

def on_message_worker(client, userdata, message):
    payload = message.payload.decode("utf-8")
    data = json.loads(payload)
    global experiment
    global running
    experiment = data.get("experiment", None)
    running = data.get("running", [])

if submit_button:
    client = create_mqtt_client(hivemq_host, hivemq_port, hivemq_username, hivemq_password)

    # Publish a message to the MQTT broker
    if client:
        client.connect(hivemq_host, hivemq_port)
        client.loop_start()

        # Subscribe to the worker topic
        client.subscribe(f"pioreactor/{reactor}/worker")
        client.message_callback_add(f"pioreactor/{reactor}/worker", on_message_worker)

        payload = {
            "command": "get_worker",
            "reactor": reactor
        }
        payload_str = json.dumps(payload)
        client.publish("pioreactor/control", payload_str)
        
        timeout = 10
        start_time = time.time()   

        while experiment is None and (time.time() - start_time) < timeout:
            time.sleep(1)
        
        client.loop_stop()
        client.unsubscribe(f"pioreactor/{reactor}/worker")

        if experiment is None:
            st.error("No experiment assigned to the reactor.")
            st.stop()
        
        for run in running:
            jobs[run] = True

if experiment is not None:
    st.title(f"Experiment: {experiment}")
    
    cols = st.columns(5)

    with cols[0]:
        st.write("Temperature")
        st.write(jobs["temperature_automation"])

    with cols[1]:
        st.write("Growth Rate")
        st.write(jobs["growth_rate_calculating"])

    with cols[2]:
        st.write("Stirring")
        st.write(jobs["stirring"])

    with cols[3]:
        st.write("OD Reading")
        st.write(jobs["od_reading"])

    with cols[4]:
        st.write("LED")
        st.write(jobs["led_automation"])