import json
import queue

import paho.mqtt.client as mqtt
import streamlit as st
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Rectangle

import secrets

from time import time

# Initialize Streamlit app
st.title("Light-mixing Control Panel")

# Description and context
st.markdown(
    """
This application accesses a public test demo located in Toronto, ON, Canada (as of 2024-07-27). 
For more context, you can refer to this [Colab notebook](https://colab.research.google.com/github/sparks-baird/self-driving-lab-demo/blob/main/notebooks/4.2-paho-mqtt-colab-sdl-demo-test.ipynb) 
and the [self-driving-lab-demo project](https://github.com/sparks-baird/self-driving-lab-demo).
You may also be interested in the Acceleration Consortium's ["Hello World" microcourse](https://ac-microcourses.readthedocs.io/en/latest/courses/hello-world/index.html) for self-driving labs.
"""
)

max_power = 0.35
max_value = round(max_power * 255)

with st.form("mqtt_form"):

    # MQTT Configuration
    HIVEMQ_HOST = st.text_input(
        "Enter your HiveMQ host:",
        "248cc294c37642359297f75b7b023374.s2.eu.hivemq.cloud",
        type="password",
    )
    HIVEMQ_USERNAME = st.text_input("Enter your HiveMQ username:", "sgbaird")
    HIVEMQ_PASSWORD = st.text_input(
        "Enter your HiveMQ password:", "D.Pq5gYtejYbU#L", type="password"
    )
    PORT = st.number_input(
        "Enter the port number:", min_value=1, max_value=65535, value=8883
    )

    # User input for the Pico ID
    PICO_ID = st.text_input("Enter your Pico ID:", "test", type="password")

    # Information about the maximum power reduction
    st.info(
        f"The upper limit for RGB power levels has been set to {max_value} instead of 255. NeoPixels are bright 😎"
    )

    # Sliders for RGB values
    R = st.slider("Select the Red value:", min_value=0, max_value=max_value, value=0)
    G = st.slider("Select the Green value:", min_value=0, max_value=max_value, value=0)
    B = st.slider("Select the Blue value:", min_value=0, max_value=max_value, value=0)

    submit_button = st.form_submit_button(label="Send RGB Command")

command_topic = f"sdl-demo/picow/{PICO_ID}/GPIO/28/"
sensor_data_topic = f"sdl-demo/picow/{PICO_ID}/as7341/"

# random session id to keep track of the session and filter out old data
experiment_id = secrets.token_hex(4)  # 4 bytes = 8 characters
sensor_data_file = f"sensor_data-{experiment_id}.json"

# TODO: Session ID using st.session_state to have history of commands and sensor data

sensor_data_queue = queue.Queue()


# Singleton: https://docs.streamlit.io/develop/api-reference/caching-and-state/st.cache_resource
@st.cache_resource
def create_paho_client(tls=True):
    client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2, protocol=mqtt.MQTTv5)
    if tls:
        client.tls_set(tls_version=mqtt.ssl.PROTOCOL_TLS_CLIENT)
    return client


# Define setup separately since sensor_data is dynamic
def setup_paho_client(
    client, sensor_data_topic, hostname, username, password=None, port=8883
):
    def on_message(client, userdata, msg):
        sensor_data_queue.put(json.loads(msg.payload))

    def on_connect(client, userdata, flags, rc, properties=None):
        if rc != 0:
            print("Connected with result code " + str(rc))
        client.subscribe(sensor_data_topic, qos=1)

    client.on_connect = on_connect
    client.on_message = on_message

    client.username_pw_set(username, password)
    client.connect(hostname, port)
    client.loop_start()  # Use a non-blocking loop

    return client


def send_command(client, command_topic, msg):
    print("Sending command...")
    result = client.publish(command_topic, json.dumps(msg), qos=2)
    result.wait_for_publish()  # Ensure the message is sent

    if result.rc == mqtt.MQTT_ERR_SUCCESS:
        print(f"Command sent: {msg} to topic {command_topic}")
    else:
        print(f"Failed to send command: {result.rc}")


# Helper function to plot discrete spectral sensor data
def plot_spectra(sensor_data):
    """https://chatgpt.com/share/210d2fee-ca64-45a5-866e-e6df6e56bd1c"""
    wavelengths = np.array([410, 440, 470, 510, 550, 583, 620, 670])
    intensities = np.array(
        [
            sensor_data["ch410"],
            sensor_data["ch440"],
            sensor_data["ch470"],
            sensor_data["ch510"],
            sensor_data["ch550"],
            sensor_data["ch583"],
            sensor_data["ch620"],
            sensor_data["ch670"],
        ]
    )

    fig, ax = plt.subplots(figsize=(10, 6))

    num_points = 100  # for "fake" color bar effect

    # Adding rectangles for color bar effect
    dense_wavelengths = np.linspace(wavelengths.min(), wavelengths.max(), num_points)
    # rect_height = max(intensities) * 0.02  # Height of the rectangles
    rect_height = 300

    for dw in dense_wavelengths:
        rect = Rectangle(
            (
                dw - (wavelengths.max() - wavelengths.min()) / num_points / 2,
                -rect_height * 2,
            ),
            (wavelengths.max() - wavelengths.min()) / num_points,
            rect_height * 3,
            color=plt.cm.rainbow(
                (dw - wavelengths.min()) / (wavelengths.max() - wavelengths.min())
            ),
            edgecolor="none",
        )
        ax.add_patch(rect)

    # Main scatter plot
    scatter = ax.scatter(
        wavelengths, intensities, c=wavelengths, cmap="rainbow", edgecolor="k"
    )

    # Adding vertical lines from the x-axis to each point
    for wavelength, intensity in zip(wavelengths, intensities):
        ax.vlines(wavelength, 0, intensity, color="gray", linestyle="--", linewidth=1)

    # Adjust limits and labels with larger font size
    ax.set_xlim(wavelengths.min() - 10, wavelengths.max() + 10)
    ax.set_ylim(
        0, max(30000, max(intensities) + 15)
    )  # Ensure the lower y limit is 0 and add buffer with a minimum upper limit of 30000
    ax.set_xticks(wavelengths)
    ax.set_xlabel("Wavelength (nm)", fontsize=14)
    ax.set_ylabel("Intensity", fontsize=14)
    ax.set_title("Spectral Intensity vs. Wavelength", fontsize=16)
    ax.tick_params(axis="both", which="major", labelsize=12)

    st.pyplot(fig)


# Publish button
if submit_button:
    if not PICO_ID or not HIVEMQ_HOST or not HIVEMQ_USERNAME or not HIVEMQ_PASSWORD:
        st.error("Please enter all required fields.")
    else:
        st.info(
            f"Please wait while the command {R, G, B} for experiment {experiment_id} is sent..."
        )

        client = create_paho_client(tls=True)

        client = setup_paho_client(
            client,
            sensor_data_topic,
            HIVEMQ_HOST,
            HIVEMQ_USERNAME,
            password=HIVEMQ_PASSWORD,
            port=int(PORT),
        )

        command_msg = {"R": R, "G": G, "B": B, "_experiment_id": experiment_id}

        session_timeout = time() + 60
        send_command(client, command_topic, command_msg)

        while True and time() < session_timeout:

            sensor_data = sensor_data_queue.get(True, timeout=15)

            input_message = sensor_data["_input_message"]
            received_experiment_id = input_message["_experiment_id"]

            if sensor_data and received_experiment_id == experiment_id:
                R1 = input_message["R"]
                G1 = input_message["G"]
                B1 = input_message["B"]
                st.success(
                    f"Command {R1, G1, B1} for experiment {experiment_id} sent successfully!"
                )
                plot_spectra(sensor_data)
                st.write("Sensor Data Received:", sensor_data)
                break
            else:
                st.warning(
                    f"Received data for experiment {received_experiment_id} instead of {experiment_id}. Retrying..."
                )