import streamlit as st
import json
from streamlit_shortcuts import add_keyboard_shortcuts
import random 

st.set_page_config(layout="wide")

# import streamlit_authenticator as stauth
# import yaml
# from yaml.loader import SafeLoader
# with open('auth.yaml') as file:
#     config = yaml.load(file, Loader=SafeLoader)
# authenticator = stauth.Authenticate(
#     config['credentials'],
#     config['cookie']['name'],
#     config['cookie']['key'],
#     config['cookie']['expiry_days'],
#     config['preauthorized']
# )

file_path = 'grid_eval_gpt4o.json'

# Load your data
@st.cache_data()
def load_data():
    with open(file_path, 'r') as file:
        data = json.load(file)
    # random.shuffle(data)
    # data = data[]
    return data

def save_data(data):
    print(file_path.split(".json")[0])
    with open(f"{file_path.split('.json')[0]}_graded.json", 'w') as file:
        json.dump(data, file, indent=4)
        
def download_json(data):
    return json.dumps(data, indent=4)

data = load_data()

for query in data:
    for result in query['results']:
        if 'agree' not in result:
            result['agree'] = True
            
# State management for current query index
if 'current_query_index' not in st.session_state:
    st.session_state.current_query_index = 0
if 'data' not in st.session_state:
    st.session_state.data = data
if 'graded_queries' not in st.session_state:
    st.session_state.graded_queries = 0
    
def truncate_text(text, length=250):
    return text if len(text) <= length else text[:length] + '...'

result_box_style = """
<style>
.rounded-box {
    border: 1px solid #ddd;
    border-radius: 10px;
    padding: 0.01px;
    margin-bottom: 10px;
}
</style>
"""
# Navigation to next query
def next_query():
    if st.session_state.current_query_index < len(data) - 1:
        st.session_state.current_query_index += 1
    st.rerun()

add_keyboard_shortcuts({
    's': 'Skip',
})

# Display current query and its results
def display_query():
    # Navigation bar
    global current_query    

    st.session_state.graded_queries = sum(query.get('status', None) is not None for query in st.session_state.data)
    print(f"Current Query Index: {st.session_state.current_query_index} | Graded Queries: {st.session_state.graded_queries} | Total Queries: {len(st.session_state.data)} | Current Query Status {current_query.get('status', None)}")
    col1, col2 = st.columns([4, 2], gap="small")
    with col1:
        if st.button('Previous'):
            if st.session_state.current_query_index > 0:
                st.session_state.current_query_index -= 1 % len(st.session_state.data)
                st.rerun()
        st.progress((st.session_state.current_query_index + 1) / len(st.session_state.data))
    with col2:
        col1, col2, col3, col4 = st.columns([1, 1, 1, 2], gap = "small")
        with col1:
            if st.button('Next'):
                if st.session_state.current_query_index < len(st.session_state.data) - 1:
                    st.session_state.current_query_index += 1 % len(st.session_state.data)
                    st.rerun()
        with col2:
            if st.button('Skip'):
                current_query['status'] = 'skipped'
                next_query()
        with col3:
            if st.button('Junk'):
                current_query['status'] = 'nonsense'
                next_query()
        with col4:
            # Example button for downloading data
            st.download_button(
                label="Download",
                data=download_json(st.session_state.data),
                file_name="graded_data.json",
                mime="application/json"
            )

        index = st.text_input(f"At index {st.session_state.current_query_index + 1}. Graded: {st.session_state.graded_queries}/{len(st.session_state.data)}", placeholder="Go to index:")
        if index:
            try:
                index = int(index) - 1
                if index < 0 or index >= len(data):
                    st.error("Invalid index.")
                else:
                    st.session_state.current_query_index = index
                    st.rerun()
            except ValueError:
                st.error("Please enter a valid integer.")
            
    if st.session_state.graded_queries >= len(data):
        save_data(st.session_state.data)
        st.success(f"{len(data)} Queries graded and data saved!")
    
    st.markdown(result_box_style, unsafe_allow_html=True)

    st.header(f"Query: {current_query['query']}")
    status_color = 'green' if current_query.get('status', None) is not None else 'red'
    st.markdown(f"{current_query['grid_pos_str']} | Query Grade: <b style='color: {status_color};'>{'Graded' if status_color == 'green' else 'Ungraded'}</b>", unsafe_allow_html = True)
    st.markdown(f"Model's Query Gen Reasoning Trace: {current_query['reasoning_trace'][0]}")
    st.subheader("Results:")
    for index, result in enumerate(current_query['results']):
        st.markdown(f"<div class='rounded-box'>", unsafe_allow_html=True)
        col1, col2 = st.columns([3, 2], gap="small")
        with col1:            
            st.markdown(f"<h5>{result['title']}</h5>", unsafe_allow_html=True)
            st.markdown(f"[<span style='font-size: 0.8em;'>{truncate_text(result['url'], length = 50)}</span>]({result['url']})  |  {result['published_date']}", unsafe_allow_html=True)
            st.markdown(f"{truncate_text(result['text'], length = len(result['model_trace']))}")
        with col2:
            grade_color = 'green' if result['grade'].lower() == 'yes' else 'red'
            st.markdown(f"Model Grade: <b style='color: {grade_color};'>{result['grade']}</b>", unsafe_allow_html=True)
            st.write(result['model_trace'])
            
            if st.checkbox("Reject", value= not result.get('agree'), key=f'verify-{index}'):
                result['agree'] = False
                
        st.markdown("</div>", unsafe_allow_html=True)
    st.markdown(f"<div class='rounded-box'>", unsafe_allow_html=True)

# Show current query and its results
current_query = st.session_state.data[st.session_state.current_query_index]

display_query()

col1, col2 = st.columns([5, 1], gap="small")
with col2:
    if st.button('Mark Done and Go to Next'):
        current_query['status'] = 'graded'
        next_query()
    
add_keyboard_shortcuts({
    'j': 'Junk',
})
add_keyboard_shortcuts({
    'p': 'Previous',
})
add_keyboard_shortcuts({
    'n': 'Next',
})