import time
from streamlit.delta_generator import DeltaGenerator
import streamlit as st
import os,sys
import importlib.util
import inspect

#from here https://github.com/harahu/st-clean-slate-rendering/blob/main/__main__.py
def get_clean_rendering_container() -> DeltaGenerator:
    """Makes sure we can render from a clean slate on state changes."""
    slot_in_use = st.session_state.get("slot_in_use", "initial")

    # Use slot 'a' initially, then cycle through 'b' and 'c'.
    if slot_in_use == "initial":
        slot_in_use = st.session_state.slot_in_use = "a"
    elif slot_in_use == "a":
        slot_in_use = st.session_state.slot_in_use = "b"
    elif slot_in_use == "b":
        slot_in_use = st.session_state.slot_in_use = "c"
    else:  # if slot_in_use == "c":
        slot_in_use = st.session_state.slot_in_use = "b"

    slot = {
        "a": st.empty(),
        "b": st.empty(),
        "c": st.empty(),
    }[slot_in_use]

    return slot.container()

from comparesite_models import models_funcs

def list_models():
    function_dict = {}
    for name, obj in inspect.getmembers(models_funcs):
        if inspect.isfunction(obj):
            function_dict[name] = obj
    return function_dict


@st.cache_data(show_spinner=True,ttl=15)
def get_model_response(model_choice, message):
    return models[model_choice](message)

st.set_page_config(
    page_title="NeurocoderTestEnv",
    page_icon="🧊",
    layout="wide",

)



st.subheader("📊 Neurocoder Model compare")
#st.caption("🚀 A streamlit chatbot powered by OpenAI LLM")
if "messages0" not in st.session_state:
    st.session_state["messages0"] = [{"role": "assistant", "content": "How can I help you?"}]
if "messages1" not in st.session_state:
    st.session_state["messages1"] = [{"role": "assistant", "content": "How can I help you?"}]
if "messages2" not in st.session_state:
    st.session_state["messages2"] = [{"role": "assistant", "content": "How can I help you?"}]
if "lastquery" not in st.session_state:
    st.session_state["lastquery"] = None



models = list_models()

columns = st.columns(3)
model_choice = [None,None,None]
empties = []
with st.container():
    with columns[0]:
        # Create a Streamlit selectbox with the model names
        model_choice[0] = st.selectbox(
            'Choose a model 0:',
            models
        )

        st.write(f"You selected: {model_choice[0]}")

    with columns[1]:
        # Create a Streamlit selectbox with the model names
        model_choice[1] = st.selectbox(
            'Choose a model 1:',
            models,
            index=1
        )
        st.write(f"You selected: {model_choice[1]}")
    with columns[2]:
        # Create a Streamlit selectbox with the model names
        model_choice[2] = st.selectbox(
            'Choose a model 2:',
            models,
            index=2
        )
        st.write(f"You selected: {model_choice[2]}")

dg = get_clean_rendering_container()
for i in range(0,len(st.session_state.messages0)):
    #print("i",i)
    msg = st.session_state.messages0[i]
    msg1 = st.session_state.messages1[i]
    msg2 = st.session_state.messages2[i]
    if msg["role"] == "user":
        dg.chat_message(msg["role"]).write(msg["content"])
    elif msg["role"] == "assistant":
        columnsa = dg.container().columns(spec=(1, 1, 1))
        columnsa[0].chat_message(msg["role"]).write(msg["content"])
        columnsa[1].chat_message(msg1["role"]).write(msg1["content"])
        columnsa[2].chat_message(msg2["role"]).write(msg2["content"])




if prompt := st.chat_input():


    st.session_state.messages0.append({"role": "user", "content": prompt})
    st.session_state.messages1.append({"role": "user", "content": prompt})
    st.session_state.messages2.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    st.session_state.lastquery = prompt
    response0 = get_model_response(model_choice[0], prompt)
    saveresponse0 = {"role": "assistant", "content": response0}
    st.session_state.messages0.append(saveresponse0)
    columnsf = st.columns(3)
    with columnsf[0]:
        st.chat_message(saveresponse0["role"]).write(saveresponse0["content"])
    response1 = get_model_response(model_choice[1], prompt)
    saveresponse1 = {"role": "assistant", "content": response1}
    st.session_state.messages1.append(saveresponse1)
    with columnsf[1]:
        st.chat_message(saveresponse0["role"]).write(saveresponse1["content"])
    response2 = get_model_response(model_choice[2], prompt)
    saveresponse2 = {"role": "assistant", "content": response2}
    st.session_state.messages2.append(saveresponse2)
    with columnsf[2]:
        st.chat_message(saveresponse0["role"]).write(saveresponse2["content"])