import streamlit as st
import pandas as pd
import os
import requests

st.set_page_config(page_icon='🦜', page_title='Text Generation Labeling Tool', layout='wide', initial_sidebar_state="collapsed")
st.markdown("<h1 style='text-align: center;'>Text Generation Labeling Tool</h1>", unsafe_allow_html=True)

def file_selector(folder_path=r'./Datasets'):
    filenames = os.listdir(folder_path)
    return filenames, folder_path

def revert_question_type_id(txt_question_type):
    if txt_question_type == 'What':
        return 0
    elif txt_question_type == 'Who':
        return 1
    elif txt_question_type == 'When':
        return 2
    elif txt_question_type == 'Where':
        return 3
    elif txt_question_type == 'Why':
        return 4
    elif txt_question_type == 'How':
        return 5
    elif txt_question_type == 'Others':
        return 6

filenames, folder_path = file_selector()
filename_input = st.sidebar.selectbox(label='Input dataset file:', options=filenames)
df = pd.read_csv(f'./{folder_path}/{filename_input}')

if 'idx' not in st.session_state:
    st.session_state.idx = 0

st.markdown(f"<h4 style='text-align: center;'>Sample {st.session_state.idx + 1}/{len(df)}</h4>", unsafe_allow_html=True)

col_1, col_2, col_3, col_4, col_5, col_6, col_7, col_8, col_9, col_10 = st.columns([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

btn_previous = col_1.button(label=':arrow_backward: Previous sample', use_container_width=True)
btn_next = col_2.button(label='Next sample :arrow_forward:', use_container_width=True)
btn_save = col_3.button(label=':heavy_check_mark: Save change', use_container_width=True)
txt_goto = col_4.selectbox(label='Sample', label_visibility='collapsed', options=list(range(1, len(df) + 1)))
btn_goto = col_5.button(label=':fast_forward: Move to', use_container_width=True)

if len(df) != 0:
    col_1, col_2, col_3, col_4, col_5, col_6, col_7, col_8, col_9, col_10 = st.columns(spec=[1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
    txt_context = st.text_area(height=200, label='Your context:', value=df['context'][st.session_state.idx])

    col_11, col_12, col_13 = st.columns([4.5, 1, 4.5])
    txt_question = col_11.text_area(height=90, label='Your question:', value=df['question'][st.session_state.idx])
    txt_question_type = col_12.selectbox(label='Your question type:', options=['What', 'Who', 'When', 'Where', 'Why', 'How', 'Others'], index=int(df['question_type'][st.session_state.idx]))
    txt_answer = col_13.text_area(height=90, label='Your answer:', value=df['answer'][st.session_state.idx])

    st.markdown(f"<p style='text-align: left; font-weight: normal; font-size: 14px'>Your distractors:</p>", unsafe_allow_html=True)

    col_21, col_22 = st.columns(spec=[9, 1])
    txt_distractors = col_21.text_area(height=90, label='Your distractors:', label_visibility='collapsed', value=df['distract'][st.session_state.idx])  
    btn_generate_distractor = col_22.button(label='Generate distractors', use_container_width=True)
    
    if btn_generate_distractor:
        if filename_input == 'BiologyQA.csv':
            expert = 'biologist'
        elif filename_input == 'GeographyQA.csv':
            expert = 'geographer'
        elif filename_input == 'HistoryQA.csv':
            expert = 'historian'
        url = 'https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent'
        headers = {'Content-Type': 'application/json'}
        data = {
            'contents': [
                {
                    'parts': [
                        {
                            'text': f"You are a great {expert}, here is the following content: context: '{txt_context}', question: '{txt_question}', answer: '{txt_answer}' generate three distract answers. Distractor answers are separated by [SEP]. Example: Distract answer 1 [SEP] Distract answer 2 [SEP] Distract answer 3"
                        }
                    ]
                }
            ]
        }
        api_key = 'AIzaSyApFAbCUA1H-VHAidzqmyStHFe92ODeO1Y'
        params = {'key': api_key}
        response = requests.post(url, headers=headers, json=data, params=params)
        if response.status_code == 200:
            correct = response.json()['candidates'][0]['content']['parts'][0]['text']
            st.success(f'3 distraction answers: {correct}')
            st.cache_data.clear()
        else:
            st.error('Failed to generate distractors. Please check API and inputs.')
            st.rerun()

    if btn_previous:
        if st.session_state.idx > 0:
            st.session_state.idx -= 1
            st.rerun()
        else:
            pass

    if btn_next:
        if st.session_state.idx < (len(df) - 1):
            st.session_state.idx += 1
            st.rerun()
        else:
            pass

    if btn_save:
        df['context'][st.session_state.idx] = txt_context
        df['question'][st.session_state.idx] = txt_question
        df['answer'][st.session_state.idx] = txt_answer
        df['distract'][st.session_state.idx] = txt_distractors
        df['question_type'][st.session_state.idx] = revert_question_type_id(txt_question_type)

        df.to_csv(f'./Datasets/{filename_input}', index=None)

    if btn_goto:
        st.session_state.idx = txt_goto - 1
        st.rerun()