from io import StringIO
import streamlit as st

from smem_token_parser import SMEM_Parser, read_tokens_from_lines
from smem_obj import SMEM_Obj, ObjType

MIN_COLS = 4

# Set up page config before any other streamlit commands
st.set_page_config(page_title="Soar Agent Memory Inspector", layout="wide")
st.title("Logic Inspector")


def get_smem_root_from_file(smem_file):
    # tokens = get_smem_tokens_from_local_file(smem_filename)
    tokens = read_tokens_from_lines(smem_file)
    if tokens == None:
        st.error("Error reading file: '"+str(smem_file)+"'")
        return None
    parser = SMEM_Parser()
    parser.parse_file(tokens)
    return parser.get_context_root()

if "col_obj_list" not in st.session_state:
    st.session_state["col_obj_list"] = None

file_upload_expander = st.expander(label="Select a file to inspect", expanded=(st.session_state.col_obj_list == None))
file = file_upload_expander.file_uploader(" ")

if file is not None:
    if st.session_state.col_obj_list is None:
        root = get_smem_root_from_file(StringIO(file.getvalue().decode("utf-8")))
        if root:
            st.session_state["col_obj_list"] = [None]*MIN_COLS
            st.session_state.col_obj_list[0] = root
            st.experimental_rerun()
else:
    st.session_state["col_obj_list"] = None

# inpath = "/Users/bstearn1/Documents/Projects/OptumGitHub/Intelligent-Solutions/AICarePlan/agent/D-RULES_082222_smem.soar"
# inpath = "SMEM Generator/test_pim_smem.soar"
# st.session_state.col_obj_list[0] = get_smem_root_from_file(inpath)

if st.session_state.col_obj_list is None:
    st.stop()

def add_col(index, obj):
    st.session_state.col_obj_list = st.session_state.col_obj_list[:index+1]
    st.session_state.col_obj_list.append(obj)
    # st.session_state.current_tab_index = index+1
    st.experimental_rerun()

def get_header_str(obj_type):
    if obj_type == ObjType.CONTEXT:
        return "START"
    elif obj_type == ObjType.COND_CONJ:
        return "CONDITION"
    elif obj_type == ObjType.OP:
        return "CHOICE"
    elif obj_type == ObjType.ACT_GROUP:
        return "RESULT"
    else:
        return " "

st.subheader("Click the buttons below to select conditions and to inspect resulting actions.")

cols = st.columns(max(MIN_COLS,len(st.session_state.col_obj_list)))

for i,col in enumerate(cols):
    try:
        if st.session_state.col_obj_list[i] == None:
            break
    except Exception as e:
        # print("ERROR checking column "+str(i+1)+" of "+str(len(st.session_state.col_obj_list))+": "+str(e))
        break
    # Build the available objects to navigate under this col's object
    obj = st.session_state.col_obj_list[i]
    obj_str,obj_label,obj_desc = obj.to_string()
    sub_objs = list(obj.get_child_objects(compact=True))
    # print(obj.obj_type)
    if len(sub_objs) > 0:
        col.markdown("**"+get_header_str(sub_objs[0].obj_type)+"**",)
    col.markdown("---")
    # Show the object's main title and description
    col.text(obj_str)
    if obj_desc != None:
        col.markdown("*"+obj_desc+"*")
    for j,sub in enumerate(sub_objs):
        sub_str,sub_label,sub_desc = sub.to_string()
        if sub.obj_type == ObjType.OP:
            subsub = list(sub.get_child_objects(compact=True))[0]
            _,group_string,_ = subsub.to_string()
            group_string = "\n * "+str(group_string).replace("AND", "\n * ")
            col.markdown(group_string)
            if col.button(sub_label, key="button"+str(i)+"-"+str(j)):
                add_col(i,sub)
        elif sub.obj_type == ObjType.ACT_GROUP:
            col.markdown("*Output Details:* ")
            col.markdown("\n * "+str(sub_str).replace(") AND", ")\n * "))
        elif col.button(sub_label, key="button"+str(i)+"-"+str(j)):
            add_col(i,sub)